| import gradio as gr | |
| import PyPDF2 | |
| import numpy as np | |
| from sentence_transformers import SentenceTransformer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| import re | |
| import io | |
| class SimpleRAG: | |
| def __init__(self): | |
| self.model = SentenceTransformer('all-MiniLM-L6-v2') | |
| self.chunks = [] | |
| self.embeddings = None | |
| self.document_processed = False | |
| def extract_text_from_pdf(self, pdf_file): | |
| try: | |
| pdf_reader = PyPDF2.PdfReader(io.BytesIO(pdf_file)) | |
| text = "" | |
| for page in pdf_reader.pages: | |
| text += page.extract_text() + "\n" | |
| return text | |
| except Exception as e: | |
| raise Exception(f"Error reading PDF: {str(e)}") | |
| def chunk_text(self, text, chunk_size=500, overlap=100): | |
| sentences = re.split(r'[.!?]+', text) | |
| chunks = [] | |
| current_chunk = "" | |
| for sentence in sentences: | |
| sentence = sentence.strip() | |
| if not sentence: | |
| continue | |
| if len(current_chunk) + len(sentence) <= chunk_size: | |
| current_chunk += sentence + ". " | |
| else: | |
| if current_chunk: | |
| chunks.append(current_chunk.strip()) | |
| current_chunk = sentence + ". " | |
| if current_chunk: | |
| chunks.append(current_chunk.strip()) | |
| return [chunk for chunk in chunks if len(chunk) > 50] | |
| def process_document(self, pdf_file): | |
| try: | |
| text = self.extract_text_from_pdf(pdf_file) | |
| if not text.strip(): | |
| return "Error: Could not extract text from PDF. Please try a different file." | |
| self.chunks = self.chunk_text(text) | |
| if not self.chunks: | |
| return "Error: No meaningful content found in the PDF." | |
| self.embeddings = self.model.encode(self.chunks) | |
| self.document_processed = True | |
| return f"Document processed successfully! Created {len(self.chunks)} text chunks. You can now ask questions." | |
| except Exception as e: | |
| return f"Error processing document: {str(e)}" | |
| def search_similar_chunks(self, query, top_k=3): | |
| if not self.document_processed: | |
| return "Please upload and process a document first." | |
| try: | |
| query_embedding = self.model.encode([query]) | |
| similarities = cosine_similarity(query_embedding, self.embeddings)[0] | |
| top_indices = np.argsort(similarities)[::-1][:top_k] | |
| results = [] | |
| for idx in top_indices: | |
| results.append({ | |
| 'chunk': self.chunks[idx], | |
| 'similarity': similarities[idx] | |
| }) | |
| return results | |
| except Exception as e: | |
| return f"Error searching: {str(e)}" | |
| def answer_question(self, question): | |
| if not self.document_processed: | |
| return "Please upload and process a document first." | |
| if not question.strip(): | |
| return "Please enter a question." | |
| results = self.search_similar_chunks(question) | |
| if isinstance(results, str): | |
| return results | |
| if not results: | |
| return "No relevant information found in the document." | |
| answer = "Based on the document, here's what I found:\n\n" | |
| for i, result in enumerate(results): | |
| answer += f"**Relevant section {i+1}** (similarity: {result['similarity']:.3f}):\n" | |
| answer += f"{result['chunk']}\n\n" | |
| return answer | |
| rag_system = SimpleRAG() | |
| def process_pdf(pdf_file): | |
| if pdf_file is None: | |
| return "Please upload a PDF file." | |
| try: | |
| with open(pdf_file.name, 'rb') as f: | |
| pdf_content = f.read() | |
| result = rag_system.process_document(pdf_content) | |
| return result | |
| except Exception as e: | |
| return f"Error: {str(e)}" | |
| def ask_question(question): | |
| return rag_system.answer_question(question) | |
| with gr.Blocks(title="RAG MVP - Document Q&A", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # RAG MVP - Document Q&A System | |
| Upload a PDF document and ask questions about its content. The system will find relevant information and provide answers based on the document. | |
| **How to use:** | |
| 1. Upload a PDF file | |
| 2. Wait for processing to complete | |
| 3. Ask questions about the document | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| pdf_input = gr.File( | |
| label="Upload PDF Document", | |
| file_types=[".pdf"], | |
| type="filepath" | |
| ) | |
| process_btn = gr.Button("Process Document", variant="primary") | |
| process_output = gr.Textbox( | |
| label="Processing Status", | |
| lines=3, | |
| interactive=False | |
| ) | |
| with gr.Column(scale=1): | |
| question_input = gr.Textbox( | |
| label="Ask a Question", | |
| placeholder="What is this document about?", | |
| lines=2 | |
| ) | |
| ask_btn = gr.Button("Get Answer", variant="secondary") | |
| answer_output = gr.Textbox( | |
| label="Answer", | |
| lines=10, | |
| interactive=False | |
| ) | |
| gr.Markdown(""" | |
| ### Sample Questions to Try: | |
| - What is the main topic of this document? | |
| - Can you summarize the key points? | |
| - What are the important details mentioned? | |
| """) | |
| process_btn.click( | |
| fn=process_pdf, | |
| inputs=[pdf_input], | |
| outputs=[process_output] | |
| ) | |
| ask_btn.click( | |
| fn=ask_question, | |
| inputs=[question_input], | |
| outputs=[answer_output] | |
| ) | |
| question_input.submit( | |
| fn=ask_question, | |
| inputs=[question_input], | |
| outputs=[answer_output] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |