Update app.py

app.py

@@ -1,81 +1,91 @@
 import streamlit as st
-import pandas as pd
-import numpy as np
-import faiss
-import os
-import zipfile
 from langdetect import detect
+import faiss
+import torch
 from sentence_transformers import SentenceTransformer
 from transformers import MBartForConditionalGeneration, MBart50Tokenizer
+import numpy as np
+import pandas as pd
+import os
 
-
-st.set_page_config(page_title="Multilingual RAG Translator/Answer Bot", layout="wide")
-
-st.title("🌍 Multilingual RAG Translator/Answer Bot")
-st.markdown("Ask in Urdu, French, Hindi, etc., and get culturally-aware, context-grounded answers.")
-
-# --- ZIP extraction ---
-zip_file = "all_languages_test.csv.zip"
-csv_file = "all_languages_test.csv"
-
-if not os.path.exists(csv_file):
-    with zipfile.ZipFile(zip_file, "r") as zip_ref:
-        zip_ref.extractall()
-
-# --- Language map and translation model ---
-lang_map = {
-    "en": "en_XX", "fr": "fr_XX", "ur": "ur_PK", "hi": "hi_IN",
-    "es": "es_XX", "de": "de_DE", "zh-cn": "zh_CN", "ar": "ar_AR"
-}
+st.set_page_config(page_title="🌍 Multilingual RAG Translator/Answer Bot", layout="centered")
 
 @st.cache_resource
 def load_resources():
-    df = pd.read_csv(csv_file).dropna()
-    df["context"] = df["premise"] + " " + df["hypothesis"]
-    corpus = df["context"].tolist()
+    embedder = SentenceTransformer("sentence-transformers/distiluse-base-multilingual-cased-v1")
+    tokenizer = MBart50Tokenizer.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+    model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
 
-    embedder = SentenceTransformer("distiluse-base-multilingual-cased-v2")
-    embeddings = embedder.encode(corpus, show_progress_bar=True)
+    # Load multilingual dataset CSV
+    df = pd.read_csv("all_languages_test.csv")
 
-    index = faiss.IndexFlatL2(embeddings.shape[1])
-    index.add(np.array(embeddings))
+    # Construct corpus
+    corpus = (df["premise"] + " " + df["hypothesis"]).fillna("").tolist()
 
-    tokenizer = MBart50Tokenizer.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+    # Compute embeddings for corpus
+    corpus_embeddings = embedder.encode(corpus, convert_to_numpy=True, show_progress_bar=True)
 
-    model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+    # Create FAISS index
+    dimension = corpus_embeddings.shape[1]
+    index = faiss.IndexFlatL2(dimension)
+    index.add(corpus_embeddings)
 
     return embedder, index, corpus, tokenizer, model
 
-embedder, index, corpus, tokenizer, model = load_resources()
+def detect_lang(text):
+    try:
+        return detect(text)
+    except:
+        return "en"
 
-# --- Answer generation ---
-def generate_answer(query, k=3):
-    lang = detect(query)
-    token_lang = lang_map.get(lang, "en_XX")
+def get_top_k_passages(query, embedder, index, corpus, k=3):
+    query_embedding = embedder.encode([query], convert_to_numpy=True)
+    distances, indices = index.search(query_embedding, k)
+    return [corpus[i] for i in indices[0] if i < len(corpus)]
 
-    query_vec = embedder.encode([query])
-    D, I = index.search(np.array(query_vec), k)
-    contexts = [corpus[i] for i in I[0]]
-    context = " ".join(contexts)
+def generate_answer(query, context, tokenizer, model, src_lang):
+    model.eval()
+    tokenizer.src_lang = src_lang
+    joined_context = " ".join(context)
 
-    full_input = f"question: {query} context: {context}"
-    tokenizer.src_lang = token_lang
-    encoded_input = tokenizer(full_input, return_tensors="pt")
+    inputs = tokenizer(query + " " + joined_context, return_tensors="pt", max_length=1024, truncation=True)
     generated_tokens = model.generate(
-        **encoded_input,
-        forced_bos_token_id=tokenizer.lang_code_to_id[token_lang]
+        **inputs,
+        forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"],
+        max_length=256,
+        num_beams=5,
+        early_stopping=True
     )
-
     return tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
 
-# --- UI ---
-user_input = st.text_area("💬 Enter your question in any supported language:")
+st.title("🌍 Multilingual RAG Translator/Answer Bot")
+st.caption("Ask in Urdu, French, Hindi, etc., and get culturally-aware, context-grounded answers.")
+
+query = st.text_input("💬 Enter your question in any supported language:")
 
-if st.button("Get Answer"):
-    if user_input.strip():
-        with st.spinner("Generating answer..."):
-            response = generate_answer(user_input)
-            st.success("Answer:")
-            st.write(response)
+if query:
+    if len(query.strip()) < 3:
+        st.warning("Please enter a more complete question for better results.")
     else:
-        st.warning("Please enter a question first.")
+        with st.spinner("Thinking..."):
+            embedder, index, corpus, tokenizer, model = load_resources()
+            lang = detect_lang(query)
+
+            lang_map = {
+                "en": "en_XX", "fr": "fr_XX", "ur": "ur_PK", "hi": "hi_IN",
+                "es": "es_XX", "de": "de_DE", "zh": "zh_CN", "ar": "ar_AR",
+                "ru": "ru_RU", "tr": "tr_TR", "it": "it_IT", "pt": "pt_XX",
+            }
+
+            src_lang = lang_map.get(lang, "en_XX")
+            context = get_top_k_passages(query, embedder, index, corpus)
+
+            if not context:
+                st.error("⚠️ Could not find any relevant context to answer your question.")
+            else:
+                try:
+                    answer = generate_answer(query, context, tokenizer, model, src_lang)
+                    st.markdown("### 📌 Answer:")
+                    st.success(answer)
+                except Exception as e:
+                    st.error(f"⚠️ Something went wrong while generating the answer.\n\n{e}")