Update src/streamlit_app.py

src/streamlit_app.py

@@ -18,7 +18,7 @@ if not HF_TOKEN:
     st.stop()
 
 client = InferenceClient(
-    model="google-t5/t5-base",   # ✅ FAST GEMMA MODEL
+    model="google-t5/t5-base",
     token=HF_TOKEN
 )
 
@@ -68,38 +68,42 @@ def speech_to_text(audio_file):
     return text
 
 # ==============================
-# TRANSLATION FUNCTION (GEMMA)
+# LANGUAGE MAP (IMPORTANT)
+# ==============================
+lang_map = {
+    "English": "en",
+    "Tamil": "ta",
+    "Hindi": "hi",
+    "French": "fr",
+    "Arabic": "ar",
+    "Spanish": "es",
+    "German": "de"
+}
+
+# ==============================
+# TRANSLATION FUNCTION (FIXED)
 # ==============================
 def translate(text, target_lang):
 
+    # Detect source language
     try:
         src_lang = detect(text)
     except:
-        src_lang = "unknown"
-
-    memory = retrieve_memory(text)
-
-    prompt = f"""
-You are an expert translator.
+        src_lang = "en"
 
-Translate the following text into {target_lang}.
+    tgt_lang = lang_map[target_lang]
 
-Text:
-{text}
-
-IMPORTANT:
-- Only return translated text
-- Do not explain anything
-"""
+    memory = retrieve_memory(text)
 
     try:
-        response = client.text_generation(
-            prompt,
-            max_new_tokens=200,
-            temperature=0.2
+        # ✅ Correct API for T5
+        response = client.translation(
+            text=text,
+            src_lang=src_lang,
+            tgt_lang=tgt_lang
         )
 
-        translated = response.strip()
+        translated = response["translation_text"]
 
     except Exception as e:
         translated = f"❌ Translation failed: {str(e)}"
@@ -113,17 +117,21 @@ IMPORTANT:
 # ==============================
 st.set_page_config(page_title="AI Translator", layout="wide")
 
-st.title("🌍 AI Translator with Voice")
+st.title("🌍 AI Translator with Voice (T5)")
 
 tab1, tab2 = st.tabs(["📝 Text Input", "🎤 Voice Input"])
 
 input_text = ""
 
+# ------------------------------
 # TEXT INPUT
+# ------------------------------
 with tab1:
     input_text = st.text_area("Enter text")
 
+# ------------------------------
 # VOICE INPUT
+# ------------------------------
 with tab2:
     audio_file = st.file_uploader("Upload audio (wav/mp3)", type=["wav", "mp3"])
 
@@ -136,13 +144,17 @@ with tab2:
                 st.success("Recognized Text:")
                 st.write(input_text)
 
+# ------------------------------
 # TARGET LANGUAGE
+# ------------------------------
 target_lang = st.selectbox(
     "Target Language",
-    ["English", "Tamil", "Hindi", "French", "Arabic", "Spanish", "German"]
+    list(lang_map.keys())
 )
 
-# TRANSLATE
+# ------------------------------
+# TRANSLATE BUTTON
+# ------------------------------
 if st.button("Translate"):
     if not input_text.strip():
         st.warning("Please provide input text or audio")