mininfradev/glossofication_v2

1import os, zipfile, shutil2from pathlib import Path3import torch4import gradio as gr5from transformers import AutoTokenizer, AutoModelForSeq2SeqLM6 7MODEL_ZIP = os.getenv("MODEL_ZIP", "model_best.zip")8MODEL_DIR = os.getenv("MODEL_DIR", "model_best")9PREFIX    = os.getenv("PREFIX", "translate English to Gloss: ")10MAX_SRC   = int(os.getenv("MAX_SRC", "96"))11MAX_TGT   = int(os.getenv("MAX_TGT", "96"))12 13def _find_model_root(base: Path) -> Path | None:14    """15    Find a directory under `base` that contains at least:16      - config.json17      - weights (pytorch_model.bin or model.safetensors)18      - tokenizer files (tokenizer.json or spiece.model or merges.txt)19    """20    candidates = []21    for cfg in base.rglob("config.json"):22        root = cfg.parent23        has_weights = any((root / f).exists() for f in ["pytorch_model.bin", "model.safetensors"])24        has_tok = any((root / f).exists() for f in ["tokenizer.json", "spiece.model", "merges.txt", "vocab.json"])25        if has_weights and has_tok:26            candidates.append(root)27    if not candidates:28        return None29    # prefer the deepest (most specific) candidate30    candidates.sort(key=lambda p: len(p.as_posix().split("/")))31    return candidates[-1]32 33def ensure_model_ready() -> Path:34    p_zip = Path(MODEL_ZIP)35    p_dir = Path(MODEL_DIR)36 37    # If dir exists but is clearly not a model root, we’ll still search inside later38    if (not p_dir.exists()) and p_zip.exists():39        p_dir.mkdir(parents=True, exist_ok=True)40        with zipfile.ZipFile(p_zip, "r") as zf:41            zf.extractall(p_dir)42 43    # If neither exists, bail with a helpful message44    if not p_dir.exists():45        raise FileNotFoundError(46            f"Model not found. Upload '{MODEL_ZIP}' or set MODEL_DIR to a folder containing the model."47        )48 49    # Try to use MODEL_DIR directly; if not a valid model root, search inside it50    if (p_dir / "config.json").exists():51        root = _find_model_root(p_dir) or p_dir52    else:53        root = _find_model_root(p_dir)54 55    if root is None:56        raise FileNotFoundError(57            f"Could not find a valid model folder under '{p_dir}'. "58            "Expected to locate 'config.json' plus weights and tokenizer files. "59            "Fix by (A) setting MODEL_DIR to the nested folder that contains these files, or "60            "(B) re-zipping so config.json is at the zip root."61        )62    return root63 64# Lazy globals65_tokenizer = None66_model = None67 68def load_model():69    global _tokenizer, _model70    if _model is not None:71        return _tokenizer, _model72 73    model_path = ensure_model_ready()74    _tokenizer = AutoTokenizer.from_pretrained(model_path)75    _model = AutoModelForSeq2SeqLM.from_pretrained(76        model_path,77        torch_dtype=torch.float32,78        low_cpu_mem_usage=True79    )80    _model.eval()81    return _tokenizer, _model82 83 84@torch.inference_mode()85def predict(text, num_beams, length_penalty, no_repeat_ngram_size, max_new_tokens, normalize):86    if not text or not text.strip():87        return ""88 89    tokenizer, model = load_model()90 91    # Build input92    inp = PREFIX + text.strip()93    enc = tokenizer(94        inp,95        return_tensors="pt",96        truncation=True,97        max_length=MAX_SRC98    )99 100    out = model.generate(101        **enc,102        max_length=max(MAX_TGT, max_new_tokens),103        num_beams=int(num_beams),104        length_penalty=float(length_penalty),105        no_repeat_ngram_size=int(no_repeat_ngram_size)106    )107    pred = tokenizer.decode(out[0], skip_special_tokens=True)108 109    if normalize:110        # Simple normalization useful for gloss (customize as you like)111        import re112        s = pred.strip().upper()113        s = s.replace("-", " ")114        s = re.sub(r"[^A-Z0-9 _]", " ", s)115        s = re.sub(r"\s+", " ", s).strip()116        pred = s117 118    return pred119 120with gr.Blocks(theme=gr.themes.Default()) as demo:121    gr.Markdown("# Text → Gloss (Seq2Seq) • CPU\nLoad your fine-tuned model_best.zip and generate gloss.")122    with gr.Row():123        with gr.Column(scale=2):124            inp = gr.Textbox(label="Input text", lines=3, placeholder="Type a sentence…")125            with gr.Accordion("Decoding & Output Options", open=False):126                beams = gr.Slider(1, 8, value=4, step=1, label="num_beams")127                lp = gr.Slider(0.5, 1.2, value=0.9, step=0.05, label="length_penalty")128                ngram = gr.Slider(1, 5, value=2, step=1, label="no_repeat_ngram_size")129                max_new = gr.Slider(16, 128, value=MAX_TGT, step=8, label="max_new_tokens")130                norm = gr.Checkbox(value=False, label="Normalize output (UPPERCASE + simple cleanup)")131            btn = gr.Button("Generate")132        with gr.Column(scale=3):133            out = gr.Textbox(label="Predicted gloss", lines=3)134 135    btn.click(predict, [inp, beams, lp, ngram, max_new, norm], [out])136    gr.Examples(137        examples=[138            ["The boy will go to school tomorrow.", 4, 0.9, 2, MAX_TGT, True],139            ["I want to visit my mother next week.", 4, 0.9, 2, MAX_TGT, True],140        ],141        inputs=[inp, beams, lp, ngram, max_new, norm],142        label="Examples"143    )144 145if __name__ == "__main__":146    # queue() is recommended for Spaces and concurrency147    demo.queue().launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))148