File size: 2,016 Bytes
2858b70 |
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 |
import PyPDF2
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
# Load models
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
# Function to process PDF
def process_pdf(pdf_file):
pdf_reader = PyPDF2.PdfReader(pdf_file)
document_text = ""
for page in pdf_reader.pages:
document_text += page.extract_text()
sentences = document_text.split('. ')
embeddings = embedding_model.encode(sentences)
faiss_index = faiss.IndexFlatL2(embeddings.shape[1])
faiss_index.add(embeddings)
return sentences, embeddings, faiss_index
# Function to get relevant context
def get_relevant_context(query, faiss_index, sentences, k=3):
query_vector = embedding_model.encode([query])
_, I = faiss_index.search(query_vector, k)
relevant_sentences = [sentences[i] for i in I[0]]
return ". ".join(relevant_sentences)
from transformers import pipeline
qa_model = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
def answer_question(query, faiss_index, sentences):
if not sentences:
return "Please upload a document first.", ""
relevant_context = get_relevant_context(query, faiss_index, sentences)
answer = qa_model(question=query, context=relevant_context)
return answer['answer'], relevant_context
import gradio as gr
def process_and_answer(pdf_file, query):
sentences, embeddings, faiss_index = process_pdf(pdf_file)
answer, context = answer_question(query, faiss_index, sentences)
return answer, context
with gr.Blocks() as demo:
gr.Markdown("# Interactive QA Bot")
pdf_input = gr.File(label="Upload PDF")
query_input = gr.Textbox(label="Ask a question about the document")
answer_output = gr.Textbox(label="Answer")
context_output = gr.Textbox(label="Relevant Context")
submit_button = gr.Button("Submit")
submit_button.click(process_and_answer, inputs=[pdf_input, query_input], outputs=[answer_output, context_output])
demo.launch()
|