Update app.py

app.py

@@ -3,99 +3,125 @@ import pinecone
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForSeq2SeqGeneration
 import torch
-from datasets import load_dataset
+import PyPDF2
+import io
+import os
+from tqdm import tqdm
 
-# Initialize models and databases
-def init_models():
-    # Load the embedding model
-    embeddings_model = SentenceTransformer('all-MiniLM-L6-v2')
-    
-    # Load the LLM for answering
-    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
-    model = AutoModelForSeq2SeqGeneration.from_pretrained("google/flan-t5-base")
+# Initialize models and Pinecone
+embeddings_model = SentenceTransformer('all-MiniLM-L6-v2')
+tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+model = AutoModelForSeq2SeqGeneration.from_pretrained("google/flan-t5-base")
+
+# Initialize Pinecone with environment variable
+PINECONE_API_KEY = os.getenv('PINECONE_API_KEY')
+pinecone.init(api_key=PINECONE_API_KEY, environment="gcp-starter")
+index = pinecone.Index("pdf-index")
+
+def process_pdf(file):
+    # Read PDF content
+    pdf_content = file.read()
+    pdf_file = io.BytesIO(pdf_content)
+    reader = PyPDF2.PdfReader(pdf_file)
     
-    # Initialize Pinecone
-    pinecone.init(api_key="your-pinecone-api-key", environment="gcp-starter")
-    index = pinecone.Index("test-index")
+    # Extract text from PDF
+    text_chunks = []
+    for page in reader.pages:
+        text = page.extract_text()
+        # Split into smaller chunks (roughly 1000 characters each)
+        chunks = [text[i:i+1000] for i in range(0, len(text), 1000)]
+        text_chunks.extend(chunks)
     
-    # Load your dataset from Hugging Face
-    dataset = load_dataset("your-username/your-dataset-name", split="train")
+    # Create embeddings and upload to Pinecone
+    processed_chunks = 0
+    for i, chunk in enumerate(text_chunks):
+        try:
+            # Create embedding
+            embedding = embeddings_model.encode(chunk)
+            
+            # Upload to Pinecone
+            index.upsert(
+                vectors=[(
+                    f"{file.name}_chunk_{i}",
+                    embedding.tolist(),
+                    {
+                        'file_name': file.name,
+                        'chunk_num': i,
+                        'text': chunk
+                    }
+                )]
+            )
+            processed_chunks += 1
+        except Exception as e:
+            print(f"Error processing chunk {i}: {str(e)}")
     
-    return embeddings_model, tokenizer, model, index, dataset
+    return f"Successfully processed {processed_chunks} chunks from {file.name}"
 
-# Generate response using retrieved context
-def generate_answer(question, context, tokenizer, model):
-    prompt = f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
-    
-    outputs = model.generate(
-        **inputs,
-        max_length=512,
-        num_beams=4,
-        temperature=0.7,
-        top_p=0.9,
-        repetition_penalty=1.2,
-        early_stopping=True
-    )
-    
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return answer
+def process_multiple_pdfs(files):
+    results = []
+    for file in files:
+        result = process_pdf(file)
+        results.append(result)
+    return "\n".join(results)
 
-def search_documents(query, embeddings_model, index, dataset, top_k=3):
+def search_documents(query):
     # Create embedding for the query
     query_embedding = embeddings_model.encode(query)
     
     # Search Pinecone
     results = index.query(
         vector=query_embedding.tolist(),
-        top_k=top_k,
+        top_k=3,
         include_metadata=True
     )
     
-    # Get full context from the dataset using metadata
-    contexts = []
-    for match in results.matches:
-        source = match.metadata['source']
-        # Find the corresponding document in the dataset
-        doc = next((item for item in dataset if item['source'] == source), None)
-        if doc:
-            contexts.append(doc['text'])
+    # Generate answer using FLAN-T5
+    context = "\n".join([match.metadata['text'] for match in results.matches])
     
-    return "\n\n".join(contexts)
-
-# Initialize all models and databases
-embeddings_model, tokenizer, model, index, dataset = init_models()
-
-def process_query(query):
-    # Search for relevant documents
-    context = search_documents(query, embeddings_model, index, dataset)
+    prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
     
-    # Generate answer
-    answer = generate_answer(query, context, tokenizer, model)
+    outputs = model.generate(
+        **inputs,
+        max_length=512,
+        num_beams=4,
+        temperature=0.7,
+        top_p=0.9
+    )
+    
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
     # Format sources
-    sources = [f"Source: {match.metadata['source']}" for match in index.query(
-        vector=embeddings_model.encode(query).tolist(),
-        top_k=3,
-        include_metadata=True
-    ).matches]
+    sources = [f"Source: {match.metadata['file_name']}" for match in results.matches]
     
     return answer, "\n".join(sources)
 
-# Create the Gradio interface
+# Create Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Document Search and Q&A")
+    gr.Markdown("# PDF Document Search and Q&A")
+    
+    with gr.Tab("Upload Documents"):
+        file_output = gr.File(
+            file_count="multiple",
+            label="Upload PDF Files"
+        )
+        upload_button = gr.Button("Process PDFs")
+        upload_output = gr.Textbox(label="Processing Results")
     
-    with gr.Row():
+    with gr.Tab("Search and Ask"):
         query_input = gr.Textbox(label="Enter your question")
         search_button = gr.Button("Search")
-    
-    with gr.Row():
         answer_output = gr.Textbox(label="Answer")
         sources_output = gr.Textbox(label="Sources")
     
+    upload_button.click(
+        process_multiple_pdfs,
+        inputs=[file_output],
+        outputs=[upload_output]
+    )
+    
     search_button.click(
-        process_query,
+        search_documents,
         inputs=[query_input],
         outputs=[answer_output, sources_output]
     )