fix: performance

app.py

@@ -17,33 +17,36 @@ DATASET_PATH = os.path.join(DATASET_DIR, "dataset")
 # Cache models and dataset
 @st.cache_resource
 def load_models():
-    model_name = "t5-base"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    model_name = "google/flan-t5-small"  # Lighter model
+    tokenizer = AutoTokenizer.from_pretrained(model_name, device_map="auto")
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
     return tokenizer, model
 
 @st.cache_data
-def load_dataset():
+def load_dataset(query):
     # 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)
+            papers = idx.fetch_arxiv_papers(f"{query} AND (cat:q-bio.NC OR cat:q-bio.QM OR cat:q-bio.GN OR cat:q-bio.CB OR cat:q-bio.MN)", max_results=50)  # More focused search
             idx.build_faiss_index(papers, dataset_dir=DATASET_DIR)
     
     # Load and convert to pandas for easier handling
     dataset = load_from_disk(DATASET_PATH)
-    return pd.DataFrame({
+    df = pd.DataFrame({
         'title': dataset['title'],
         'text': dataset['text']
     })
+    return df
 
 def generate_answer(question, context, max_length=200):
     tokenizer, model = load_models()
     
-    # Encode the question and context
+    # Add context about medical information
+    prompt = f"Based on scientific research about autism and health: question: {question} context: {context}"
+    
     inputs = tokenizer(
-        f"question: {question} context: {context}",
+        prompt,
         add_special_tokens=True,
         return_tensors="pt",
         max_length=512,
@@ -72,7 +75,7 @@ query = st.text_input("Please ask me anything about autism ✨")
 if query:
     with st.status("Searching for answers..."):
         # Load dataset
-        df = load_dataset()
+        df = load_dataset(query)
         
         # Get relevant context
         context = "\n".join([

faiss_index/index.py

@@ -18,9 +18,9 @@ def fetch_arxiv_papers(query, max_results=10):
     """Fetch papers from arXiv and format them for RAG"""
     client = arxiv.Client()
     search = arxiv.Search(
-        query=query,
+        query=f"{query} AND (cat:q-bio.NC OR cat:q-bio.QM OR cat:q-bio.GN OR cat:q-bio.CB OR cat:q-bio.MN)",  # Focus on biology and medical categories
         max_results=max_results,
-        sort_by=arxiv.SortCriterion.SubmittedDate
+        sort_by=arxiv.SortCriterion.Relevance  # Changed to relevance-based sorting
     )
     results = list(client.results(search))
     papers = [{"id": str(i), "text": result.summary, "title": result.title} for i, result in enumerate(results)]
@@ -29,21 +29,24 @@ def fetch_arxiv_papers(query, max_results=10):
 
 def build_faiss_index(papers, dataset_dir=DATASET_DIR):
     """Build and save dataset with FAISS index for RAG"""
-    # Initialize DPR encoder
-    ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
+    # Initialize smaller DPR encoder
+    ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base", device_map="auto", load_in_8bit=True)
     ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
     
-    # Create embeddings
+    # Create embeddings with smaller batches and memory optimization
     texts = [p["text"] for p in papers]
     embeddings = []
-    batch_size = 8
+    batch_size = 4  # Smaller batch size
+    
     for i in range(0, len(texts), batch_size):
         batch = texts[i:i + batch_size]
-        inputs = ctx_tokenizer(batch, max_length=512, padding=True, truncation=True, return_tensors="pt")
+        inputs = ctx_tokenizer(batch, max_length=256, padding=True, truncation=True, return_tensors="pt")  # Reduced max_length
         with torch.no_grad():
             outputs = ctx_encoder(**inputs)
             batch_embeddings = outputs.pooler_output.cpu().numpy()
             embeddings.append(batch_embeddings)
+            del outputs  # Explicit cleanup
+            torch.cuda.empty_cache()  # Clear GPU memory
     
     embeddings = np.vstack(embeddings)
     logging.info(f"Created embeddings with shape {embeddings.shape}")
@@ -58,7 +61,9 @@ def build_faiss_index(papers, dataset_dir=DATASET_DIR):
     
     # Create FAISS index
     dimension = embeddings.shape[1]
-    index = faiss.IndexFlatL2(dimension)
+    quantizer = faiss.IndexFlatL2(dimension)
+    index = faiss.IndexQuantizer(dimension, quantizer, 8)
+    index.train(embeddings.astype(np.float32))
     index.add(embeddings.astype(np.float32))
     
     # Save dataset and index

requirements.txt

@@ -4,5 +4,5 @@ datasets
 sentence-transformers
 faiss-cpu
 arxiv
-torch
+torch --index-url https://download.pytorch.org/whl/cpu
 accelerate>=0.26.0