Updated now with langauge codes

app.py

@@ -28,13 +28,27 @@ def load_translation_models():
         print("Loading English to Siswati model...")
         en_ss_tokenizer = AutoTokenizer.from_pretrained("dsfsi/en-ss-m2m100-combo")
         en_ss_model = AutoModelForSeq2SeqLM.from_pretrained("dsfsi/en-ss-m2m100-combo")
-        en_ss_pipeline = pipeline("translation", model=en_ss_model, tokenizer=en_ss_tokenizer)
+        # Fix: Add src_lang and tgt_lang parameters
+        en_ss_pipeline = pipeline(
+            "translation", 
+            model=en_ss_model, 
+            tokenizer=en_ss_tokenizer,
+            src_lang="en",
+            tgt_lang="ss"
+        )
         
         # Siswati to English
         print("Loading Siswati to English model...")
         ss_en_tokenizer = AutoTokenizer.from_pretrained("dsfsi/ss-en-m2m100-combo")
         ss_en_model = AutoModelForSeq2SeqLM.from_pretrained("dsfsi/ss-en-m2m100-combo")
-        ss_en_pipeline = pipeline("translation", model=ss_en_model, tokenizer=ss_en_tokenizer)
+        # Fix: Add src_lang and tgt_lang parameters
+        ss_en_pipeline = pipeline(
+            "translation", 
+            model=ss_en_model, 
+            tokenizer=ss_en_tokenizer,
+            src_lang="ss",
+            tgt_lang="en"
+        )
         
         # Cache the models
         _model_cache['en_ss_pipeline'] = en_ss_pipeline
@@ -56,6 +70,77 @@ def get_translators():
     
     return _model_cache['en_ss_pipeline'], _model_cache['ss_en_pipeline']
 
+def translate_with_fallback(text, direction):
+    """Translation function with fallback method if pipeline fails"""
+    try:
+        # Get translators
+        en_ss_translator, ss_en_translator = get_translators()
+        
+        if direction == "English → Siswati":
+            if en_ss_translator is None:
+                raise Exception("English to Siswati model not loaded")
+            
+            # Try with pipeline first
+            try:
+                result = en_ss_translator(text, max_length=512)
+                return result[0]['translation_text']
+            except Exception as pipeline_error:
+                print(f"Pipeline failed, trying direct model approach: {pipeline_error}")
+                
+                # Fallback: Use model directly
+                tokenizer = AutoTokenizer.from_pretrained("dsfsi/en-ss-m2m100-combo")
+                model = AutoModelForSeq2SeqLM.from_pretrained("dsfsi/en-ss-m2m100-combo")
+                
+                # Set language tokens
+                tokenizer.src_lang = "en"
+                encoded = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
+                
+                # Force target language token
+                forced_bos_token_id = tokenizer.get_lang_id("ss")
+                
+                with torch.no_grad():
+                    generated_tokens = model.generate(
+                        **encoded,
+                        forced_bos_token_id=forced_bos_token_id,
+                        max_length=512
+                    )
+                
+                return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+                
+        else:  # Siswati → English
+            if ss_en_translator is None:
+                raise Exception("Siswati to English model not loaded")
+            
+            # Try with pipeline first
+            try:
+                result = ss_en_translator(text, max_length=512)
+                return result[0]['translation_text']
+            except Exception as pipeline_error:
+                print(f"Pipeline failed, trying direct model approach: {pipeline_error}")
+                
+                # Fallback: Use model directly
+                tokenizer = AutoTokenizer.from_pretrained("dsfsi/ss-en-m2m100-combo")
+                model = AutoModelForSeq2SeqLM.from_pretrained("dsfsi/ss-en-m2m100-combo")
+                
+                # Set language tokens
+                tokenizer.src_lang = "ss"
+                encoded = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
+                
+                # Force target language token
+                forced_bos_token_id = tokenizer.get_lang_id("en")
+                
+                with torch.no_grad():
+                    generated_tokens = model.generate(
+                        **encoded,
+                        forced_bos_token_id=forced_bos_token_id,
+                        max_length=512
+                    )
+                
+                return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+                
+    except Exception as e:
+        raise Exception(f"Translation failed: {str(e)}")
+
 def analyze_siswati_features(text):
     """Analyze Siswati-specific linguistic features"""
     features = {}
@@ -117,32 +202,17 @@ def translate_text(text, direction):
     start_time = time.time()
     
     try:
-        # Get translators (will load if not cached)
-        en_ss_translator, ss_en_translator = get_translators()
+        # Perform translation using the fallback method
+        translated_text = translate_with_fallback(text, direction)
+        
+        # Analyze source and target text
+        source_metrics = calculate_linguistic_metrics(text)
+        target_metrics = calculate_linguistic_metrics(translated_text)
         
-        # Perform translation
+        # Analyze Siswati features based on direction
         if direction == "English → Siswati":
-            if en_ss_translator is None:
-                return "Translation model not loaded. Please try again.", "Model loading failed.", create_empty_metrics_table()
-            
-            result = en_ss_translator(text, max_length=512)
-            translated_text = result[0]['translation_text']
-            
-            # Analyze source (English) and target (Siswati)
-            source_metrics = calculate_linguistic_metrics(text)
-            target_metrics = calculate_linguistic_metrics(translated_text)
             siswati_features = analyze_siswati_features(translated_text)
-            
-        else:  # Siswati → English
-            if ss_en_translator is None:
-                return "Translation model not loaded. Please try again.", "Model loading failed.", create_empty_metrics_table()
-            
-            result = ss_en_translator(text, max_length=512)
-            translated_text = result[0]['translation_text']
-            
-            # Analyze source (Siswati) and target (English)
-            source_metrics = calculate_linguistic_metrics(text)
-            target_metrics = calculate_linguistic_metrics(translated_text)
+        else:
             siswati_features = analyze_siswati_features(text)
         
         processing_time = time.time() - start_time
@@ -279,23 +349,9 @@ def secure_file_processing(file_obj, direction):
             if len(text) > 1000:
                 text = text[:1000] + "..."
             
-            # Get translators for batch processing
-            en_ss_translator, ss_en_translator = get_translators()
-            
-            # Perform translation based on direction
+            # Perform translation using the fallback method
             try:
-                if direction == "English → Siswati":
-                    if en_ss_translator is None:
-                        translated = "Model not available"
-                    else:
-                        result = en_ss_translator(text, max_length=512)
-                        translated = result[0]['translation_text']
-                else:  # Siswati → English
-                    if ss_en_translator is None:
-                        translated = "Model not available"
-                    else:
-                        result = ss_en_translator(text, max_length=512)
-                        translated = result[0]['translation_text']
+                translated = translate_with_fallback(text, direction)
             except Exception as e:
                 translated = f"Translation error: {str(e)}"
             
@@ -356,7 +412,6 @@ def create_gradio_interface():
         # Header Section
         gr.HTML("""
         <div class="main-header">
-            <img src="https://www.dsfsi.co.za/images/logo_transparent_expanded.png" width="400" alt="DSFSI Logo" style="margin-bottom: 1rem;">
             <h1>🔬 Siswati-English Linguistic Translation Tool</h1>
             <p style="font-size: 1.1em; color: #666; max-width: 800px; margin: 0 auto;">
                 Advanced AI-powered translation system with comprehensive linguistic analysis features,