fix: faiss error

app.py

@@ -23,32 +23,23 @@ def fetch_arxiv_papers(query, max_results=5):
     papers = [{"title": result.title, "summary": result.summary, "pdf_url": result.pdf_url} for result in results]
     return papers
 
+# Load FAISS index
+def load_faiss_index(index_file="faiss_index.index"):
+    return faiss.read_index(index_file)
+
 # RAG Pipeline
-def rag_pipeline(query, papers):
+def rag_pipeline(query, papers, index):
     # Load pre-trained RAG model
     tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
-    retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="custom")
+    retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="custom", passages=papers, index=index)
     model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)
 
-    # Encode papers into embeddings
-    embedder = SentenceTransformer('all-MiniLM-L6-v2')
-    paper_embeddings = embedder.encode([paper["summary"] for paper in papers])
-
-    # Build FAISS index
-    index = faiss.IndexFlatL2(paper_embeddings.shape[1])
-    index.add(paper_embeddings)
-
-    # Retrieve relevant papers
-    query_embedding = embedder.encode([query])
-    distances, indices = index.search(query_embedding, k=2)  # Top 2 relevant papers
-    relevant_papers = [papers[i] for i in indices[0]]
-
     # 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]
 
-    return answer, relevant_papers
+    return answer
 
 # Run the app
 if query:
@@ -56,14 +47,17 @@ if query:
     papers = fetch_arxiv_papers(query)
     st.write(f"Found {len(papers)} papers.")
 
+    st.write("Loading FAISS index...")
+    index = load_faiss_index()
+
     st.write("Running RAG pipeline...")
-    answer, relevant_papers = rag_pipeline(query, papers)
+    answer = rag_pipeline(query, papers, index)
 
     st.write("### Answer:")
     st.write(answer)
 
     st.write("### Relevant Papers:")
-    for paper in relevant_papers:
+    for paper in papers:
         st.write(f"**Title:** {paper['title']}")
         st.write(f"**Summary:** {paper['summary']}")
         st.write(f"**PDF URL:** {paper['pdf_url']}")

faiss.index.py

@@ -0,0 +1,36 @@
+import numpy as np
+import faiss
+from sentence_transformers import SentenceTransformer
+import arxiv
+
+# 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, "summary": result.summary, "pdf_url": result.pdf_url} for result in results]
+    return papers
+
+# Build and save FAISS index
+def build_faiss_index(papers, index_file="faiss_index.index"):
+    embedder = SentenceTransformer('all-MiniLM-L6-v2')
+    paper_embeddings = embedder.encode([paper["summary"] for paper in papers])
+
+    # Create FAISS index
+    dimension = paper_embeddings.shape[1]
+    index = faiss.IndexFlatL2(dimension)
+    index.add(paper_embeddings)
+
+    # Save index to disk
+    faiss.write_index(index, index_file)
+    print(f"FAISS index saved to {index_file}")
+
+# Example usage
+if __name__ == "__main__":
+    query = "quantum computing"
+    papers = fetch_arxiv_papers(query)
+    build_faiss_index(papers)