fix: dimension error

app.py

@@ -1,8 +1,9 @@
 import streamlit as st
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration, DPRQuestionEncoder, DPRQuestionEncoderTokenizer
 import faiss
 import os
 from datasets import load_from_disk
+import torch
 
 # Title
 st.title("AMA Austim 🧩")
@@ -29,33 +30,43 @@ def load_rag_dataset(dataset_dir="rag_dataset"):
 # RAG Pipeline
 def rag_pipeline(query, dataset, index):
     # Load pre-trained RAG model and configure retriever
-    tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
+    model_name = "facebook/rag-sequence-nq"
+    tokenizer = RagTokenizer.from_pretrained(model_name)
+    
+    # Configure retriever with correct paths and question encoder
     retriever = RagRetriever.from_pretrained(
-        "facebook/rag-sequence-nq",
+        model_name,
         index_name="custom",
         passages_path=os.path.join("rag_dataset", "dataset"),
-        index_path=os.path.join("rag_dataset", "embeddings.faiss")
+        index_path=os.path.join("rag_dataset", "embeddings.faiss"),
+        use_dummy_dataset=False
     )
-    model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)
+    
+    # Initialize the model with the configured retriever
+    model = RagSequenceForGeneration.from_pretrained(model_name, retriever=retriever)
 
     # Generate answer using RAG
     inputs = tokenizer(query, return_tensors="pt")
-    generated_ids = model.generate(inputs["input_ids"])
-    answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    with torch.no_grad():
+        generated_ids = model.generate(inputs["input_ids"], max_length=200)
+        answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
     return answer
 
 # Run the app
 if query:
-    st.write("Loading or creating RAG dataset...")
-    dataset, index = load_rag_dataset()
-    
-    st.write("Running RAG pipeline...")
+    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 = load_rag_dataset()
+        st.write("Found the best sources!")
+        status.update(
+            label="Download complete!", 
+            state="complete", 
+            expanded=False
+        )
     answer = rag_pipeline(query, dataset, index)
-
     st.write("### Answer:")
     st.write(answer)
-
     st.write("### Retrieved Papers:")
     for i in range(min(5, len(dataset))):
         st.write(f"**Title:** {dataset[i]['title']}")

faiss_index/index.py

@@ -1,9 +1,10 @@
 import numpy as np
 import faiss
-from sentence_transformers import SentenceTransformer
 import arxiv
 from datasets import Dataset
 import os
+from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
+import torch
 
 # Fetch arXiv papers
 def fetch_arxiv_papers(query, max_results=10):
@@ -26,15 +27,29 @@ def build_faiss_index(papers, dataset_dir="rag_dataset"):
         "text": [p["text"] for p in papers],
     })
 
+    # 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
-    embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-    embeddings = embedder.encode(dataset["text"], show_progress_bar=True)
+    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)
     
     # Add embeddings to dataset
     dataset = dataset.add_column("embeddings", [emb.tolist() for emb in embeddings])
     
     # Create FAISS index
-    dimension = embeddings.shape[1]
+    dimension = embeddings.shape[1]  # Should be 768 for DPR
     index = faiss.IndexFlatL2(dimension)
     index.add(embeddings.astype(np.float32))