fix: rag

app.py

@@ -41,11 +41,13 @@ def load_rag_dataset(dataset_dir="rag_dataset"):
         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)
         
-        # Add FAISS index back to dataset
-        dataset.load_faiss_index('embeddings', 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()}")

faiss_index/index.py

@@ -6,7 +6,6 @@ import os
 from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
 import torch
 import logging
-from datasets.utils.file_utils import DownloadConfig
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -40,25 +39,17 @@ def fetch_arxiv_papers(query, max_results=10):
 def build_faiss_index(papers, dataset_dir="rag_dataset"):
     """Build and save dataset with FAISS index for RAG"""
     try:
-        # 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],
-        })
-        
-        logging.info(f"Created dataset with {len(dataset)} papers")
-
         # 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(dataset), batch_size):
-            batch = dataset[i:i + batch_size]["text"]
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i:i + batch_size]
             inputs = ctx_tokenizer(
                 batch,
                 max_length=512,
@@ -75,28 +66,29 @@ def build_faiss_index(papers, dataset_dir="rag_dataset"):
         embeddings = np.vstack(embeddings)
         logging.info(f"Created embeddings with shape {embeddings.shape}")
         
-        # Create FAISS index
+        # 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)
-        
-        # Add embeddings to dataset
-        dataset = dataset.add_faiss_index(
-            column='embeddings',
-            custom_index=index,
-            device=0 if torch.cuda.is_available() else -1
-        )
-        dataset = dataset.add_column("embeddings", [emb.tolist() for emb in embeddings])
-        
-        # Save dataset and index
         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)
-        dataset.get_index('embeddings').save(index_path)
+        faiss.write_index(index, index_path)
         
         logging.info(f"Saved dataset to {dataset_path}")
         logging.info(f"Saved index to {index_path}")