Update app.py

app.py

@@ -1,3 +1,5 @@
+# app.py — FINAL WORKING VERSION (deploy this now)
+
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from peft import PeftModel
@@ -8,102 +10,90 @@ import os
 
 # === HF Login ===
 hf_token = os.environ.get("HF_TOKEN")
-if hf_token:
-    login(token=hf_token)
-else:
-    raise ValueError("HF_TOKEN not set! Add it as a Space secret.")
-
-# === MODEL CONFIG ===
-# You currently have ONLY the LoRA adapter uploaded
-# So we load the base model first, then apply your LoRA on top
+if not hf_token:
+    raise ValueError("Add HF_TOKEN as a secret in your Space!")
+login(token=hf_token)
+
+# === MODEL ===
 BASE_MODEL = "Sunbird/translate-nllb-1.3b-salt"
-LORA_ADAPTER = "KMayanja/sunbird-medical-luganda-bidirectional"   # ← your repo
+LORA_ADAPTER = "KMayanja/sunbird-medical-luganda-bidirectional"
 
+print("Downloading models...")
 snapshot_download(repo_id=BASE_MODEL, token=hf_token)
 snapshot_download(repo_id=LORA_ADAPTER, token=hf_token)
 
-print("Loading tokenizer and base model...")
+print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
 
-print("Loading base model (this takes ~15 seconds)...")
+print("Loading base model...")
 base_model = AutoModelForSeq2SeqLM.from_pretrained(
     BASE_MODEL,
-    torch_dtype=torch.float32,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     low_cpu_mem_usage=True,
     trust_remote_code=True
 )
 
-print("Applying your medical LoRA adapter...")
+print("Applying your LoRA adapter...")
 model = PeftModel.from_pretrained(base_model, LORA_ADAPTER)
-
-# DO NOT .to(device) here — @spaces.GPU will handle it automatically
 model.eval()
-print("Model ready! (LoRA successfully applied)")
-
-# === LANGUAGE CODES (correct FLORES-200 codes) ===
-supported_langs = ["eng_Latn", "lug_Latn"]
-lang_names = {"eng_Latn": "English", "lug_Latn": "Luganda"}
-
-# === FALLBACK TO OLD CODE (just uncomment if you ever need it) ===
-"""
-# model_name = "Sunbird/translate-nllb-1.3b-salt"
-# tokenizer = NllbTokenizer.from_pretrained(model_name)
-# model = M2M100ForConditionalGeneration.from_pretrained(model_name)
-# language_tokens = {'eng': 256047, 'lug': 256110, ...}
-"""
-
-# === TRANSLATION FUNCTION (GPU → CPU auto-fallback via @spaces.GPU) ===
-@spaces.GPU(duration=180)  # 3-minute GPU, then falls back to CPU
-def translate(text, source_language="eng_Latn", target_language="lug_Latn"):
-    if not text.strip():
-        return "Please enter some text."
 
-    tokenizer.src_lang = source_language
-    tokenizer.tgt_lang = target_language
+# === FIXED: Correct way to get language token IDs (works with fast tokenizer) ===
+def get_lang_id(lang_code: str) -> int:
+    return tokenizer.convert_tokens_to_ids(lang_code)
+
+print("Model ready on:", "GPU" if torch.cuda.is_available() else "CPU")
+
+# === Translation function ===
+@spaces.GPU(duration=180)
+def translate(text, src="eng_Latn", tgt="lug_Latn"):
+    if not text.strip():
+        return "Please enter text to translate."
 
+    tokenizer.src_lang = src          # only needed for some NLLB versions
     inputs = tokenizer(
         text,
         return_tensors="pt",
         padding=True,
         truncation=True,
         max_length=512
-    ).to(model.device)  # automatically uses GPU or CPU
+    ).to(model.device)
 
     with torch.no_grad():
         generated = model.generate(
             **inputs,
-            forced_bos_token_id=tokenizer.lang_code_to_id[target_language],
+            forced_bos_token_id=get_lang_id(tgt),   # ← FIXED LINE
             max_length=512,
             num_beams=5,
             early_stopping=True,
-            no_repeat_ngram_size=3
+            no_repeat_ngram_size=3,
+            repetition_penalty=1.1
         )
 
     return tokenizer.decode(generated[0], skip_special_tokens=True)
 
 
-# === GRADIO INTERFACE ===
-iface = gr.Interface(
-    fn=translate,
-    inputs=[
-        gr.Textbox(label="Text to translate", lines=5, placeholder="Enter medical text..."),
-        gr.Dropdown(choices=supported_langs, value="eng_Latn", label="Source Language"),
-        gr.Dropdown(choices=supported_langs, value="lug_Latn", label="Target Language"),
-    ],
-    outputs=gr.Textbox(label="Translation", lines=5),
-    title="Uganda Medical Translator (English ↔ Luganda)",
-    description="""
-    **Best available medical translator for Luganda** — fine-tuned on 6.8k high-quality medical sentences.<br>
-    Trained by KMayanja using Sunbird 1.3B + LoRA.<br>
-    BLEU ≈ 20, chrF ≈ 36 (excellent real-world quality despite low-resource metrics).
-    """,
-    examples=[
-        ["The patient has severe malaria and needs immediate artesunate.", "eng_Latn", "lug_Latn"],
-        ["Take 2 tablets three times daily after meals.", "eng_Latn", "lug_Latn"],
+# === Gradio UI ===
+with gr.Blocks(title="Medical Translator") as iface:
+    gr.Markdown("# Uganda Medical Translator (English ↔ Luganda)")
+    gr.Markdown("**Luganda medical model** — fine-tuned on 6.8k sentences by KMayanja")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            textbox = gr.Textbox(lines=6, label="Input Text", placeholder="Enter medical text here...")
+        with gr.Column(scale=2):
+            output = gr.Textbox(lines=6, label="Translation")
+
+    with gr.Row():
+        src_lang = gr.Dropdown(["eng_Latn", "lug_Latn"], value="eng_Latn", label="Source Language")
+        tgt_lang = gr.Dropdown(["lug_Latn", "eng_Latn"], value="lug_Latn", label="Target Language")
+        btn = gr.Button("Translate", variant="primary")
+
+    btn.click(translate, inputs=[textbox, src_lang, tgt_lang], outputs=output)
+
+    gr.Examples([
+        ["The patient has severe malaria and needs immediate artesunate injection.", "eng_Latn", "lug_Latn"],
+        ["Take two tablets three times daily after meals.", "eng_Latn", "lug_Latn"],
         ["Omulwadde alina omusujja ogw’ekizungu era akennyamba okunywa amazzi.", "lug_Latn", "eng_Latn"],
-    ],
-    allow_flagging="never"
-)
+    ], inputs=[textbox, src_lang, tgt_lang])
 
-if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+iface.launch(server_name="0.0.0.0", server_port=7860)