Update app.py

app.py

@@ -12,7 +12,7 @@ import os
 import requests
 import torch
 from flask import Flask, request, render_template_string, jsonify
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
 # Configuration
@@ -35,17 +35,25 @@ tokenizer = None
 device = None
 
 def load_model():
-    """Load MarianMT model once and cache it"""
+    """Load MarianMT model once and cache it using Auto classes"""
     global model, tokenizer, device
     
     if model is None:
         print("Loading MarianMT model...")
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         
-        tokenizer = MarianTokenizer.from_pretrained(MODEL_ID)
-        model = MarianMTModel.from_pretrained(MODEL_ID).to(device).eval()
-        
-        print(f"Model loaded on {device}")
+        try:
+            # Try Auto classes first (more flexible)
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
+            model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID).to(device).eval()
+            print(f"Model loaded successfully on {device} using Auto classes")
+        except Exception as e:
+            print(f"Auto classes failed: {e}")
+            print("Trying legacy Marian classes...")
+            from transformers import MarianMTModel, MarianTokenizer
+            tokenizer = MarianTokenizer.from_pretrained(MODEL_ID, use_fast=False)
+            model = MarianMTModel.from_pretrained(MODEL_ID).to(device).eval()
+            print(f"Model loaded on {device} using legacy classes")
     
     return model, tokenizer, device
 
@@ -57,6 +65,7 @@ def translate_marian(text):
     try:
         model, tokenizer, device = load_model()
         
+        # Prepare inputs
         inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)
         inputs = {k: v.to(device) for k, v in inputs.items()}
         
@@ -74,7 +83,7 @@ def translate_marian(text):
         translations = []
         for output in outputs:
             trans = tokenizer.decode(output, skip_special_tokens=True)
-            if trans not in translations:  # Avoid duplicates
+            if trans and trans not in translations:  # Avoid duplicates and empty
                 translations.append(trans)
         
         return translations if translations else ["[Error: No translation generated]"]
@@ -83,7 +92,7 @@ def translate_marian(text):
         print(f"MarianMT translation error: {e}")
         import traceback
         traceback.print_exc()
-        return [f"[Error: {str(e)}"]
+        return [f"[Error: {str(e)}]"]
 
 def translate_libre_variant(text, variant_code):
     """Translate using a specific LibreTranslate variant"""
@@ -597,28 +606,20 @@ HTML_TEMPLATE = """
         let selectedText = '';
         
         function selectTranslation(element, text) {
-            // Remove selected class from all items in this list
             element.parentElement.querySelectorAll('.translation-item').forEach(item => {
                 item.classList.remove('selected');
             });
-            // Add selected class to clicked item
             element.classList.add('selected');
             selectedText = text;
-            
-            // Copy automatically on selection
             copyText(text, false);
         }
         
         function selectVariant(element, text) {
-            // Remove selected class from all variants
             document.querySelectorAll('.variant-item').forEach(item => {
                 item.classList.remove('selected');
             });
-            // Add selected class to clicked item
             element.classList.add('selected');
             selectedText = text;
-            
-            // Copy automatically on selection
             copyText(text, false);
         }
         
@@ -629,7 +630,6 @@ HTML_TEMPLATE = """
                     showToastMessage('Copied to clipboard!');
                 }
             } catch (err) {
-                // Fallback
                 const textArea = document.createElement('textarea');
                 textArea.value = text;
                 document.body.appendChild(textArea);
@@ -670,10 +670,7 @@ def index():
     if request.method == "POST":
         source_text = request.form.get("text", "").strip()
         if source_text:
-            # Get Marian translations
             marian = translate_marian(source_text)
-            
-            # Get Libre translations
             libre_results = translate_libre_all_variants(source_text)
             for name, data in libre_results.items():
                 data['code'] = KABYLE_VARIANTS[name]
@@ -688,13 +685,11 @@ def index():
 
 @app.route("/health")
 def health():
-    """Health check endpoint"""
     return jsonify({
         "status": "healthy",
         "model_loaded": model is not None
     })
 
 if __name__ == "__main__":
-    # Hugging Face Spaces expects port 7860
     port = int(os.environ.get("PORT", 7860))
     app.run(host="0.0.0.0", port=port, debug=False)