Spaces:
Running
Running
File size: 3,538 Bytes
be1ce5d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | 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()
|