Update app.py

app.py

@@ -2,31 +2,71 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import time
+import threading
 
 # Global variables for model and tokenizer
 model = None
 tokenizer = None
+model_loading = False
+model_loaded = False
+loading_error = None
 
 def load_model():
     """Load the model and tokenizer"""
-    global model, tokenizer
+    global model, tokenizer, model_loading, model_loaded, loading_error
+    
+    model_loading = True
+    loading_error = None
     
     try:
         model_name = "UnarineLeo/nllb_eng_ven_terms"
         print(f"Loading model: {model_name}")
         
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        # Try loading with different configurations
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = AutoModelForSeq2SeqLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                device_map="auto" if torch.cuda.is_available() else None
+            )
+        except Exception as e1:
+            print(f"First attempt failed: {e1}")
+            # Fallback: try without optimizations
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        
+        # Test if model works
+        test_input = tokenizer("Hello", return_tensors="pt")
+        with torch.no_grad():
+            _ = model.generate(**test_input, max_length=10)
         
+        model_loaded = True
+        model_loading = False
         print("Model loaded successfully!")
         return True
+        
     except Exception as e:
+        loading_error = str(e)
+        model_loading = False
+        model_loaded = False
         print(f"Error loading model: {e}")
         return False
 
+def get_model_status():
+    """Get current model loading status"""
+    if model_loaded:
+        return "✅ Model loaded and ready"
+    elif model_loading:
+        return "⏳ Model is loading, please wait..."
+    elif loading_error:
+        return f"❌ Model loading failed: {loading_error}"
+    else:
+        return "⏳ Initializing model..."
+
 def translate_text(text, max_length=512, num_beams=5):
     """
-    Translate English text to Venda
+    Translate English text to Venda using the fine-tuned NLLB model
     
     Args:
         text (str): Input English text
@@ -36,45 +76,88 @@ def translate_text(text, max_length=512, num_beams=5):
     Returns:
         tuple: (translated_text, status_message)
     """
-    global model, tokenizer
+    global model, tokenizer, model_loaded, model_loading
     
     if not text.strip():
         return "", "Please enter some text to translate."
     
-    if model is None or tokenizer is None:
-        return "", "Model not loaded. Please wait while the model loads."
+    if not model_loaded:
+        if model_loading:
+            return "", "⏳ Model is still loading, please wait a moment and try again."
+        else:
+            return "", f"❌ Model not available. {loading_error if loading_error else 'Please refresh the page.'}"
     
     try:
-        # Set source language
-        tokenizer.src_lang = "eng_Latn"
+        # Language codes as used in training
+        source_lang = "eng_Latn"
+        target_lang = "ven_Latn"
+        
+        # Format input exactly like in training: "eng_Latn: {text}"
+        formatted_input = f"{source_lang}: {text}"
+        
+        # Set source language for tokenizer
+        if hasattr(tokenizer, 'src_lang'):
+            tokenizer.src_lang = source_lang
         
         # Tokenize input
-        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+        inputs = tokenizer(
+            formatted_input, 
+            return_tensors="pt", 
+            padding=True, 
+            truncation=True, 
+            max_length=128  # Match training max_length
+        )
         
         # Generate translation
         start_time = time.time()
         with torch.no_grad():
             generated_tokens = model.generate(
                 **inputs,
-                forced_bos_token_id=tokenizer.lang_code_to_id["ven_Latn"],
                 max_length=max_length,
                 num_beams=num_beams,
                 early_stopping=True,
-                do_sample=False
+                do_sample=False,
+                pad_token_id=tokenizer.pad_token_id if hasattr(tokenizer, 'pad_token_id') else tokenizer.eos_token_id
             )
         
         # Decode translation
