fix: rag

app.py

@@ -2,124 +2,60 @@ import streamlit as st
 from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
 import faiss
 import os
-from datasets import load_from_disk, Dataset
+from datasets import load_from_disk
 import torch
 import logging
-import traceback
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 
-# Title
-st.title("🧩 AMA Austim")
-
-# Input: Query
-query = st.text_input("Please ask me anything about autism ✨")
-
-def prepare_rag_passages(dataset):
-    """Convert dataset to the format expected by RAG"""
-    return [
-        {
-            "id": str(i),
-            "text": row["text"],
-            "title": row["title"],
-            "document_id": int(row["id"])
-        }
-        for i, row in enumerate(dataset)
-    ]
-
-# Load or create RAG dataset
-def load_rag_dataset(dataset_dir="rag_dataset"):
-    try:
-        if not os.path.exists(dataset_dir):
-            with st.spinner("Building initial dataset from autism research papers..."):
-                import faiss_index.index as faiss_index_index
-                initial_papers = faiss_index_index.fetch_arxiv_papers("autism research", max_results=100)
-                dataset_dir = faiss_index_index.build_faiss_index(initial_papers, dataset_dir)
-        
-        # Load the dataset and index
-        dataset_path = os.path.join(dataset_dir, "dataset")
-        index_path = os.path.join(dataset_dir, "embeddings.faiss")
-        
-        if not os.path.exists(dataset_path) or not os.path.exists(index_path):
-            raise ValueError("Dataset or index not found")
-            
-        dataset = load_from_disk(dataset_path)
-        index = faiss.read_index(index_path)
-        
-        logging.info("Successfully loaded dataset and index")
-        return dataset, dataset_path, index_path
-    except Exception as e:
-        st.error(f"Error loading dataset: {str(e)}\n{traceback.format_exc()}")
-        return None, None, None
+# Cache models and dataset
+@st.cache_resource  # Cache models in memory
+def load_models():
+    tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
+    retriever = RagRetriever.from_pretrained(
+        "facebook/rag-sequence-nq",
+        index_name="custom",
+        passages_path="/data/rag_dataset/dataset",
+        index_path="/data/rag_dataset/embeddings.faiss"
+    )
+    model = RagSequenceForGeneration.from_pretrained(
+        "facebook/rag-sequence-nq",
+        retriever=retriever,
+        device_map="auto"
+    )
+    return tokenizer, retriever, model
+
+@st.cache_data  # Cache dataset on disk
+def load_dataset():
+    return load_from_disk("/data/rag_dataset/dataset")
 
 # RAG Pipeline
