fix: try to lightweight it

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
-import faiss
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import os
 from datasets import load_from_disk
 import torch
@@ -13,51 +12,37 @@ logging.basicConfig(level=logging.INFO)
 DATA_DIR = "/data" if os.path.exists("/data") else "."
 DATASET_DIR = os.path.join(DATA_DIR, "rag_dataset")
 DATASET_PATH = os.path.join(DATASET_DIR, "dataset")
-INDEX_PATH = os.path.join(DATASET_DIR, "embeddings.faiss")
 
 # Cache models and dataset
-@st.cache_resource  # Cache models in memory
+@st.cache_resource
 def load_models():
-    tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
-    retriever = RagRetriever.from_pretrained(
-        "facebook/rag-sequence-nq",
-        index_name="custom",
-        passages_path=DATASET_PATH,
-        index_path=INDEX_PATH
-    )
-    model = RagSequenceForGeneration.from_pretrained(
-        "facebook/rag-sequence-nq",
-        retriever=retriever
-    )
-    # Move to CPU (since we're in a CPU environment)
-    model = model.cpu()
-    return tokenizer, retriever, model
+    model_name = "t5-base"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    return tokenizer, model
 
-@st.cache_data  # Cache dataset on disk
-def load_dataset():
-    # Create initial dataset if it doesn't exist
-    if not os.path.exists(DATASET_PATH):
-        with st.spinner("Building initial dataset from autism research papers..."):
-            import faiss_index.index as idx
-            papers = idx.fetch_arxiv_papers("autism research", max_results=100)
-            idx.build_faiss_index(papers, dataset_dir=DATASET_DIR)
-    return load_from_disk(DATASET_PATH)
-
-# RAG Pipeline
-def rag_pipeline(query, dataset, index):
-    tokenizer, retriever, model = load_models()
-    inputs = tokenizer(query, return_tensors="pt", max_length=512, truncation=True)
+def generate_answer(question, context, max_length=200):
+    tokenizer, model = load_models()
+    
+    # Encode the question and context
+    inputs = tokenizer(
+        f"question: {question} context: {context}",
+        add_special_tokens=True,
+        return_tensors="pt",
+        max_length=512,
+        truncation=True,
+        padding=True
+    )
+    
+    # Get model predictions
     with torch.no_grad():
-        outputs = model.generate(
-            inputs["input_ids"],
-            max_length=200,
-            min_length=50,
-            num_beams=5,
-            early_stopping=True,
-            no_repeat_ngram_size=3
-        )
-        answer = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-    return answer
+        outputs = model(**inputs)
+        answer_ids = torch.argmax(outputs.logits, dim=-1)
+        
+        # Convert token positions to text
+        answer = tokenizer.decode(answer_ids[0], skip_special_tokens=True)
+        
+    return answer if answer and not answer.isspace() else "I cannot find a specific answer to this question in the provided context."
 
 # Streamlit App
 st.title("🧩 AMA Autism")
@@ -65,10 +50,26 @@ query = st.text_input("Please ask me anything about autism ✨")
 
 if query:
     with st.status("Searching for answers..."):
+        # Load dataset
         dataset = load_dataset()
-        answer = rag_pipeline(query, dataset, index=None)
-        if answer:
+        
+        # Get relevant context
+        context = "\n".join([
+            f"{paper['text'][:1000]}"  # Use more context for better answers
+            for paper in dataset[:3]
+        ])
+        
+        # Generate answer
+        answer = generate_answer(query, context)
+        
+        if answer and not answer.isspace():
             st.success("Answer found!")
             st.write(answer)
+            
+            st.write("### Sources Used:")
+            for i in range(min(3, len(dataset))):
+                st.write(f"**Title:** {dataset[i]['title']}")
+                st.write(f"**Summary:** {dataset[i]['text'][:200]}...")
+                st.write("---")
         else:
-            st.error("Failed to generate an answer.")
+            st.warning("I couldn't find a specific answer in the research papers. Try rephrasing your question.")