Update app.py

app.py

@@ -5,18 +5,21 @@ import requests
 import gradio as gr
 from langdetect import detect, LangDetectException
 
-# Groq SDK (install in requirements.txt)
+# Hugging Face Transformers
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+
+# Groq SDK
 try:
     from groq import Groq
 except Exception:
     Groq = None
 
-# Config (set these in your Space secrets)
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")         # your Groq key
+# Config
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")         
 GROQ_MODEL = os.getenv("GROQ_MODEL", "mixtral-8x7b-32768")
-HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")  # hf_... token (optional but recommended)
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 
-# Init clients (if keys provided)
+# Init Groq
 groq_client = None
 if GROQ_API_KEY and Groq is not None:
     try:
@@ -24,95 +27,62 @@ if GROQ_API_KEY and Groq is not None:
     except Exception as e:
         print("Groq client init failed:", repr(e))
 
-# Simple mapping for UI -> code
+# Universal translation model
+m2m_model_name = "facebook/m2m100_418M"
+m2m_tokenizer = M2M100Tokenizer.from_pretrained(m2m_model_name)
+m2m_model = M2M100ForConditionalGeneration.from_pretrained(m2m_model_name)
+
+# UI mapping
 LANG_UI_TO_CODE = {"English": "en", "Spanish": "es", "French": "fr"}
 
 SYSTEM_PROMPT = """
 You are a multilingual translation assistant.
-Supported languages: English, Spanish, French.
-
 Task:
 1. Detect the input language automatically.
-2. Translate the input into the exact target language requested by the user.
+2. Translate into the requested target language.
 3. Preserve meaning, tone, and formatting.
-4. Keep numbers, symbols, names and special characters unchanged.
+4. Keep numbers, symbols, names, and special characters unchanged.
 5. If the input is already in the target language, return it unchanged.
-6. Output ONLY the translated text, no commentary.
 """
 
 def call_groq(user_text, target_lang_ui):
     if not groq_client:
         raise RuntimeError("Groq client not configured")
-    # Put the target explicitly for determinism
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": f"Target language: {target_lang_ui}\n\n{user_text}"},
     ]
-    # Best-effort: extract content from different response shapes
     chat = groq_client.chat.completions.create(
         model=GROQ_MODEL,
         messages=messages,
         temperature=0,
         max_tokens=2048,
     )
-    # SDK usually returns .choices[0].message.content
     try:
         return chat.choices[0].message.content.strip()
     except Exception:
-        # fallback for dict-like response
         try:
             return chat["choices"][0]["message"]["content"].strip()
         except Exception as e:
-            print("Unexpected Groq response shape:", repr(e))
+            print("Unexpected Groq response:", repr(e))
             raise
 
-def call_hf_opus(user_text, src_code, tgt_code):
-    # map most common pairs to explicit model ids
-    model_map = {
-        ("en","es"): "Helsinki-NLP/opus-mt-en-es",
-        ("es","en"): "Helsinki-NLP/opus-mt-es-en",
-        ("en","fr"): "Helsinki-NLP/opus-mt-en-fr",
-        ("fr","en"): "Helsinki-NLP/opus-mt-fr-en",
-        ("es","fr"): "Helsinki-NLP/opus-mt-es-fr",
-        ("fr","es"): "Helsinki-NLP/opus-mt-fr-es",
-    }
-    model_id = model_map.get((src_code, tgt_code)) or f"Helsinki-NLP/opus-mt-{src_code}-{tgt_code}"
-
-    url = f"https://api-inference.huggingface.co/models/{model_id}"
-    headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"} if HUGGINGFACE_TOKEN else {}
-    payload = {"inputs": user_text}
-
-    resp = requests.post(url, headers=headers, json=payload, timeout=30)
-    if resp.status_code != 200:
-        raise RuntimeError(f"Hugging Face fallback failed: {resp.status_code} {resp.text}")
-
-    data = resp.json()
-    # handle common response shapes:
-    if isinstance(data, list) and len(data) > 0:
-        first = data[0]
-        if isinstance(first, dict):
-            # prefer common keys
-            for k in ("translation_text", "generated_text", "text"):
-                if k in first:
-                    return first[k]
-            # otherwise return first value
-            return next(iter(first.values()))
-        else:
-            return str(first)
-    if isinstance(data, dict):
-        for k in ("translation_text", "generated_text", "text"):
-            if k in data:
-                return data[k]
-        return json.dumps(data)
-    return str(data)
-
-def detect_lang_code(text):
+def call_m2m(user_text, target_code):
     try:
-        code = detect(text)
-        # keep only en/es/fr; if another, default to 'en' for fallback routing
-        return code if code in ("en","es","fr") else "en"
+        src_code = detect(user_text)
     except LangDetectException:
-        return "en"
+        src_code = "en"  # fallback
+
+    # if already target language → return as-is
+    if src_code == target_code:
+        return user_text
+
+    m2m_tokenizer.src_lang = src_code
+    encoded = m2m_tokenizer(user_text, return_tensors="pt")
+    generated = m2m_model.generate(
+        **encoded, forced_bos_token_id=m2m_tokenizer.get_lang_id(target_code)
+    )
+    return m2m_tokenizer.decode(generated[0], skip_special_tokens=True)
 
 def translate_text(user_text, target_lang_ui):
     user_text = (user_text or "").strip()
@@ -120,7 +90,7 @@ def translate_text(user_text, target_lang_ui):
         return "⚠️ Please enter some text to translate."
     target_code = LANG_UI_TO_CODE.get(target_lang_ui, "en")
 
-    # First: try Groq (preferred)
+    # Try Groq first
     try:
         if groq_client:
             out = call_groq(user_text, target_lang_ui)
@@ -129,21 +99,16 @@ def translate_text(user_text, target_lang_ui):
     except Exception as e:
         print("Groq call failed:", repr(e))
 
-    # Fallback: detect source and call HF OPUS-MT models
+    # Fallback → M2M100 universal translator
     try:
-        src_code = detect_lang_code(user_text)
-        if src_code == target_code:
-            # already same language
-            return user_text
-        out = call_hf_opus(user_text, src_code, target_code)
-        return out.strip()
+        return call_m2m(user_text, target_code)
     except Exception as e:
-        print("Hugging Face fallback failed:", repr(e))
-        return f"Error: translation failed. Check Space logs for details."
+        print("M2M100 translation failed:", repr(e))
+        return "❌ Translation failed. Check logs."
 
 # ----------------- Gradio UI -----------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 🌐 Hackathon Translator (EN/ES/FR)")
+    gr.Markdown("## 🌐 Hackathon Translator (Universal)")
 
     with gr.Row():
         txt = gr.Textbox(label="Enter your text", lines=6, placeholder="Type or paste text here...")