Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,194 +1,86 @@
 import streamlit as st
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from langdetect import detect
-import numpy as np
-import faiss
-import tempfile
-import speech_recognition as sr
-from sentence_transformers import SentenceTransformer
-import os
 
 # ==============================
-# CONFIG
+# LOAD MODEL (ONLY ONCE)
 # ==============================
-HF_TOKEN = os.environ.get("HF_TOKEN")
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+    model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
+    return tokenizer, model
 
-if not HF_TOKEN:
-    st.error("❌ HF_TOKEN not found. Add it in Hugging Face Secrets.")
-    st.stop()
-
-client = InferenceClient(
-    model="google/gemma-7b-it",
-    token=HF_TOKEN
-)
-
-embed_model = SentenceTransformer("all-MiniLM-L6-v2")
-
-# ==============================
-# FAISS MEMORY
-# ==============================
-dimension = 384
-index = faiss.IndexFlatL2(dimension)
-memory_texts = []
-
-def embed(text):
-    return embed_model.encode(text).astype("float32")
-
-def store_memory(src, tgt):
-    text_pair = f"{src} -> {tgt}"
-    vec = embed(text_pair)
-    index.add(np.array([vec]))
-    memory_texts.append(text_pair)
-
-def retrieve_memory(query):
-    if len(memory_texts) == 0:
-        return None
-    vec = embed(query)
-    D, I = index.search(np.array([vec]), k=1)
-    return memory_texts[I[0][0]]
+tokenizer, model = load_model()
 
 # ==============================
-# SAFE LANGUAGE DETECTION
+# LANGUAGE MAP
 # ==============================
-def safe_detect(text):
-    text = text.strip()
-
-    # Avoid wrong detection for short text
-    if len(text) < 5:
-        return "auto"
-
-    try:
-        return detect(text)
-    except:
-        return "auto"
-
-# ==============================
-# SPEECH TO TEXT
-# ==============================
-def speech_to_text(audio_file):
-    recognizer = sr.Recognizer()
-
-    with tempfile.NamedTemporaryFile(delete=False) as temp_audio:
-        temp_audio.write(audio_file.read())
-        temp_audio_path = temp_audio.name
-
-    with sr.AudioFile(temp_audio_path) as source:
-        audio = recognizer.record(source)
-
-    try:
-        text = recognizer.recognize_google(audio)
-    except:
-        text = ""
-
-    return text
+lang_map = {
+    "English": "eng_Latn",
+    "Tamil": "tam_Taml",
+    "Hindi": "hin_Deva",
+    "French": "fra_Latn",
+    "Arabic": "arb_Arab",
+    "Spanish": "spa_Latn",
+    "German": "deu_Latn"
+}
+
+detect_map = {
+    "en": "eng_Latn",
+    "ta": "tam_Taml",
+    "hi": "hin_Deva",
+    "fr": "fra_Latn",
+    "ar": "arb_Arab",
+    "es": "spa_Latn",
+    "de": "deu_Latn"
+}
 
 # ==============================
 # TRANSLATION FUNCTION
 # ==============================
 def translate(text, target_lang):
 
-    src_lang = safe_detect(text)
-    memory = retrieve_memory(text)
-
-    # 🚨 Handle very short input
-    if len(text.split()) <= 1:
-        return "⚠️ Please enter a full sentence for better translation.", src_lang, memory
-
-    # Prompt design
-    if src_lang == "auto":
-        prompt = f"""
-You are a professional multilingual translator.
-
-Detect the language and translate into {target_lang}.
-
-Text:
-{text}
-
-Rules:
-- Only return translated text
-- No explanation
-"""
-    else:
-        prompt = f"""
-You are a professional multilingual translator.
-
-Translate from {src_lang} to {target_lang}.
-
-Text:
-{text}
-
-Rules:
-- Only return translated text
-- No explanation
-"""
-
     try:
-        response = client.text_generation(
-            prompt,
-            max_new_tokens=150,
-            temperature=0.2,
-            top_p=0.9
-        )
+        detected = detect(text)
+        src_lang = detect_map.get(detected, "eng_Latn")
+    except:
+        src_lang = "eng_Latn"
+
+    tgt_lang = lang_map[target_lang]
 
-        translated = response.strip()
+    tokenizer.src_lang = src_lang
 
-        # Handle empty or bad output
-        if not translated or len(translated) < 2:
-            translated = "❌ Unable to translate. Try a clearer sentence."
+    encoded = tokenizer(text, return_tensors="pt")
 
-    except Exception as e:
-        translated = f"❌ Translation failed: {str(e)}"
+    generated_tokens = model.generate(
+        **encoded,
+        forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang),
+        max_length=200
+    )
 
-    store_memory(text, translated)
+    translated = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
-    return translated, src_lang, memory
+    return translated, src_lang
 
 # ==============================
 # UI
 # ==============================
-st.set_page_config(page_title="AI Translator", layout="wide")
-
-st.title("🌍 AI Translator with Voice (Gemma 7B)")
-
-tab1, tab2 = st.tabs(["📝 Text Input", "🎤 Voice Input"])
-
-input_text = ""
+st.title("🌍 NLLB Translator (Transformers)")
 
-# TEXT INPUT
-with tab1:
-    input_text = st.text_area("Enter text", height=150)
+input_text = st.text_area("Enter text")
 
-# VOICE INPUT
-with tab2:
-    audio_file = st.file_uploader("Upload audio (wav/mp3)", type=["wav", "mp3"])
+target_lang = st.selectbox("Target Language", list(lang_map.keys()))
 
-    if audio_file:
-        st.audio(audio_file)
-
-        if st.button("Convert Speech to Text"):
-            with st.spinner("Processing audio..."):
-                input_text = speech_to_text(audio_file)
-                st.success("Recognized Text:")
-                st.write(input_text)
-
-# TARGET LANGUAGE
-target_lang = st.selectbox(
-    "Target Language",
-    ["English", "Tamil", "Hindi", "French", "Arabic", "Spanish", "German"]
-)
-
-# TRANSLATE
 if st.button("Translate"):
     if not input_text.strip():
-        st.warning("Please provide input text or audio")
+        st.warning("Enter text")
     else:
         with st.spinner("Translating..."):
-            output, src_lang, memory = translate(input_text, target_lang)
+            output, src_lang = translate(input_text, target_lang)
 
             st.success("✅ Translation")
             st.write(output)
 
-            st.info(f"Detected Language: {src_lang}")
-
-            if memory:
-                st.caption(f"💡 Similar past translation: {memory}")
+            st.info(f"Detected Language: {src_lang}")