-        translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+        raw_translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+        
+        # Clean up translation - remove language prefixes if present
+        translation = raw_translation
+        
+        # Remove source language prefix if it appears in output
+        if translation.startswith(f"{source_lang}:"):
+            translation = translation[len(f"{source_lang}:"):].strip()
+        
+        # Remove target language prefix if it appears in output
+        if translation.startswith(f"{target_lang}:"):
+            translation = translation[len(f"{target_lang}:"):].strip()
+        
+        # Remove original input if it appears at the start
+        if translation.lower().startswith(text.lower()):
+            translation = translation[len(text):].strip()
+        
+        # Remove any remaining colons or prefixes at the start
+        translation = translation.lstrip(': ')
         
         end_time = time.time()
         processing_time = round(end_time - start_time, 2)
         
-        status = f"✅ Translation completed in {processing_time} seconds"
+        if translation and translation != formatted_input:
+            status = f"✅ Translation completed in {processing_time} seconds"
+        else:
+            status = "⚠️ Translation completed but result may be incomplete"
+            if not translation:
+                translation = "[No translation generated]"
         
         return translation, status
         
     except Exception as e:
         error_msg = f"❌ Translation error: {str(e)}"
+        print(f"Translation error: {e}")
+        import traceback
+        print(f"Full traceback: {traceback.format_exc()}")
         return "", error_msg
 
 def translate_batch(text_list):
@@ -116,9 +199,11 @@ def translate_batch(text_list):
     except Exception as e:
         return "", f"❌ Batch translation error: {str(e)}"
 
-# Load model on startup
+# Start loading model in background thread
 print("Initializing model...")
-model_loaded = load_model()
+loading_thread = threading.Thread(target=load_model)
+loading_thread.daemon = True
+loading_thread.start()
 
 # Create Gradio interface
 with gr.Blocks(title="English to Venda Translator", theme=gr.themes.Soft()) as demo:
@@ -132,6 +217,21 @@ with gr.Blocks(title="English to Venda Translator", theme=gr.themes.Soft()) as d
     **Model:** `UnarineLeo/nllb_eng_ven_terms`
     """)
     
+    # Model status indicator
+    status_indicator = gr.Textbox(
+        value=get_model_status(),
+        label="Model Status",
+        interactive=False,
+        max_lines=1
+    )
+    
+    # Auto-refresh status every 3 seconds while loading
+    def update_status():
+        return get_model_status()
+    
+    # Set up periodic status updates
+    demo.load(lambda: get_model_status(), outputs=status_indicator, every=3)
+    
     with gr.Tab("Single Translation"):
         with gr.Row():
             with gr.Column():
@@ -175,20 +275,22 @@ with gr.Blocks(title="English to Venda Translator", theme=gr.themes.Soft()) as d
                     lines=1
                 )
         
-        # Examples
+        # Examples based on statistical terminology the model was trained on
         gr.Examples(
             examples=[
-                ["Hello, how are you?"],
-                ["Good morning, everyone."],
-                ["Thank you for your help."],
-                ["What is your name?"],
-                ["I am learning Venda."],
-                ["Welcome to our school."],
-                ["The weather is beautiful today."],
-                ["Can you help me please?"]
+                ["Area planted for grain"],
+                ["Population census"],
+                ["Economic growth rate"],
+                ["Statistical survey"],
+                ["Data collection"],
+                ["Income distribution"],
+                ["Employment rate"],
+                ["Agricultural production"],
+                ["Household size"],
+                ["Rural development"]
             ],
             inputs=[input_text],
-            label="Try these examples:"
+            label="Try these statistical terms (model was trained on statistical terminology):"
         )
     
     with gr.Tab("Batch Translation"):
@@ -263,6 +365,13 @@ with gr.Blocks(title="English to Venda Translator", theme=gr.themes.Soft()) as d
         inputs=[input_text, max_length_slider, num_beams_slider],
         outputs=[output_text, status_text]
     )
+    
+    # Refresh status button
+    refresh_btn = gr.Button("🔄 Refresh Status", size="sm")
+    refresh_btn.click(
+        fn=update_status,
+        outputs=[status_indicator]
+    )
 
 # Launch the app
 if __name__ == "__main__":