fix: hf format

app.py

@@ -38,29 +38,32 @@ def load_rag_dataset(dataset_dir="rag_dataset"):
                 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
+        dataset_path = os.path.join(dataset_dir, "dataset")
+        index_path = os.path.join(dataset_dir, "embeddings.faiss")
+        
+        dataset = load_from_disk(dataset_path)
+        
+        # Add FAISS index back to dataset
+        dataset.load_faiss_index('embeddings', index_path)
+        
+        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
+        return None, None, None
 
 # RAG Pipeline
-def rag_pipeline(query, dataset, index):
+def rag_pipeline(query, dataset, dataset_path, index_path):
     try:
         # Initialize components
         model_name = "facebook/rag-sequence-nq"
         tokenizer = RagTokenizer.from_pretrained(model_name)
         
-        # Convert dataset to passages format
-        passages = prepare_rag_passages(dataset)
-        
-        # Initialize retriever with passages
+        # Initialize retriever with correct paths
         retriever = RagRetriever.from_pretrained(
             model_name,
             index_name="custom",
-            passages=passages,
-            index=index
+            passages_path=dataset_path,
+            index_path=index_path
         )
         
         # Initialize model with retriever
@@ -93,15 +96,15 @@ 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, index = load_rag_dataset()
+            dataset, dataset_path, index_path = load_rag_dataset()
             
-            if dataset is None or index is None:
+            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, index)
+                answer = rag_pipeline(query, dataset, dataset_path, index_path)
                 
                 if answer:
                     status.update(label="Search complete!", state="complete", expanded=False)

faiss_index/index.py

@@ -6,6 +6,7 @@ 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)
@@ -24,9 +25,9 @@ def fetch_arxiv_papers(query, max_results=10):
         
         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
+                "id": str(i),
+                "text": result.summary,
+                "title": result.title,
             })
         
         logging.info(f"Fetched {len(papers)} papers from arXiv")
@@ -39,16 +40,16 @@ 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 with required fields for RAG
+        # Create dataset
         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
+            "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 (same as used by RAG)
+        # 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")
         
@@ -74,17 +75,32 @@ 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 (L2 distance)
-        dimension = embeddings.shape[1]  # Should be 768 for DPR
+        # Create FAISS index
+        dimension = embeddings.shape[1]
         index = faiss.IndexFlatL2(dimension)
         index.add(embeddings.astype(np.float32))
         
-        # Save dataset and index
+        # Save everything
         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}")
+        # 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")
+        
+        dataset.save_to_disk(dataset_path)
+        dataset.get_index('embeddings').save(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: