review.py (4439B)
1 import sys 2 import os 3 import json 4 import subprocess 5 import re 6 from datetime import timedelta 7 8 # --- Configuration --- 9 HIGHLIGHT_SCRIPT_NAME = "highlight.py" 10 GOOD_DATASET_FILE = "dataset.jsonl" 11 BAD_FILES_LOG = "errors.txt" 12 PROGRESS_FILE = ".curation_progress.log" 13 14 # --- Unbuffered Input (Cross-Platform) --- 15 try: 16 # Unix-like systems (Linux, macOS) 17 import tty, termios 18 def getch(): 19 fd = sys.stdin.fileno() 20 old_settings = termios.tcgetattr(fd) 21 try: 22 tty.setraw(sys.stdin.fileno()) 23 ch = sys.stdin.read(1) 24 finally: 25 termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) 26 return ch 27 except ImportError: 28 # Windows 29 import msvcrt 30 def getch(): 31 return msvcrt.getch().decode() 32 33 # --- Main Curation Logic --- 34 35 def review_files(file_paths: list[str]): 36 """ 37 Iterates through files, displays them for review, and sorts them based on user input. 38 """ 39 if not os.path.exists(HIGHLIGHT_SCRIPT_NAME): 40 print(f"Error: The visualization script '{HIGHLIGHT_SCRIPT_NAME}' was not found in this directory.") 41 sys.exit(1) 42 43 # Load the set of already processed files 44 processed_files = set() 45 if os.path.exists(PROGRESS_FILE): 46 with open(PROGRESS_FILE, 'r', encoding='utf-8') as f: 47 processed_files = {line.strip() for line in f} 48 49 total_files = len(file_paths) 50 for i, filepath in enumerate(file_paths): 51 if filepath in processed_files: 52 print(f"Skipping already processed file: {filepath}") 53 continue 54 55 print("-" * 80) 56 print(f"Reviewing file ({i + 1}/{total_files}): {filepath}") 57 58 # 1. Run highlight.py and pipe its output to the user's pager 59 try: 60 pager_cmd = os.environ.get("PAGER", "less").split() 61 highlight_process = subprocess.Popen( 62 ["python", HIGHLIGHT_SCRIPT_NAME, filepath], 63 stdout=subprocess.PIPE, 64 text=True, 65 encoding='utf-8' 66 ) 67 pager_process = subprocess.Popen(pager_cmd, stdin=highlight_process.stdout) 68 highlight_process.stdout.close() 69 return_code = pager_process.wait() 70 if return_code != 0: 71 print(f"Warning: Pager exited with code {return_code}.") 72 except FileNotFoundError: 73 print(f"Error: Pager command not found. Is '{os.environ.get('PAGER', 'less')}' installed?") 74 continue 75 except Exception as e: 76 print(f"An error occurred while running the subprocess: {e}") 77 continue 78 79 # 2. Prompt for user decision 80 while True: 81 print("\nWas this file's labeling correct? (y/n): ", end="", flush=True) 82 choice = getch().lower() 83 print(choice) 84 85 if choice == 'y': 86 print(f" -> Appending {filepath} to {GOOD_DATASET_FILE}") 87 try: 88 with open(filepath, 'r', encoding='utf-8') as infile, \ 89 open(GOOD_DATASET_FILE, 'a', encoding='utf-8') as outfile: 90 for line in infile: 91 outfile.write(line) 92 except Exception as e: 93 print(f"Error appending to dataset file: {e}") 94 break 95 elif choice == 'n': 96 print(f" -> Logging {filepath} to {BAD_FILES_LOG}") 97 try: 98 with open(BAD_FILES_LOG, 'a', encoding='utf-8') as errfile: 99 errfile.write(filepath + '\n') 100 except Exception as e: 101 print(f"Error writing to log file: {e}") 102 break 103 elif choice in ['\x03', 'q']: # Handle Ctrl+C or 'q' to quit 104 print("\nExiting review process.") 105 sys.exit(0) 106 else: 107 print("Invalid input. Please press 'y' for yes or 'n' for no.") 108 109 # Log the file as processed after a decision has been made 110 with open(PROGRESS_FILE, 'a', encoding='utf-8') as f: 111 f.write(filepath + '\n') 112 113 if __name__ == "__main__": 114 if len(sys.argv) < 2: 115 print("Usage: python curate_dataset.py <file1.jsonl> <file2.jsonl> ...") 116 print("Or: python curate_dataset.py /path/to/*/*.jsonl") 117 sys.exit(1) 118 119 review_files(sys.argv[1:]) 120 print("-" * 80) 121 print("Curation complete.")