diff --git a/annotating.py b/annotating.py index 8f32034..a41d4e1 100644 --- a/annotating.py +++ b/annotating.py @@ -236,7 +236,8 @@ def render_real_latex_text(text, width_px, bg_color=(255, 255, 255, 255), max_li \\usepackage[T1]{{fontenc}} \\usepackage{{lmodern}} % Enables arbitrary font scaling \\usepackage{{amsmath, amssymb}} -%\\usepackage{{anyfontsize}} % replaces by lmodern +\\usepackage{{commands}} +%\\usepackage{{anyfontsize}} % replaced by lmodern \\begin{{document}} \\fontsize{{{fontsize}}}{{{line_spacing}}}\\selectfont {text} diff --git a/correction.py b/correction.py index 7d31f64..382ae0c 100644 --- a/correction.py +++ b/correction.py @@ -217,8 +217,8 @@ completed_tasks = [] # --- Lock for thread-safe file writing --- io_lock = threading.Lock() -pro_lock = threading.Lock() # New lock for counter -pro_count = 0 # New counter +pro_lock = threading.Lock() +pro_count = 0 flash_count = 0 if overwrite: @@ -230,63 +230,68 @@ else: if progress_path.exists(): with open(progress_path, "r", encoding="utf-8") as f: completed_tasks = json.load(f) - # Reload existing results to avoid overwriting them with partial data if output_path.exists(): with open(output_path, "r", encoding="utf-8") as f: results = json.load(f) -# Create a set for O(1) lookup. Normalize paths to strings. completed_set = set((str(f), l) for f, l in completed_tasks) - -# Filter tasks first to avoid overhead in threads tasks_to_process = [t for t in tasks if (str(t[0]), t[1]) not in completed_set] +def call_gemini_with_retries(model_id, contents, config): + """Handles requests to Gemini with a 1min and 5min retry mechanism.""" + delays = [60, 300] + for attempt in range(3): + try: + full_response_text = "" + for chunk in client.models.generate_content_stream( + model=model_id, + contents=contents, + config=config, + ): + if chunk.text: + full_response_text += chunk.text + return full_response_text + except Exception as e: + if attempt < 2: + print(f"\tGemini API failure: {e}. Retrying in {delays[attempt]} seconds...") + time.sleep(delays[attempt]) + else: + print(f"\tGemini API failure: {e}. Maximum retries reached.") + raise + def process_single_task(task_tuple): global pro_count, flash_count file_path, label = task_tuple group_name = os.path.splitext(file_path)[0] json_path = group_name + '.json' + with open(json_path, 'r') as f: - # List of (groupid, start, end), in pixels group_data = json.load(f) + n = len(group_data) - # l[3] is ratio of width to width of group d_data = {l[0]: (l[1], l[2], l[3]) for l in group_data} total_height = group_data[-1][2] use_flash = n >= 4 or total_height <= 500 + if not use_flash: with pro_lock: if limit is None or pro_count < limit: pro_count += 1 else: - # Limit reached, force switch to Flash use_flash = True + if use_flash: with pro_lock: flash_count += 1 try: contents, config = generate_request(file_path, label) - if use_flash: - print(f"Asking Gemini Flash: {label} {group_name}") - else: - print(f"Asking Gemini Pro : {label} {group_name}") + model_to_use = MODEL_ID_flash if use_flash else MODEL_ID_pro + print(f"Asking Gemini {'Flash' if use_flash else 'Pro '}: {label} {group_name}") - full_response_text = "" - # Assuming client is thread-safe (usually is). - # If not, create a new client instance inside this function. - for chunk in client.models.generate_content_stream( - model=MODEL_ID_flash if use_flash else MODEL_ID_pro, - contents=contents, - config=config, - ): - if chunk.text: - full_response_text += chunk.text - - # Parse JSON + full_response_text = call_gemini_with_retries(model_to_use, contents, config) json_data = json.loads(full_response_text) - # print("Debug : ", json_data) # Ensure consistency of answer placements for p in json_data: pid = p["id"] @@ -303,23 +308,20 @@ def process_single_task(task_tuple): if pid not in d_data: print("Error : Gemini answered a copie id not present", pid, label, group_name) continue - yming,ymaxg, width_r = d_data[pid] + yming, ymaxg, width_r = d_data[pid] if ymin < yming-50 or ymax > ymaxg+50: - print("Error : Gemini answered box2d too low/up", - pid, label, group_name) + print("Error : Gemini answered box2d too low/up", pid, label, group_name) if ymax < yming or ymin > ymaxg: print("Removing the box.") f["box_2d"] = None continue nymin = max(ymin, yming) * 1000 // total_height nymax = min(ymax, ymaxg) * 1000 // total_height - f["box_2d"] = [nymin, xmin, nymax, xmax] - # print("Group :", yming, ymaxg, "Answered:", ymin, ymax) - if xmax / 1000 > width_r: - print("Error : Gemini answered box2d too right", - pid, label, group_name) - if xmin /1000 > width_r: + + if f["box_2d"] and xmax / 1000 > width_r: + print("Error : Gemini answered box2d too right", pid, label, group_name) + if xmin / 1000 > width_r: print("Removing the box.") f["box_2d"] = None continue @@ -328,10 +330,9 @@ def process_single_task(task_tuple): # --- Use Lock for writing shared data --- with io_lock: if label not in results: - results[label] = [] # Ensure key exists if not using defaultdict + results[label] = [] results[label].append(json_data) - # Save Results with open(output_path, "w", encoding="utf-8") as f: json.dump(results, f, indent=2)