Spaces:
Sleeping
Sleeping
File size: 2,567 Bytes
f1586e3 99637f2 1577a49 13a46cd 99637f2 f1586e3 99637f2 f1586e3 99637f2 f1586e3 99637f2 f1586e3 99637f2 f1586e3 99637f2 f1586e3 99637f2 f1586e3 |
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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
import streamlit as st
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import arxiv
# Title
st.title("arXiv RAG with Streamlit")
# Input: Query
query = st.text_input("Enter your query:")
# 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):
st.warning("FAISS index not found. Building index from scratch...")
# Import the build function from the other file
import 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)
st.success("FAISS index built successfully!")
return faiss.read_index(index_file)
# RAG Pipeline
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", passages=papers, index=index)
model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", 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]
return answer
# 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("Running RAG pipeline...")
answer = rag_pipeline(query, papers, 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("---") |