Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import PyPDF2 | |
| import cohere | |
| from pinecone import Pinecone | |
| from sentence_transformers import SentenceTransformer | |
| import io | |
| # Initialize Pinecone and connect to the index | |
| pc = Pinecone(api_key="0f78bc1b-81f7-4a15-9af3-0fbcf0acdb4e") | |
| index = pc.Index("quickstart") | |
| # Load the sentence transformer model | |
| model = SentenceTransformer('all-MiniLM-L6-v2') | |
| # Initialize Cohere with your API key | |
| co = cohere.Client("CxIrucBVA8NNJJOBUnxwRWq488MVydBku1DlqP1u") | |
| def extract_text_from_pdf(pdf_file): | |
| """Extracts text from the uploaded PDF, with error handling.""" | |
| try: | |
| if pdf_file is None: | |
| return "No file uploaded." | |
| # Read the PDF content from the file path | |
| with open(pdf_file.name, 'rb') as f: | |
| pdf_reader = PyPDF2.PdfReader(f) | |
| text = "" | |
| for page_num in range(len(pdf_reader.pages)): | |
| text += pdf_reader.pages[page_num].extract_text() or "" | |
| if not text.strip(): | |
| return "The uploaded PDF is empty or has no readable content." | |
| return text | |
| except PyPDF2.errors.PdfReadError: | |
| return "The uploaded PDF is encrypted or unreadable." | |
| except Exception as e: | |
| return f"Error reading PDF: {str(e)}" | |
| def store_pdf_embeddings(pdf_text): | |
| """Generate and store embeddings for the uploaded PDF content.""" | |
| segments = [pdf_text[i:i + 512] for i in range(0, len(pdf_text), 512)] | |
| embeddings = model.encode(segments) | |
| vectors = [(f"seg-{i}", embed.tolist()) for i, embed in enumerate(embeddings)] | |
| index.upsert(vectors=vectors) | |
| return "PDF uploaded and stored successfully!" | |
| def ask_question(query): | |
| """Handle user questions and generate answers based on the PDF content.""" | |
| query_embedding = model.encode(query).tolist() | |
| # Retrieve the most relevant segment from Pinecone | |
| result = index.query(top_k=1, vector=query_embedding) | |
| retrieved_seg_id = result['matches'][0]['id'] | |
| segment_text = f"Segment: {retrieved_seg_id}" | |
| # Generate the answer using the retrieved segment as context | |
| prompt = f"{segment_text}\nQuestion: {query}\nAnswer:" | |
| response = co.generate( | |
| model="command-xlarge-nightly", | |
| prompt=prompt, | |
| max_tokens=50 | |
| ) | |
| # Return both the segment and the answer | |
| answer = response.generations[0].text.strip() | |
| return segment_text, answer | |
| # Gradio Interface Setup | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Interactive QA Bot with PDF Support") | |
| # PDF Upload Section | |
| pdf_input = gr.File(label="Upload PDF", type="filepath", file_types=[".pdf"]) | |
| upload_status = gr.Textbox(label="Upload Status", interactive=False) | |
| upload_button = gr.Button("Upload and Store") | |
| # Handle PDF Upload | |
| upload_button.click( | |
| lambda pdf: store_pdf_embeddings(extract_text_from_pdf(pdf)) | |
| if pdf is not None else "Please upload a valid PDF.", | |
| inputs=pdf_input, outputs=upload_status | |
| ) | |
| # Question and Answer Section | |
| query_input = gr.Textbox(label="Enter your question") | |
| segment_output = gr.Textbox(label="Retrieved Segment", interactive=False) | |
| answer_output = gr.Textbox(label="Answer", interactive=False) | |
| query_button = gr.Button("Ask") | |
| # Handle User Questions | |
| query_button.click( | |
| ask_question, inputs=query_input, outputs=[segment_output, answer_output] | |
| ) | |
| demo.launch(share=True) # Set share=True if you want a public link |