feat: add dataset

app.py

@@ -1,46 +1,41 @@
 import streamlit as st
 from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
-from sentence_transformers import SentenceTransformer
 import faiss
-import numpy as np
-import arxiv
+import os
+from datasets import load_from_disk
 
 # Title
-st.title("arXiv RAG with Streamlit")
+st.title("AMA Austim 🧩")
 
 # Input: Query
-query = st.text_input("Enter your query:")
+query = st.text_input("Please ask me anything about autism ✨")
 
-# Fetch arXiv papers
-def fetch_arxiv_papers(query, max_results=5):
-    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
-
-# Load FAISS index
-def load_faiss_index(index_file="faiss_index.index"):
-    import os
-    if not os.path.exists(index_file):
+# Load or create RAG dataset
+def load_rag_dataset(dataset_dir="rag_dataset"):
+    if not os.path.exists(dataset_dir):
         # Import the build function from the other file
         import faiss_index.index as faiss_index_index
         
         # Fetch some initial papers to build the index
         initial_papers = faiss_index_index.fetch_arxiv_papers("autism research", max_results=100)
-        faiss_index_index.build_faiss_index(initial_papers, index_file)
+        dataset_dir = faiss_index_index.build_faiss_index(initial_papers, dataset_dir)
     
-    return faiss.read_index(index_file)
+    # 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
 
 # RAG Pipeline
-def rag_pipeline(query, papers, index):
-    # Load pre-trained RAG model
+def rag_pipeline(query, dataset, index):
+    # Load pre-trained RAG model and configure retriever
     tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
-    retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="custom", passages=papers, index=index)
+    retriever = RagRetriever.from_pretrained(
+        "facebook/rag-sequence-nq",
+        index_name="custom",
+        passages_path=os.path.join("rag_dataset", "dataset"),
+        index_path=os.path.join("rag_dataset", "embeddings.faiss")
+    )
     model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)
 
     # Generate answer using RAG
@@ -52,22 +47,17 @@ def rag_pipeline(query, papers, index):
 
 # Run the app
 if query:
-    st.write("Fetching arXiv papers...")
-    papers = fetch_arxiv_papers(query)
-    st.write(f"Found {len(papers)} papers.")
-
-    st.write("Loading FAISS index...")
-    index = load_faiss_index()
-
+    st.write("Loading or creating RAG dataset...")
+    dataset, index = load_rag_dataset()
+    
     st.write("Running RAG pipeline...")
-    answer = rag_pipeline(query, papers, index)
+    answer = rag_pipeline(query, dataset, index)
 
     st.write("### Answer:")
     st.write(answer)
 
-    st.write("### 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']}")
+    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("---")

faiss_index/index.py

@@ -2,6 +2,8 @@ import numpy as np
 import faiss
 from sentence_transformers import SentenceTransformer
 import arxiv
+from datasets import Dataset
+import os
 
 # Fetch arXiv papers
 def fetch_arxiv_papers(query, max_results=10):
@@ -12,22 +14,36 @@ def fetch_arxiv_papers(query, max_results=10):
         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]
+    papers = [{"title": result.title, "text": result.summary, "id": str(i)} for i, result in enumerate(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])
+# Build and save dataset with FAISS index
+def build_faiss_index(papers, dataset_dir="rag_dataset"):
+    # Create dataset
+    dataset = Dataset.from_dict({
+        "id": [p["id"] for p in papers],
+        "title": [p["title"] for p in papers],
+        "text": [p["text"] for p in papers],
+    })
 
+    # Create embeddings
+    embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+    embeddings = embedder.encode(dataset["text"], show_progress_bar=True)
+    
+    # Add embeddings to dataset
+    dataset = dataset.add_column("embeddings", [emb.tolist() for emb in embeddings])
+    
     # Create FAISS index
-    dimension = paper_embeddings.shape[1]
+    dimension = embeddings.shape[1]
     index = faiss.IndexFlatL2(dimension)
-    index.add(paper_embeddings)
+    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"))
 
-    # Save index to disk
-    faiss.write_index(index, index_file)
-    print(f"FAISS index saved to {index_file}")
+    return dataset_dir
 
 # Example usage
 if __name__ == "__main__":

requirements.txt

@@ -3,4 +3,5 @@ transformers
 datasets
 sentence-transformers
 faiss-cpu
-arxiv
+arxiv
+torch