Spaces:
Sleeping
Sleeping
File size: 1,677 Bytes
99637f2 f91cc3b 99637f2 f91cc3b 99637f2 f91cc3b 99637f2 f91cc3b 99637f2 f91cc3b 99637f2 f91cc3b 99637f2 f91cc3b 99637f2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
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):
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, "text": result.summary, "id": str(i)} for i, result in enumerate(results)]
return 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 = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
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"))
return dataset_dir
# Example usage
if __name__ == "__main__":
query = "quantum computing"
papers = fetch_arxiv_papers(query)
build_faiss_index(papers) |