File size: 2,151 Bytes
5f8fbdc |
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 |
import streamlit as st
import pdfplumber
from transformers import pipeline, RagTokenizer, RagRetriever, RagSequenceForGeneration
def preprocess_text(text):
# Remove extra whitespace and normalize line breaks
text = text.replace('\n', ' ').replace('\r', '')
text = ' '.join(text.split())
return text
st.title("Chat with Your PDF")
uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
if uploaded_file is not None:
with st.spinner('Reading PDF...'):
# Extract text from PDF using pdfplumber
with pdfplumber.open(uploaded_file) as pdf:
text = ""
for page in pdf.pages:
text += page.extract_text()
text = preprocess_text(text)
st.success('PDF successfully read and preprocessed!')
# Display some text from the PDF
st.text_area("Extracted Text", text[:1000], height=300)
# Initialize the RAG model
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
rag_model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
# Tokenize the text for RAG
input_texts = text.split('. ')
input_ids = tokenizer(input_texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
# Build context embeddings for retrieval
context_input_ids = retriever(input_ids.input_ids, input_ids.input_ids, num_beams=2)
question = st.text_input("Ask a question about the PDF:")
if question:
with st.spinner('Searching for answer...'):
# Tokenize the question
question_ids = tokenizer(question, return_tensors="pt")['input_ids']
# Generate answer using RAG
generated = rag_model.generate(input_ids=context_input_ids.input_ids, context_input_ids=question_ids, num_beams=2)
rag_answer = tokenizer.decode(generated[0], skip_special_tokens=True)
st.write(rag_answer)
|