Spaces:
Sleeping
Sleeping
| import spaces # Must precede torch for ZeroGPU. | |
| import difflib | |
| import logging | |
| import gradio as gr | |
| import torch | |
| from transformers import AutoTokenizer, T5ForConditionalGeneration | |
| MODEL_ID = "aimgo/Emendator" | |
| MODEL_REVISION = "52040eff93121a80d63e3ee25a56714f578a21d2" | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| MAX_CHARACTERS = 250 | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| logger.info("Loading %s on %s", MODEL_ID, DEVICE) | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, revision=MODEL_REVISION) | |
| model = T5ForConditionalGeneration.from_pretrained( | |
| MODEL_ID, revision=MODEL_REVISION, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True | |
| ).to(DEVICE).eval() | |
| logger.info("Emendator loaded successfully.") | |
| def highlight_changes(original: str, corrected: str) -> list[tuple[str, str | None]]: | |
| differences = [] | |
| for operation, i1, i2, j1, j2 in difflib.SequenceMatcher(None, original, corrected, autojunk=False).get_opcodes(): | |
| if operation == "equal": | |
| differences.append((original[i1:i2], None)) | |
| else: | |
| if operation in ("delete", "replace"): | |
| differences.append((original[i1:i2], "Removed")) | |
| if operation in ("insert", "replace"): | |
| differences.append((corrected[j1:j2], "Added")) | |
| return differences | |
| def correct_latin(text: str) -> tuple[str, list[tuple[str, str | None]]]: | |
| if not text or not text.strip(): | |
| raise gr.Error("Input a Latin passage or choose an example.") | |
| if len(text) > MAX_CHARACTERS: | |
| raise gr.Error(f"Your passage contains {len(text)} characters. Shorten it to 250 characters or fewer.") | |
| # ByT5 tokenizes UTF-8 bytes; avoid silently truncating accented text. | |
| inputs = tokenizer(text, return_tensors="pt", truncation=False).to(DEVICE) | |
| with torch.inference_mode(): | |
| output = model.generate( | |
| **inputs, max_new_tokens=inputs["input_ids"].shape[1] + 32, | |
| num_beams=4, do_sample=False, early_stopping=True, repetition_penalty=1.15, use_cache=True | |
| ) | |
| token_ids = output[0].detach().cpu().tolist() | |
| if tokenizer.eos_token_id not in token_ids: | |
| raise gr.Error("Generation reached its length limit. Try a shorter passage.") | |
| corrected = tokenizer.decode(token_ids, skip_special_tokens=True) | |
| if not corrected.strip(): | |
| raise gr.Error("The model returned an empty result. Try another passage.") | |
| return corrected, highlight_changes(text, corrected) | |
| EXAMPLES = [ | |
| ["Nil igirur rnors cft ad nos ncq;pcrtinct hilurn»"], | |
| ["Vt quod ali cibus eft aliis fuat acre uenenurn."], | |
| ["atqu optimo viro, peterem, superavi tamen dignitate Catilinam, qatia Gdbam. Quod si id crimen homini novo esse deberet, protecto"], | |
| ["/Egidinslufitanusolufiponc natusS.Thcol.Mag.fcripfitCorn rncntariurninlibros>Scntcntiarurn«"], | |
| ["Diligamtc Do rnincfortitudomca:Dominusfirma rncntumrncurn&rcfugiurnmcurn&libcratorrnc9«"], | |
| ] | |
| with gr.Blocks(title="Emendator — Latin OCR Correction") as demo: | |
| gr.Markdown( | |
| "# Latin OCR Correction\n" | |
| "Correct Latin OCR errors with [Emendator](https://huggingface.co/aimgo/Emendator).\n\n" | |
| "Use passages of 250 characters. Shorter or longer sequences lower performance." | |
| "Carefully review the output, as the model may change spelling, punctuation, or wording." | |
| ) | |
| if DEVICE == "cpu": | |
| gr.Markdown("**Running on CPU.** Inference will be slow.") | |
| with gr.Row(): | |
| source = gr.Textbox(label="Latin OCR input", placeholder="Paste your Latin OCR text here...", lines=7) | |
| result = gr.Textbox(label="Corrected Latin", lines=7, interactive=False) | |
| with gr.Row(): | |
| submit = gr.Button("Correct Latin", variant="primary") | |
| clear = gr.ClearButton(value="Clear") | |
| differences = gr.HighlightedText( | |
| label="Changes", show_legend=True, combine_adjacent=True, | |
| color_map={"Removed": "#fecaca", "Added": "#bbf7d0"} | |
| ) | |
| gr.Examples(examples=EXAMPLES, inputs=[source], cache_examples=False, label="Examples from the model card") | |
| submit.click(fn=correct_latin, inputs=[source], outputs=[result, differences], api_name="correct", concurrency_limit=1) | |
| clear.add([source, result, differences]) | |
| gr.Markdown( | |
| "Model: Aidan McCarthy, *Emendator* (2026). Model and examples licensed under " | |
| "[CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/)." | |
| ) | |
| demo.queue(max_size=16, default_concurrency_limit=1) | |
| if __name__ == "__main__": | |
| demo.launch() |