Spaces:
Runtime error
Runtime error
File size: 3,355 Bytes
2210613 6c72794 2210613 6c72794 2210613 c07449b 2210613 |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
# Import necessary libraries
import PyPDF2
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
from sentence_transformers import SentenceTransformer
import torch
from tkinter import Tk, filedialog
# Step 1: Ask for the PDF file via file upload dialog
def upload_file():
"""Open a file dialog to select a PDF file."""
root = Tk()
root.withdraw() # Hide the root window
file_path = filedialog.askopenfilename(
title="Select PDF File",
filetypes=[("PDF Files", "*.pdf")],
)
if not file_path:
raise ValueError("No file was selected. Please upload a PDF.")
return file_path
print("Please select the PDF file containing the chapter.")
try:
file_path = upload_file()
print(f"File selected: {file_path}")
except Exception as e:
print(f"Error: {e}")
exit()
# Step 2: Extract text from the PDF
def extract_text_from_pdf(file_path):
text = ""
with open(file_path, 'rb') as file:
pdf_reader = PyPDF2.PdfReader(file)
for page in pdf_reader.pages:
text += page.extract_text()
return text
chapter_text = extract_text_from_pdf(file_path)
print("Text extracted from the PDF successfully!")
# Step 3: Split the text into smaller passages
def split_text_into_chunks(text, chunk_size=500):
"""Split the text into chunks of size chunk_size."""
words = text.split()
chunks = []
for i in range(0, len(words), chunk_size):
chunk = " ".join(words[i:i + chunk_size])
chunks.append(chunk)
return chunks
passages = split_text_into_chunks(chapter_text)
print(f"Chapter split into {len(passages)} passages for RAG processing.")
# Step 4: Initialize the RAG model and tokenizer
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load the RAG model, tokenizer, and retriever
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
retriever = RagRetriever.from_pretrained(
"facebook/rag-token-nq",
index_name="custom",
passages=passages, # Set passages as the custom index
use_dummy_dataset=True, # Dummy dataset required for custom index
)
model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq").to(device)
# Step 5: Encode passages into embeddings for retrieval
sentence_model = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")
passage_embeddings = sentence_model.encode(passages, convert_to_tensor=True)
retriever.index.set_passages(passages, passage_embeddings.cpu().detach().numpy())
print("Passages indexed successfully!")
# Step 6: Define a function to generate answers
def generate_answer(question, passages):
inputs = tokenizer.prepare_seq2seq_batch(
questions=[question], return_tensors="pt"
).to(device)
generated_ids = model.generate(**inputs)
answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
return answer
# Step 7: Interactive Question-Answering
print("\nChapter is ready. You can now ask questions!")
while True:
user_question = input("\nEnter your question (or type 'exit' to quit): ")
if user_question.lower() == "exit":
print("Exiting the application. Thank you!")
break
try:
answer = generate_answer(user_question, passages)
print(f"Answer: {answer}")
except Exception as e:
print(f"An error occurred: {e}")
|