Update true_pre_2level.py

true_pre_2level.py

@@ -9,6 +9,7 @@ import csv
 # Set global variables
 NUM_EXAMPLES = 900  # 要处理的 JSONL 文件的行数（即例子数量）The number of lines (i.e., examples) in the JSONL file to be processed.
 QUESTIONS_PER_EXAMPLE = 2  # 每个例子的标准问题数量 The standard number of questions per example
+Trans_level = 1  # 定义全局变量，决定是否切换到跨级判别框架，0为不启用跨级判别框架，1为启用 Define a global variable to decide whether to switch to the trans-level discrimination framework. 0 means disabling the trans-level discrimination framework, and 1 means enabling it.
 
 model = '/path/model'
 jsonl_file = '/path/val.jsonl'
@@ -72,18 +73,10 @@ def process_conversations(data, session=None):
     for i in range(0, len(conv), 2):
         human_question = conv[i]['value']
 
-        # if '<image>' in human_question:
-            
-        # else:
-        #     response = pipe.chat(human_question, session=session, gen_config=gen_config)
-
-        # generated_answer = response.response.text.strip()
-
         if i == 0:  # First question
             response = pipe.chat((human_question, image), session=session, gen_config=gen_config)
             generated_answer = response.response.text.strip()
             first_level = extract_level(conv[i + 1]['value'], r'The FirstLevel is (.+)$')
-            # predicted_first_level = extract_level(generated_answer, r'The FirstLevel is (.+)$')
             predicted_first_level = extract_level(generated_answer, r'(?:The FirstLevel is )?(.+)$')
             # Update first_level_accuracy
             if first_level in first_level_accuracy:
@@ -106,16 +99,22 @@ def process_conversations(data, session=None):
 
         elif i == 2:  # Second question
             # 提取POI信息 POI
+            # Extract POI information (Points of Interest)
             poi_info = extract_poi_info(conv[i]['value'])
             # 提取行人密度信息 People 
+            # Extract pedestrian density information
             pedestrian_density = extract_pedestrian_density(conv[i]['value'])
+            
+            if Trans_level == 0:
+                current_first_level = predicted_first_level
+            else:
+                current_first_level = first_level
 
-            if predicted_first_level in first_level_to_second_levels:
-                second_levels = first_level_to_second_levels[predicted_first_level]
+            if current_first_level in first_level_to_second_levels:
+                second_levels = first_level_to_second_levels[current_first_level]
                 # 构建下一个问题，包括POI和行人密度信息
-                # 如果使用跨级判别框架，请修改{predicted_first_level}变量为{first_level}变量
-                # If using a cross-level discriminative framework, please modify the {predicted_first_level} variable to the {first_level} variable.
-                human_question = (f"The FirstLevel category of this image is {predicted_first_level}. "
+                # Construct the next question, including POI and pedestrian density information
+                human_question = (f"The FirstLevel category of this image is {current_first_level}. "
                                   f"Please select the most likely SecondLevel among {', '.join(second_levels)}. "
                                   "This image contains some POI (Point of Interest) information, "
                                   "which is now provided to you. You can refer to this POI information "
@@ -133,22 +132,25 @@ def process_conversations(data, session=None):
 
             second_level = extract_level(conv[i + 1]['value'], r'The SecondLevel is (.+)$')
             predicted_second_level = extract_level(generated_answer, r'(?:The SecondLevel is )?(.+)$')
-            # predicted_second_level = extract_level(generated_answer, r'The SecondLevel is (.+)$')
 
             # Update second_level_accuracy
             # 确保真实的第一级分类存在于结构中
+            # Ensure the real first-level category exists in the structure
             real_first_level = extract_level(conv[i - 1]['value'], r'The FirstLevel is (.+)$')
 
             if real_first_level in second_level_accuracy:
                 # 使用真实的一级分类查找对应的二级分类结构
+                # Use the real first-level category to find the corresponding second-level structure
                 second_level_data = second_level_accuracy[real_first_level]
 
                 # 更新统计总数，确保真实的二级分类存在于结构中
+                # Update statistics, ensuring the real second-level category exists in the structure
                 if second_level in second_level_data:
                     second_level_data[second_level]['total'] += 1
                     second_level_writer.writerow([data['id'], real_first_level, second_level, predicted_second_level])
 
                     # 比较和记录正确性
+                    # Compare and record correctness
                     if second_level == predicted_second_level:
                         correct_second += 1
                         second_level_data[second_level]['correct'] += 1
@@ -172,9 +174,9 @@ correct_first_total = 0
 correct_second_total = 0
 error_logs = []
 
-
 with open(jsonl_file, 'r') as f:
     lines = [next(f) for _ in range(NUM_EXAMPLES)]  # 只读取前 NUM_EXAMPLES 行
+                                                    # Only read the first NUM_EXAMPLES lines
 
 with open('first_level_results_true_0_9K.csv', 'w', newline='') as first_level_csv_file, \
      open('second_level_results_true_0_9K.csv', 'w', newline='') as second_level_csv_file:
@@ -194,6 +196,7 @@ with open('first_level_results_true_0_9K.csv', 'w', newline='') as first_level_c
         error_logs.extend(errors)
 
 # 计算和打印正确率
+# Calculate and print accuracy
 total_questions = NUM_EXAMPLES * QUESTIONS_PER_EXAMPLE
 
 first_accuracy = correct_first_total / NUM_EXAMPLES
@@ -204,6 +207,7 @@ print(f'Second question accuracy: {second_accuracy * 100:.2f}%')
 print(f'Overall accuracy: {((correct_first_total + correct_second_total) / total_questions) * 100:.2f}%')
 
 # 计算一级分类正确率
+# Calculate FirstLevel accuracy
 for first_level in first_level_accuracy:
     correct = first_level_accuracy[first_level]['correct']
     total = first_level_accuracy[first_level]['total']
@@ -212,6 +216,7 @@ for first_level in first_level_accuracy:
         print(f'Accuracy for FirstLevel "{first_level}": {accuracy * 100:.2f}% right/total: {correct} / {total} ')
 
 # 计算二级分类正确率
+# Calculate SecondLevel accuracy
 for first_level, second_levels in second_level_accuracy.items():
     for second_level in second_levels:
         correct = second_levels[second_level]['correct']
@@ -221,5 +226,6 @@ for first_level, second_levels in second_level_accuracy.items():
             print(f'Accuracy for SecondLevel "{second_level}" under FirstLevel "{first_level}": {accuracy * 100:.2f}% right/total: {correct} / {total}')
 
 # 将错误记录写入日志文件
+# Write error logs to a file
 with open('error_log_0_9K', 'w') as outfile:
     json.dump(error_logs, outfile, indent=4)