resource issue

app.py

@@ -1,23 +1,23 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoProcessor
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Language configuration with proper model handling
+# Language configuration with optimized model selection
 LANGUAGE_CONFIG = {
     "Amharic": {
         "code": "amh",
-        "model_type": "seamless",
-        "seamless_code": "amh"
+        "model_type": "nllb",
+        "nllb_code": "amh_Ethi"
     },
     "Swahili": {
         "code": "swh", 
-        "model_type": "seamless",
-        "seamless_code": "swh"
+        "model_type": "nllb",
+        "nllb_code": "swh_Latn"
     },
     "Somali": {
         "code": "som",
-        "model_type": "seamless",
-        "seamless_code": "som"
+        "model_type": "nllb",
+        "nllb_code": "som_Latn"
     },
     "Afan Oromo": {
         "code": "gaz",
@@ -37,110 +37,74 @@ LANGUAGE_CONFIG = {
 }
 
 # Model instances
-models = {}
-tokenizers = {}
-processors = {}
+model = None
+tokenizer = None
 
-print("🚀 Initializing translation models...")
+print("🚀 Initializing translation model for Hugging Face Spaces...")
 
-# Load SeamlessM4T model for Amharic, Swahili, Somali
+# Load a smaller, more efficient NLLB model
 try:
-    print("📥 Loading SeamlessM4T model...")
-    seamless_model_id = "facebook/seamless-m4t-v2-large"
-    processors['seamless'] = AutoProcessor.from_pretrained(seamless_model_id)
-    models['seamless'] = AutoModelForSeq2SeqLM.from_pretrained(seamless_model_id)
-    print("✅ SeamlessM4T model loaded successfully!")
-except Exception as e:
-    print(f"❌ Failed to load SeamlessM4T model: {e}")
-    models['seamless'] = None
-    processors['seamless'] = None
-
-# Load NLLB model for other languages
-try:
-    print("📥 Loading NLLB model...")
-    nllb_model_id = "facebook/nllb-200-distilled-600M"
-    tokenizers['nllb'] = AutoTokenizer.from_pretrained(nllb_model_id)
-    models['nllb'] = AutoModelForSeq2SeqLM.from_pretrained(nllb_model_id)
+    print("📥 Loading NLLB-200-1.3B model...")
+    model_id = "facebook/nllb-200-1.3B"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16,  # Use half precision to save memory
+        device_map="auto"
+    )
     print("✅ NLLB model loaded successfully!")
 except Exception as e:
-    print(f"❌ Failed to load NLLB model: {e}")
-    models['nllb'] = None
-    tokenizers['nllb'] = None
-
-def translate_with_seamless(text, source_lang_code):
-    """Translate text using SeamlessM4T model"""
+    print(f"❌ Failed to load NLLB-200-1.3B: {e}")
     try:
-        if models['seamless'] is None or processors['seamless'] is None:
-            return "SeamlessM4T model not available"
-        
-        # Preprocess text
-        inputs = processors['seamless'](text=text, src_lang=source_lang_code, return_tensors="pt")
-        
-        # Get BOS token for target language (English)
-        forced_bos_token_id = processors['seamless'].tokenizer.convert_tokens_to_ids("<|eng|>")
-        
-        # Generate translation
-        with torch.no_grad():
-            generated_tokens = models['seamless'].generate(
-                **inputs,
-                forced_bos_token_id=forced_bos_token_id,
-                max_length=256
-            )
-        
-        # Decode and return
-        translation = processors['seamless'].batch_decode(generated_tokens, skip_special_tokens=True)[0]
-        return translation
-        
-    except Exception as e:
-        print(f"SeamlessM4T translation error: {e}")
-        return f"Translation failed: {str(e)[:200]}"
+        # Fallback to even smaller model
+        print("🔄 Trying smaller model: NLLB-200-distilled-600M...")
+        model_id = "facebook/nllb-200-distilled-600M"
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+        print("✅ NLLB distilled model loaded successfully!")
+    except Exception as e2:
+        print(f"❌ All models failed to load: {e2}")
+        model = None
+        tokenizer = None
 
-def translate_with_nllb(text, source_lang_code):
-    """Translate text using NLLB model"""
+def translate_text(text, source_language):
+    """Main translation function"""
+    if not text.strip():
+        return "Please enter text to translate"
+    
+    if source_language not in LANGUAGE_CONFIG:
+        return f"Translation for {source_language} is not supported"
+    
+    if model is None or tokenizer is None:
+        return "Translation model is not available. Please try again later."
+    
+    config = LANGUAGE_CONFIG[source_language]
+    
     try:
-        if models['nllb'] is None or tokenizers['nllb'] is None:
-            return "NLLB model not available"
-        
         # Tokenize input
-        inputs = tokenizers['nllb'](text, return_tensors="pt")
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+        
+        # Move to same device as model
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
         
         # Define target language (English)
-        forced_bos_token_id = tokenizers['nllb'].convert_tokens_to_ids("eng_Latn")
+        forced_bos_token_id = tokenizer.convert_tokens_to_ids("eng_Latn")
         
-        # Generate translation using beam search for better quality
+        # Generate translation with optimized settings for HF Spaces
         with torch.no_grad():
-            generated_tokens = models['nllb'].generate(
+            generated_tokens = model.generate(
                 **inputs,
                 forced_bos_token_id=forced_bos_token_id,
                 max_length=256,
-                num_beams=5,
-                early_stopping=True
+                num_beams=3,  # Reduced for faster inference
+                early_stopping=True,
+                no_repeat_ngram_size=2
             )
         
         # Decode
-        translation = tokenizers['nllb'].batch_decode(generated_tokens, skip_special_tokens=True)[0]
+        translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
         return translation
         
-    except Exception as e:
-        print(f"NLLB translation error: {e}")
-        return f"Translation failed: {str(e)[:200]}"
-
-def translate_text(text, source_language):
-    """Main translation function"""
-    if not text.strip():
-        return "Please enter text to translate"
-    
-    if source_language not in LANGUAGE_CONFIG:
-        return f"Translation for {source_language} is not supported"
-    
-    config = LANGUAGE_CONFIG[source_language]
-    
-    try:
-        if config["model_type"] == "seamless":
-            return translate_with_seamless(text, config["seamless_code"])
-        else:  # nllb
-            return translate_with_nllb(text, config["nllb_code"])
-            
     except Exception as e:
         print(f"Translation error for {source_language}: {e}")
         return f"Translation failed: {str(e)[:200]}"
@@ -155,17 +119,18 @@ EXAMPLE_TEXTS = {
     "Chichewa": "Alipo wina aliyense ali ndi ufulu wachibadwidwe."
 }
 
-# Test the models on startup
-def test_models():
-    print("🧪 Testing translation models...")
+# Test the model on startup
+def test_model():
+    if model is None:
+        print("❌ No model available for testing")
+        return
+    
+    print("🧪 Testing translation model...")
     
     test_cases = [
         ("Swahili", "Habari za asubuhi"),
         ("Somali", "Maanta waa maalin fiican"),
         ("Amharic", "ሰላም"),
-        ("Afan Oromo", "Akkam jirta"),
-        ("Tigrinya", "ሰላም"),
-        ("Chichewa", "Moni")
     ]
     
     for lang, text in test_cases:
@@ -175,8 +140,9 @@ def test_models():
         except Exception as e:
             print(f"❌ {lang} test failed: {e}")
 
-# Run tests on startup
-test_models()
+# Run test if model is loaded
+if model is not None:
+    test_model()
 
 # Create Gradio interface
 with gr.Blocks(
@@ -261,10 +227,9 @@ with gr.Blocks(
             gr.Markdown("### 🔧 Model Information")
             
             # Create status display
-            seamless_status = "✅ Loaded" if models.get('seamless') else "❌ Failed"
-            nllb_status = "✅ Loaded" if models.get('nllb') else "❌ Failed"
+            model_status = "✅ Loaded" if model is not None else "❌ Failed to load"
             
-            status_text = f"SeamlessM4T: {seamless_status} | NLLB: {nllb_status}"
+            status_text = f"NLLB-200 Model: {model_status}"
             gr.Textbox(
                 value=status_text, 
                 label="Model Status", 
@@ -272,18 +237,16 @@ with gr.Blocks(
             )
             
             # Create model info
-            seamless_langs = [lang for lang, config in LANGUAGE_CONFIG.items() if config["model_type"] == "seamless"]
-            nllb_langs = [lang for lang, config in LANGUAGE_CONFIG.items() if config["model_type"] == "nllb"]
-            
             gr.Markdown(f"""
-            **Advanced Models (SeamlessM4T):** {', '.join(seamless_langs)}  
-            **Standard Models (NLLB-200):** {', '.join(nllb_langs)}  
+            **Supported Languages:** {', '.join(LANGUAGE_CONFIG.keys())}
+            
+            **Model:** NLLB-200 (No Language Left Behind)
             
             **Features:**
             - High-quality translations for African languages
             - Support for text input and copy-paste functionality
-            - Fast and accurate results using beam search
-            - Proper tokenization for each language family
+            - Fast and accurate results
+            - Optimized for Hugging Face Spaces
             """)
 
     # Add CSS for better styling
@@ -304,5 +267,4 @@ if __name__ == "__main__":
         server_port=7860,
         share=False,
         show_error=True
-    )
-
+    )