Update app.py

Updated app.py to use medical model finetuned from the sunbird/salt-nllb-200-1.3B

app.py

@@ -1,84 +1,120 @@
 import gradio as gr
-from transformers import NllbTokenizer, M2M100ForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import spaces
 from huggingface_hub import login, snapshot_download
 import os
 
-# Fix: Retrieve HF token from environment (set as a Space secret)
+# === HF Login ===
 hf_token = os.environ.get("HF_TOKEN")
 if hf_token:
     login(token=hf_token)
-    os.environ["HF_TOKEN"] = hf_token
 else:
-    raise ValueError("HF_TOKEN environment variable not set. Add it as a secret in your Space settings.")
+    raise ValueError("HF_TOKEN not set! Add it as a Space secret.")
 
-# Model name
-model_name = "Sunbird/translate-nllb-1.3b-salt"
+# === MODEL CONFIG ===
+model_name = "KMayanja/sunbird-medical-luganda-bidirectional"
 
-# Download the model files first to avoid issues during loading
+# Optional: cache model locally on first load
 snapshot_download(repo_id=model_name, token=hf_token)
 
-# Load the tokenizer and model once at startup
+# === LOAD TOKENIZER & MODEL ONCE AT STARTUP ===
 try:
-    tokenizer = NllbTokenizer.from_pretrained(model_name, token=hf_token)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = M2M100ForConditionalGeneration.from_pretrained(model_name, token=hf_token)
-    model.to(device)  # Fix: Move to device once here, not per request
-    model.eval()  # Set to evaluation mode for inference
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+    
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float32,  # Safe default (GPU will auto-upgrade to bfloat16 if possible)
+        low_cpu_mem_usage=True
+    )
+    
+    # Let @spaces.GPU handle device placement — do NOT move model here
+    # model.to(device)  ← removed on purpose
+    
+    model.eval()
+    print("Model loaded successfully.")
+
 except Exception as e:
     print(f"Error loading model: {e}")
     raise
 
-# Supported languages and their tokens
-language_tokens = {
-    'eng': 256047,
-    'ach': 256111,
-    'lgg': 256008,
-    'lug': 256110,
-    'nyn': 256002,
-    'teo': 256006,
+# === LANGUAGE CODES (correct ones for your fine-tuned model) ===
+# These are the official FLORES-200 codes used by Sunbird & NLLB
+lang_code_to_id = {
+    "eng_Latn": tokenizer.lang_code_to_id["eng_Latn"],
+    "lug_Latn": tokenizer.lang_code_to_id["lug_Latn"],
 }
 
-supported_languages = list(language_tokens.keys())
+supported_langs = ["eng_Latn", "lug_Latn"]
+lang_names = {"eng_Latn": "English", "lug_Latn": "Luganda"}
 
-@spaces.GPU
-def translate(text, source_language, target_language):
-    if source_language not in supported_languages:
-        raise ValueError(f"Source language '{source_language}' not supported. Supported: {supported_languages}")
-    if target_language not in supported_languages:
-        raise ValueError(f"Target language '{target_language}' not supported. Supported: {supported_languages}")
 
-    # Fix: No need to move model here—it's already on device
+# === FALLBACK: Old working code (commented out — just uncomment to revert) ===
+"""
+# model_name = "Sunbird/translate-nllb-1.3b-salt"
+# tokenizer = NllbTokenizer.from_pretrained(model_name, token=hf_token)
+# model = M2M100ForConditionalGeneration.from_pretrained(model_name, token=hf_token)
+# language_tokens = {'eng': 256047, 'lug': 256110, ...}
+"""
+
+# === MAIN TRANSLATION FUNCTION WITH GPU AUTO-FALLBACK ===
+@spaces.GPU(duration=120)  # 2 minutes GPU, then auto-fallback to CPU
+def translate(text, source_language, target_language):
+    if text.strip() == "":
+        return "Please enter text to translate."
 
-    inputs = tokenizer(text, return_tensors="pt").to(device)
-    inputs['input_ids'][0][0] = language_tokens[source_language]
+    # Set source & target language
+    tokenizer.src_lang = source_language
+    tokenizer.tgt_lang = target_language
 
-    translated_tokens = model.generate(
-        **inputs,
-        forced_bos_token_id=language_tokens[target_language],
-        max_length=100,
-        num_beams=5,
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=512
     )
 
-    result = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
-    return result
+    # Move inputs to correct device (GPU if available, else CPU)
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
-# Create Gradio interface
+    with torch.no_grad():
+        generated_ids = model.generate(
+            **inputs,
+            forced_bos_token_id=tokenizer.lang_code_to_id[target_language],
+            max_length=512,
+            num_beams=5,
+            early_stopping=True,
+            no_repeat_ngram_size=3
+        )
+
+    translation = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    return translation
+
+
+# === GRADIO INTERFACE ===
 iface = gr.Interface(
     fn=translate,
     inputs=[
-        gr.Textbox(label="Text to translate"),
-        gr.Dropdown(choices=supported_languages, label="Source language (e.g., 'eng')", value='eng'),  # Optional: Dropdown for easier UX
-        gr.Dropdown(choices=supported_languages, label="Target language (e.g., 'lug')", value='lug'),  # Optional: Dropdown for easier UX
+        gr.Textbox(label="Text to translate", lines=4, placeholder="Enter medical text here..."),
+        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=4),
+    title="Luganda Medical Translator (Sunbird 1.3B Fine-tuned)",
+    description="""
+    State-of-the-art bidirectional English ↔ Luganda medical translator.<br>
+    Trained on 6.8k high-quality medical sentences. Best available model for healthcare in Uganda.
+    """,
+    examples=[
+        ["The patient has severe malaria and needs immediate treatment.", "eng_Latn", "lug_Latn"],
+        ["Omulwadde alina omusujja ogw’ekizungu era akennyamba okunywa amazzi.", "lug_Latn", "eng_Latn"],
+        ["Take two tablets three times daily after meals.", "eng_Latn", "lug_Latn"],
     ],
-    outputs=gr.Textbox(label="Translated text"),
-    title="Test Translation API",
-    description="Translate text using Sunbird/translate-nllb-1.3b-salt model(To be replaced later). Supported languages: eng (English), lug (Luganda).",
+    allow_flagging="never"
 )
 
-# Fix: Remove share=True—HF Spaces handles this
-# Launch the application
+# === LAUNCH ===
 if __name__ == "__main__":
     iface.launch(
         server_name="0.0.0.0",