-def rag_pipeline(query, dataset, dataset_path, index_path):
-    try:
-        # Initialize components
-        model_name = "facebook/rag-sequence-nq"
-        tokenizer = RagTokenizer.from_pretrained(model_name)
-        
-        # Initialize retriever with correct paths
-        retriever = RagRetriever.from_pretrained(
-            model_name,
-            index_name="custom",
-            passages_path=dataset_path,
-            index_path=index_path
-        )
-        
-        # Initialize model with retriever
-        model = RagSequenceForGeneration.from_pretrained(
-            model_name,
-            retriever=retriever,
-            use_auth_token=False
+def rag_pipeline(query, dataset, index):
+    tokenizer, retriever, model = load_models()
+    inputs = tokenizer(query, return_tensors="pt", max_length=512, truncation=True)
+    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
 
-        # Generate answer
-        inputs = tokenizer(query, return_tensors="pt", max_length=512, truncation=True)
-        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
-    except Exception as e:
-        st.error(f"Error generating answer: {str(e)}\n{traceback.format_exc()}")
-        return None
+# Streamlit App
+st.title("🧩 AMA Autism")
+query = st.text_input("Please ask me anything about autism ✨")
 
-# Run the app
 if query:
-    with st.status("Looking for data in the best sources...", expanded=True) as status:
-        try:
-            st.write("Still looking... this may take a while as we look at some prestigious papers...")
-            dataset, dataset_path, index_path = load_rag_dataset()
-            
-            if dataset is None:
-                st.error("Failed to load or create the dataset.")
-                status.update(label="Error loading data", state="error")
-            else:
-                st.write("Found the best sources!")
-                st.write("Now answering your question...")
-                answer = rag_pipeline(query, dataset, dataset_path, index_path)
-                
-                if answer:
-                    status.update(label="Search complete!", state="complete", expanded=False)
-                    st.write("### Answer:")
-                    st.write_stream(answer)
-                    st.write("### Retrieved Papers:")
-                    for i in range(min(5, len(dataset))):
-                        st.write(f"**Title:** {dataset[i]['title']}")
-                        st.write(f"**Summary:** {dataset[i]['text'][:200]}...")
-                        st.write("---")
-                else:
-                    status.update(label="Error generating answer", state="error")
-        
-        except Exception as e:
-            st.error(f"Unexpected error: {str(e)}\n{traceback.format_exc()}")
-            status.update(label="Error", state="error")
+    with st.status("Searching for answers..."):
+        dataset = load_dataset()
+        answer = rag_pipeline(query, dataset, index=None)
+        if answer:
+            st.success("Answer found!")
+            st.write(answer)
+        else:
+            st.error("Failed to generate an answer.")

faiss_index/index.py

@@ -12,95 +12,63 @@ logging.basicConfig(level=logging.INFO)
 
 def fetch_arxiv_papers(query, max_results=10):
     """Fetch papers from arXiv and format them for RAG"""
-    try:
-        client = arxiv.Client()
-        search = arxiv.Search(
-            query=query,
-            max_results=max_results,
-            sort_by=arxiv.SortCriterion.SubmittedDate
-        )
-        results = list(client.results(search))
-        papers = []
-        
-        for i, result in enumerate(results):
-            papers.append({
-                "id": str(i),
-                "text": result.summary,
-                "title": result.title,
-            })
-        
-        logging.info(f"Fetched {len(papers)} papers from arXiv")
-        return papers
-    
-    except Exception as e:
-        logging.error(f"Error fetching papers: {str(e)}")
-        raise
+    client = arxiv.Client()
+    search = arxiv.Search(
+        query=query,
+        max_results=max_results,
+        sort_by=arxiv.SortCriterion.SubmittedDate
+    )
+    results = list(client.results(search))
+    papers = [{"id": str(i), "text": result.summary, "title": result.title} for i, result in enumerate(results)]
+    logging.info(f"Fetched {len(papers)} papers from arXiv")
+    return papers
 
-def build_faiss_index(papers, dataset_dir="rag_dataset"):
+def build_faiss_index(papers, dataset_dir="/data/rag_dataset"):
     """Build and save dataset with FAISS index for RAG"""
-    try:
-        # Initialize DPR encoder
-        ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
-        ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
-        
-        # Create embeddings in batches
-        texts = [p["text"] for p in papers]
-        embeddings = []
-        batch_size = 8
-        
-        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"
-            )
-            
-            with torch.no_grad():
-                outputs = ctx_encoder(**inputs)
-                batch_embeddings = outputs.pooler_output.cpu().numpy()
-                embeddings.append(batch_embeddings)
-        
-        embeddings = np.vstack(embeddings)
-        logging.info(f"Created embeddings with shape {embeddings.shape}")
-        
-        # Create dataset with embeddings
-        dataset = Dataset.from_dict({
-            "id": [p["id"] for p in papers],
-            "text": [p["text"] for p in papers],
-            "title": [p["title"] for p in papers],
-            "embeddings": [emb.tolist() for emb in embeddings],
-        })
-        
-        logging.info(f"Created dataset with {len(dataset)} papers")
-        
-        # Create FAISS index from embeddings
-        dimension = embeddings.shape[1]
-        index = faiss.IndexFlatL2(dimension)
-        index.add(embeddings.astype(np.float32))
-        
-        # Save everything
-        os.makedirs(dataset_dir, exist_ok=True)
-        dataset_path = os.path.join(dataset_dir, "dataset")
-        index_path = os.path.join(dataset_dir, "embeddings.faiss")
-        
-        # Save dataset and index
-        dataset.save_to_disk(dataset_path)
-        faiss.write_index(index, index_path)
-        
-        logging.info(f"Saved dataset to {dataset_path}")
-        logging.info(f"Saved index to {index_path}")
-        
-        return dataset_dir
-        
-    except Exception as e:
-        logging.error(f"Error building index: {str(e)}")
-        raise
+    # Initialize DPR encoder
+    ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
+    ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
+    
+    # Create embeddings
+    texts = [p["text"] for p in papers]
+    embeddings = []
+    batch_size = 8
+    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")
+        with torch.no_grad():
+            outputs = ctx_encoder(**inputs)
+            batch_embeddings = outputs.pooler_output.cpu().numpy()
+            embeddings.append(batch_embeddings)
+    
+    embeddings = np.vstack(embeddings)
+    logging.info(f"Created embeddings with shape {embeddings.shape}")
+    
+    # Create dataset
+    dataset = Dataset.from_dict({
+        "id": [p["id"] for p in papers],
+        "text": [p["text"] for p in papers],
+        "title": [p["title"] for p in papers],
+        "embeddings": [emb.tolist() for emb in embeddings],
+    })
+    
+    # Create FAISS index
+    dimension = embeddings.shape[1]
+    index = faiss.IndexFlatL2(dimension)
+    index.add(embeddings.astype(np.float32))
+    
+    # Save dataset and index
+    os.makedirs(dataset_dir, exist_ok=True)
+    dataset_path = os.path.join(dataset_dir, "dataset")
+    index_path = os.path.join(dataset_dir, "embeddings.faiss")
+    dataset.save_to_disk(dataset_path)
+    faiss.write_index(index, index_path)
+    logging.info(f"Saved dataset to {dataset_path}")
+    logging.info(f"Saved index to {index_path}")
+    return dataset_dir
 
 # Example usage
 if __name__ == "__main__":
-    query = "quantum computing"
-    papers = fetch_arxiv_papers(query)
+    query = "autism research"
+    papers = fetch_arxiv_papers(query, max_results=100)
     build_faiss_index(papers)