"""Gradio UI for the GEC inline-edit fine-tuned model. Deploys to HuggingFace Spaces on ZeroGPU: the heavy generation function is wrapped with @spaces.GPU so the model is loaded on a free H100 burst. Model dropdown lets you switch between: - "Base + few-shot" -> Qwen/Qwen2.5-3B-Instruct with a 3-shot prompt (fair comparison baseline — see report.md) - "SFT" -> base + the SFT LoRA adapter - "DPO" -> merged SFT model + the DPO LoRA adapter (the DPO adapter was trained on top of the merged 16-bit SFT model, not the raw base) Environment variables: GEC_BASE_MODEL default 'Qwen/Qwen2.5-3B-Instruct' GEC_SFT_ADAPTER HF Hub id of the SFT adapter (required for SFT) GEC_DPO_ADAPTER HF Hub id of the DPO adapter (required for DPO) GEC_DPO_BASE_MODEL HF Hub id of the merged SFT model the DPO adapter sits on (required for DPO, e.g. 'Lopato4ka/qwen2.5-3b-gec-sft-merged') """ from __future__ import annotations import html import os import re import time from functools import lru_cache import gradio as gr try: import spaces HAS_SPACES = True except ImportError: HAS_SPACES = False from gec.inference import generate_batch, load_model from gec.parse import parse_inline INTRO = """\ # ✏️ GEC inline-edit — fine-tuned Qwen2.5-3B This space corrects English grammar by emitting **inline edits** in the format `I {goes=>go} to school` — the exact syntax requested in the LPNLP fine-tuning assignment. Pick a model variant on the right and paste any English sentence. The left panel shows the model's raw bracketed output; the right one shows the cleaned-up corrected sentence after parsing the brackets back out. """ EXAMPLES = [ "I goes to school every day .", "She have did her homework already .", "The childrens was playing in park yesterday .", "He don't likes apples and oranges .", "We was supposed to met them at noon ; but they didn't came .", "Although she is tired but she continue working .", ] BASE_MODEL = os.environ.get("GEC_BASE_MODEL", "Qwen/Qwen2.5-3B-Instruct") SFT_ADAPTER = os.environ.get("GEC_SFT_ADAPTER", "") DPO_ADAPTER = os.environ.get("GEC_DPO_ADAPTER", "") DPO_BASE_MODEL = os.environ.get("GEC_DPO_BASE_MODEL", "") VARIANTS = ["Base + few-shot"] if SFT_ADAPTER: VARIANTS.append("SFT") if DPO_ADAPTER and DPO_BASE_MODEL: VARIANTS.append("DPO") DEFAULT_VARIANT = "DPO" if "DPO" in VARIANTS else ("SFT" if "SFT" in VARIANTS else VARIANTS[0]) @lru_cache(maxsize=4) def _load(variant: str): """Load (tokenizer, model) lazily. Cached per variant.""" print(f"[boot] loading variant={variant} …") t0 = time.time() base, adapter = BASE_MODEL, None if variant == "SFT": adapter = SFT_ADAPTER elif variant == "DPO": # The DPO adapter was trained on the merged 16-bit SFT model. base, adapter = DPO_BASE_MODEL, DPO_ADAPTER tok, model = load_model(base, adapter_id=adapter) # On a GPU host (e.g. the Modal demo) move the weights over; on CPU # Spaces this is a no-op and on ZeroGPU the @spaces.GPU decorator # owns device placement (cuda is not visible at load time there). import torch if torch.cuda.is_available(): model = model.to("cuda") print(f"[boot] loaded {variant} in {time.time() - t0:.1f}s") return tok, model _BRACKET_RE = re.compile(r"\{([^{}=]*?)=>([^{}=]*?)\}") def _highlight_edits(raw: str) -> str: """Render the bracketed string as HTML with red strike + green replacement.""" def repl(m: re.Match) -> str: src = html.escape(m.group(1).strip()) tgt = html.escape(m.group(2).strip()) if not src: return f'+{tgt}' if not tgt: return f'{src}' return ( f'{src}' f'{tgt}' ) escaped = html.escape(raw) # Re-allow the `{x=>y}` markup we control. escaped = re.sub( r"\{([^{}=]*?)=>([^{}=]*?)\}", lambda m: "{" + m.group(1) + "=>" + m.group(2) + "}", escaped, ) return _BRACKET_RE.sub(repl, escaped) def _correct_impl(sentence: str, variant: str) -> tuple[str, str, str]: """Run the model and return (highlighted HTML, raw bracketed, cleaned text).""" sentence = (sentence or "").strip() if not sentence: return "_Type a sentence first._", "", "" if variant not in VARIANTS: return f"_Variant `{variant}` unavailable — check env vars._", "", "" tok, model = _load(variant) few_shot = variant == "Base + few-shot" result = generate_batch([sentence], tok, model, include_few_shot=few_shot, batch_size=1)[0] cleaned = result.corrected if not result.parse_ok: cleaned = f"{cleaned} ⚠️ (output had unmatched braces)" return _highlight_edits(result.raw), result.raw, cleaned if HAS_SPACES: @spaces.GPU(duration=60) def correct(sentence: str, variant: str): return _correct_impl(sentence, variant) else: def correct(sentence: str, variant: str): return _correct_impl(sentence, variant) def build_ui() -> gr.Blocks: with gr.Blocks(title="GEC inline-edit") as demo: gr.Markdown(INTRO) with gr.Row(): with gr.Column(scale=3): sentence = gr.Textbox( label="Sentence to correct", placeholder="I goes to school every day .", lines=2, autofocus=True, ) with gr.Row(): correct_btn = gr.Button("Correct", variant="primary", scale=2) clear_btn = gr.Button("Clear", scale=1) highlighted = gr.HTML(label="Edits (highlighted)") raw_box = gr.Textbox(label="Raw model output (bracketed)", interactive=False) clean_box = gr.Textbox(label="Cleaned corrected sentence", interactive=False) gr.Examples(EXAMPLES, inputs=sentence, label="Try these") with gr.Column(scale=1, min_width=260): variant = gr.Dropdown( choices=VARIANTS, value=DEFAULT_VARIANT, label="Model variant", ) gr.Markdown( "**Base model:** " f"`{BASE_MODEL}`\n\n" f"**SFT adapter:** `{SFT_ADAPTER or '(not configured)'}`\n\n" f"**DPO adapter:** `{DPO_ADAPTER or '(not configured)'}`\n\n" "---\n\n" "First request may take 30–60 s while ZeroGPU warms up. " "Subsequent requests within the 60 s window are fast." ) correct_btn.click(correct, inputs=[sentence, variant], outputs=[highlighted, raw_box, clean_box]) sentence.submit(correct, inputs=[sentence, variant], outputs=[highlighted, raw_box, clean_box]) clear_btn.click(lambda: ("", "", "", ""), outputs=[sentence, highlighted, raw_box, clean_box]) return demo if __name__ == "__main__": build_ui().queue().launch( server_name="0.0.0.0", server_port=int(os.environ.get("PORT", "7860")), theme=gr.themes.Soft(), ssr_mode=False, )