fix: rag

app.py

@@ -2,15 +2,13 @@ import streamlit as st
 from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
 import faiss
 import os
-from datasets import load_from_disk
+from datasets import load_from_disk, Dataset
 import torch
 import logging
-import warnings
-from pathlib import Path
+import traceback
 
 # Configure logging
-logging.basicConfig(level=logging.WARNING)
-warnings.filterwarnings('ignore')
+logging.basicConfig(level=logging.INFO)
 
 # Title
 st.title("🧩 AMA Austim")
@@ -18,92 +16,105 @@ st.title("🧩 AMA Austim")
 # Input: Query
 query = st.text_input("Please ask me anything about autism ✨")
 
-@st.cache_resource
-def load_rag_components(_dataset_path=None, _index_path=None):
-    """Load and cache RAG components to avoid reloading."""
-    model_name = "facebook/rag-sequence-nq"
-    tokenizer = RagTokenizer.from_pretrained(model_name)
-    
-    retriever_config = {
-        "index_name": "custom",
-        "use_dummy_dataset": True
-    }
-    
-    if _dataset_path and _index_path:
-        retriever_config.update({
-            "passages_path": _dataset_path,
-            "index_path": _index_path
-        })
-    
-    retriever = RagRetriever.from_pretrained(model_name, **retriever_config)
-    model = RagSequenceForGeneration.from_pretrained(model_name)
-    return tokenizer, retriever, model
+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"):
-    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)
-    
-    dataset_path = os.path.join(dataset_dir, "dataset")
-    index_path = os.path.join(dataset_dir, "embeddings.faiss")
-    
-    # Load the dataset and index
-    dataset = load_from_disk(dataset_path)
-    index = faiss.read_index(index_path)
-    
-    return dataset, index, dataset_path, index_path
+    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 = load_from_disk(os.path.join(dataset_dir, "dataset"))
+        index = faiss.read_index(os.path.join(dataset_dir, "embeddings.faiss"))
+        return dataset, index
+    except Exception as e:
+        st.error(f"Error loading dataset: {str(e)}\n{traceback.format_exc()}")
+        return None, None
 
 # RAG Pipeline
-def rag_pipeline(query, dataset, index, dataset_path, index_path):
+def rag_pipeline(query, dataset, index):
     try:
-        # Load cached components with paths
-        tokenizer, retriever, model = load_rag_components(dataset_path, index_path)
+        # Initialize components
+        model_name = "facebook/rag-sequence-nq"
+        tokenizer = RagTokenizer.from_pretrained(model_name)
         
-        # Configure retriever with our dataset
-        retriever.index.dataset = dataset
-        retriever.index.index = index
-        model.retriever = retriever
+        # Convert dataset to passages format
+        passages = prepare_rag_passages(dataset)
+        
+        # Initialize retriever with passages
+        retriever = RagRetriever.from_pretrained(
+            model_name,
+            index_name="custom",
+            passages=passages,
+            index=index
+        )
+        
+        # Initialize model with retriever
+        model = RagSequenceForGeneration.from_pretrained(
+            model_name,
+            retriever=retriever,
+            use_auth_token=False
+        )
 
         # Generate answer
         inputs = tokenizer(query, return_tensors="pt", max_length=512, truncation=True)
         with torch.no_grad():
-            generated_ids = model.generate(
+            outputs = model.generate(
                 inputs["input_ids"],
                 max_length=200,
                 min_length=50,
                 num_beams=5,
-                early_stopping=True
+                early_stopping=True,
+                no_repeat_ngram_size=3
             )
-            answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+            answer = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
         
         return answer
-    
     except Exception as e:
-        st.error(f"An error occurred while processing your query: {str(e)}")
+        st.error(f"Error generating answer: {str(e)}\n{traceback.format_exc()}")
         return None
 
 # Run the app
 if query:
     with st.status("Looking for data in the best sources...", expanded=True) as status:
-        st.write("Still looking... this may take a while as we look at some prestigious papers...")
-        dataset, index, dataset_path, index_path = load_rag_dataset()
-        st.write("Found the best sources!")
-        answer = rag_pipeline(query, dataset, index, dataset_path, index_path)
-        st.write("Now answering your question...")
-        status.update(
-            label="Searching complete!", 
-            state="complete", 
-            expanded=False
-        )
-    
-    if answer:
-        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("---")
+        try:
+            st.write("Still looking... this may take a while as we look at some prestigious papers...")
+            dataset, index = load_rag_dataset()
+            
+            if dataset is None or index 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, index)
+                
+                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")

