Samp21 commited on
Commit
5f8fbdc
1 Parent(s): b91d544

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +50 -0
app.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import pdfplumber
3
+ from transformers import pipeline, RagTokenizer, RagRetriever, RagSequenceForGeneration
4
+
5
+ def preprocess_text(text):
6
+ # Remove extra whitespace and normalize line breaks
7
+ text = text.replace('\n', ' ').replace('\r', '')
8
+ text = ' '.join(text.split())
9
+ return text
10
+
11
+ st.title("Chat with Your PDF")
12
+
13
+ uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
14
+
15
+ if uploaded_file is not None:
16
+ with st.spinner('Reading PDF...'):
17
+ # Extract text from PDF using pdfplumber
18
+ with pdfplumber.open(uploaded_file) as pdf:
19
+ text = ""
20
+ for page in pdf.pages:
21
+ text += page.extract_text()
22
+
23
+ text = preprocess_text(text)
24
+ st.success('PDF successfully read and preprocessed!')
25
+
26
+ # Display some text from the PDF
27
+ st.text_area("Extracted Text", text[:1000], height=300)
28
+
29
+ # Initialize the RAG model
30
+ tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
31
+ retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
32
+ rag_model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
33
+
34
+ # Tokenize the text for RAG
35
+ input_texts = text.split('. ')
36
+ input_ids = tokenizer(input_texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
37
+
38
+ # Build context embeddings for retrieval
39
+ context_input_ids = retriever(input_ids.input_ids, input_ids.input_ids, num_beams=2)
40
+
41
+ question = st.text_input("Ask a question about the PDF:")
42
+ if question:
43
+ with st.spinner('Searching for answer...'):
44
+ # Tokenize the question
45
+ question_ids = tokenizer(question, return_tensors="pt")['input_ids']
46
+
47
+ # Generate answer using RAG
48
+ generated = rag_model.generate(input_ids=context_input_ids.input_ids, context_input_ids=question_ids, num_beams=2)
49
+ rag_answer = tokenizer.decode(generated[0], skip_special_tokens=True)
50
+ st.write(rag_answer)