Spaces:
Running
Running
| import json | |
| import os | |
| import sys | |
| import difflib | |
| import time | |
| from pathlib import Path | |
| from google import genai | |
| from google.genai import types | |
| def main(): | |
| if not os.environ.get("GEMINI_API_KEY"): | |
| print("Error: GEMINI_API_KEY environment variable not set.") | |
| print("Please set it before running this script: export GEMINI_API_KEY='your-key'") | |
| sys.exit(1) | |
| client = genai.Client() | |
| base_dir = Path(__file__).parent.parent | |
| input_file = base_dir / "data" / "clean" / "mansfield.jsonl" | |
| output_file = base_dir / "data" / "clean" / "mansfield_polished.jsonl" | |
| if not input_file.exists(): | |
| print(f"Error: Could not find {input_file}") | |
| sys.exit(1) | |
| print(f"Reading {input_file}...") | |
| entries = [] | |
| with open(input_file, 'r', encoding='utf-8') as f: | |
| for line in f: | |
| entries.append(json.loads(line)) | |
| print(f"Found {len(entries)} entries. Starting polish pass...") | |
| system_instruction = """You are an expert OCR corrector and editor. | |
| Your task is to fix OCR errors, typos, and stray spacing in a diary entry by Katherine Mansfield. | |
| DO NOT rewrite her prose. | |
| DO NOT alter her punctuation style, fragmented sentences, or voice. | |
| DO NOT add commentary or explanation. | |
| Output ONLY the corrected text of the entry.""" | |
| changes_made = 0 | |
| with open(output_file, 'w', encoding='utf-8') as out_f: | |
| for i, entry in enumerate(entries): | |
| original_text = entry.get('text', '') | |
| if not original_text.strip(): | |
| out_f.write(json.dumps(entry) + '\n') | |
| continue | |
| try: | |
| # Sleep briefly to avoid hammering the API | |
| time.sleep(0.1) | |
| response = client.models.generate_content( | |
| model='gemini-2.5-flash', | |
| contents=original_text, | |
| config=types.GenerateContentConfig( | |
| system_instruction=system_instruction, | |
| temperature=0.1 | |
| ) | |
| ) | |
| polished_text = response.text.strip() | |
| # Compare similarity to ensure no wild rewrites | |
| similarity = difflib.SequenceMatcher(None, original_text, polished_text).ratio() | |
| if similarity < 0.7: | |
| # Too different, LLM might have rewritten it entirely | |
| print(f"[{i+1}/{len(entries)}] Warning: High edit distance (similarity {similarity:.2f}). Keeping original.") | |
| final_text = original_text | |
| elif original_text != polished_text: | |
| changes_made += 1 | |
| print(f"[{i+1}/{len(entries)}] Polished entry (similarity {similarity:.2f})") | |
| final_text = polished_text | |
| else: | |
| final_text = original_text | |
| except Exception as e: | |
| print(f"[{i+1}/{len(entries)}] Error processing entry: {e}. Keeping original.") | |
| final_text = original_text | |
| entry['text'] = final_text | |
| out_f.write(json.dumps(entry) + '\n') | |
| out_f.flush() | |
| print(f"Done! Polished {changes_made} out of {len(entries)} entries.") | |
| print(f"Polished corpus saved to {output_file}") | |
| print("You can use a diff tool to compare the original and polished files.") | |
| if __name__ == "__main__": | |
| main() | |