faiss_index/index.py

@@ -5,63 +5,91 @@ from datasets import Dataset
 import os
 from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
 import torch
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
 
-# Fetch arXiv papers
 def fetch_arxiv_papers(query, max_results=10):
-    client = arxiv.Client()
-    search = arxiv.Search(
-        query=query,
-        max_results=max_results,
-        sort_by=arxiv.SortCriterion.SubmittedDate
-    )
-    results = list(client.results(search))
-    papers = [{"title": result.title, "text": result.summary, "id": str(i)} for i, result in enumerate(results)]
-    return papers
+    """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),  # Unique identifier
+                "text": result.summary,  # Main content for embedding
+                "title": result.title,  # Title for display
+            })
+        
+        logging.info(f"Fetched {len(papers)} papers from arXiv")
+        return papers
+    
+    except Exception as e:
+        logging.error(f"Error fetching papers: {str(e)}")
+        raise
 
-# Build and save dataset with FAISS index
 def build_faiss_index(papers, dataset_dir="rag_dataset"):
-    # Create dataset with required columns for RAG
-    dataset = Dataset.from_dict({
-        "id": [p["id"] for p in papers],
-        "text": [p["text"] for p in papers],  # RAG expects 'text' field
-        "title": [p["title"] for p in papers],
-        "embeddings": None,  # Will be filled later
-    })
-
-    # Initialize DPR context encoder (same as used by RAG)
-    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
-    embeddings = []
-    batch_size = 8
-    
-    for i in range(0, len(dataset), batch_size):
-        batch = dataset[i:i + batch_size]["text"]
-        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)
-    
-    # Create FAISS index
-    dimension = embeddings.shape[1]  # Should be 768 for DPR
-    index = faiss.IndexFlatL2(dimension)
-    index.add(embeddings.astype(np.float32))
-    
-    # Save dataset and index
-    os.makedirs(dataset_dir, exist_ok=True)
-    
-    # Save dataset with embeddings
-    dataset = dataset.add_column("embeddings", [emb.tolist() for emb in embeddings])
-    dataset.save_to_disk(os.path.join(dataset_dir, "dataset"))
-    
-    # Save FAISS index
-    faiss.write_index(index, os.path.join(dataset_dir, "embeddings.faiss"))
+    """Build and save dataset with FAISS index for RAG"""
+    try:
+        # Create dataset with required fields for RAG
+        dataset = Dataset.from_dict({
+            "id": [p["id"] for p in papers],
+            "text": [p["text"] for p in papers],  # Main content field
+            "title": [p["title"] for p in papers],  # Additional metadata
+        })
+        
+        logging.info(f"Created dataset with {len(dataset)} papers")
 
-    return dataset_dir
+        # Initialize DPR encoder (same as used by RAG)
+        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
+        embeddings = []
+        batch_size = 8
+        
+        for i in range(0, len(dataset), batch_size):
+            batch = dataset[i:i + batch_size]["text"]
+            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 FAISS index (L2 distance)
+        dimension = embeddings.shape[1]  # Should be 768 for DPR
+        index = faiss.IndexFlatL2(dimension)
+        index.add(embeddings.astype(np.float32))
+        
+        # Save dataset and index
+        os.makedirs(dataset_dir, exist_ok=True)
+        dataset.save_to_disk(os.path.join(dataset_dir, "dataset"))
+        faiss.write_index(index, os.path.join(dataset_dir, "embeddings.faiss"))
+        
+        logging.info(f"Saved dataset and index to {dataset_dir}")
+        return dataset_dir
+        
+    except Exception as e:
+        logging.error(f"Error building index: {str(e)}")
+        raise
 
 # Example usage
 if __name__ == "__main__":