Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import pdfplumber | |
| from sentence_transformers import SentenceTransformer | |
| import faiss | |
| import numpy as np | |
| from transformers import pipeline | |
| # Load models | |
| embedding_model = SentenceTransformer('all-mpnet-base-v2') # Better embedding model | |
| qa_pipeline = pipeline("question-answering", model="deepset/roberta-large-squad2") # Larger QA model | |
| # Initialize FAISS index | |
| dimension = 768 # Dimension of the embedding model | |
| index = faiss.IndexFlatL2(dimension) | |
| # Store text chunks and their embeddings | |
| text_chunks = [] | |
| def extract_text_from_pdf(pdf_file): | |
| """Extract text from a PDF file.""" | |
| text = "" | |
| try: | |
| with pdfplumber.open(pdf_file) as pdf: | |
| for page in pdf.pages: | |
| text += page.extract_text() or "" # Handle empty pages | |
| except Exception as e: | |
| return f"Error extracting text: {e}" | |
| return text | |
| def index_text_chunks(pdf_file): | |
| """Split text into chunks, generate embeddings, and index them.""" | |
| global text_chunks, index | |
| # Extract text from the uploaded PDF | |
| text = extract_text_from_pdf(pdf_file) | |
| if not text: | |
| return "No text extracted from the PDF. Please upload a valid PDF file." | |
| # Split text into chunks (e.g., paragraphs) | |
| text_chunks = [chunk for chunk in text.split("\n\n") if chunk.strip()] | |
| # Generate embeddings for the chunks | |
| embeddings = embedding_model.encode(text_chunks) | |
| # Build the FAISS index | |
| index = faiss.IndexFlatL2(dimension) | |
| index.add(np.array(embeddings)) | |
| return f"Paper uploaded and indexed successfully! Found {len(text_chunks)} chunks." | |
| def answer_question(question): | |
| """Retrieve relevant chunks and generate an answer.""" | |
| global text_chunks, index | |
| if not text_chunks: | |
| return "Please upload a paper first." | |
| # Embed the question | |
| question_embedding = embedding_model.encode([question]) | |
| # Retrieve top-k relevant chunks (increase k for more context) | |
| k = 5 # Retrieve more chunks for better context | |
| distances, indices = index.search(question_embedding, k=k) | |
| relevant_chunks = [text_chunks[i] for i in indices[0]] | |
| # Use the QA model to generate an answer | |
| context = " ".join(relevant_chunks) | |
| result = qa_pipeline(question=question, context=context) | |
| # Post-process the answer | |
| answer = result['answer'] | |
| if answer.strip() == "": | |
| return "The paper does not provide enough information to answer this question." | |
| return answer | |
| # Gradio Interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Chat with Your Paper 📄") | |
| gr.Markdown("Upload a PDF of your research paper and ask questions about it.") | |
| with gr.Row(): | |
| pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"]) | |
| upload_status = gr.Textbox(label="Upload Status", interactive=False) | |
| with gr.Row(): | |
| question_input = gr.Textbox(label="Ask a Question", placeholder="What is the main contribution of this paper?") | |
| answer_output = gr.Textbox(label="Answer", interactive=False) | |
| # Buttons | |
| upload_button = gr.Button("Upload and Index Paper") | |
| ask_button = gr.Button("Ask Question") | |
| # Define actions | |
| upload_button.click( | |
| fn=index_text_chunks, | |
| inputs=pdf_input, | |
| outputs=upload_status | |
| ) | |
| ask_button.click( | |
| fn=answer_question, | |
| inputs=question_input, | |
| outputs=answer_output | |
| ) | |
| # Launch the app | |
| demo.launch() |