Update app.py

app.py

@@ -1,42 +1,61 @@
 import streamlit as st
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
 from datasets import load_dataset
-from transformers import XLMRobertaTokenizer, XLMRobertaForSequenceClassification
-
-# Load a multilingual dataset (use "xnli" or "tydi_qa")
-try:
-    dataset = load_dataset("xnli", "en", split="validation")  # Using English subset as an example
-except Exception as e:
-    st.error(f"Error loading the dataset: {e}")
-
-# Initialize tokenizer and retriever for multilingual support (using XLM-Roberta)
-tokenizer = XLMRobertaTokenizer.from_pretrained("xlm-roberta-base")
-retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="compressed", passages_path="./path_to_multilingual_dataset")
-
-# Initialize the RAG model
-model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
-
-# Define Streamlit app
-st.title('Multilingual RAG Translator/Answer Bot')
-
-st.markdown("This app uses a multilingual RAG model to answer your questions in the language of the query. Ask questions in languages like Urdu, Hindi, or French!")
-
-# User input for query
-user_query = st.text_input("Ask a question in Urdu, Hindi, or French:")
-
-if user_query:
-    # Tokenize the input question
-    inputs = tokenizer(user_query, return_tensors="pt", padding=True, truncation=True)
-    input_ids = inputs['input_ids']
-
-    # Use the retriever to get relevant context
-    retrieved_docs = retriever.retrieve(input_ids)
-
-    # Generate an answer using the context
-    generated_ids = model.generate(input_ids, context_input_ids=retrieved_docs)
-    answer = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-
-    # Display the answer
-    st.write(f"Answer: {answer}")
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
 
-    # Display the most relevant documents
+# Load a multilingual dataset (xnli or tydi_qa)
+def load_data():
+    try:
+        # Load the 'xnli' dataset, validation split
+        dataset = load_dataset("xnli", split="validation")  
+        st.write(f"Loaded {len(dataset)} examples from the 'validation' split.")
+        return dataset
+    except Exception as e:
+        st.write(f"Error loading 'xnli' dataset: {e}")
+        return None
+
+# Initialize RAG model components
+def initialize_rag():
+    try:
+        # Initialize tokenizer and retriever
+        tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
+        retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="compressed", passages_path="./path_to_data")
+        model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
+        return tokenizer, retriever, model
+    except Exception as e:
+        st.write(f"Error initializing RAG components: {e}")
+        return None, None, None
+
+# Main function to run the app
+def main():
+    st.title("Multilingual RAG Translator/Answer Bot")
+    
+    # Load the dataset
+    dataset = load_data()
+    if dataset is None:
+        st.write("Dataset could not be loaded.")
+        return
+
+    # Initialize RAG model components
+    tokenizer, retriever, model = initialize_rag()
+    if tokenizer is None or retriever is None or model is None:
+        st.write("RAG components could not be initialized.")
+        return
+
+    # UI to input a query
+    query = st.text_input("Enter your question in Urdu, Hindi, or French:")
+    
+    if query:
+        # Tokenize the input query
+        inputs = tokenizer(query, return_tensors="pt")
+
+        # Retrieve relevant documents
+        retrieved_docs = retriever.retrieve(query)
+        # Generate an answer using the model
+        generated = model.generate(input_ids=inputs['input_ids'], context_input_ids=retrieved_docs['input_ids'])
+        answer = tokenizer.decode(generated[0], skip_special_tokens=True)
+
+        st.write("Answer:", answer)
+
+# Run the Streamlit app
+if __name__ == "__main__":
+    main()