Update app.py

app.py

@@ -1,40 +1,44 @@
+import os
 import gradio as gr
 import PyPDF2
-from transformers import AutoTokenizer, AutoModel
 import torch
 import weaviate
-from weaviate import WeaviateClient
-from weaviate.auth import AuthApiKey
+from transformers import AutoTokenizer, AutoModel
 from weaviate.classes.init import Auth
 import cohere
 
-weaviate_url = "vgwhgmrlqrqqgnlb1avjaa.c0.us-west3.gcp.weaviate.cloud"
-weaviate_api_key = "7VoeYTjkOS4aHINuhllGpH4JPgE2QquFmSMn"
+# Load credentials from environment variables
+WEAVIATE_URL = "vgwhgmrlqrqqgnlb1avjaa.c0.us-west3.gcp.weaviate.cloud"
+WEAVIATE_API_KEY = "7VoeYTjkOS4aHINuhllGpH4JPgE2QquFmSMn"
+COHERE_API_KEY = "LEvCVeZkqZMW1aLYjxDqlstCzWi4Cvlt9PiysqT8"
 
+# Connect to Weaviate
 client = weaviate.connect_to_weaviate_cloud(
-    cluster_url=weaviate_url,
-    auth_credentials=Auth.api_key(weaviate_api_key),
+    cluster_url=WEAVIATE_URL,
+    auth_credentials=Auth.api_key(WEAVIATE_API_KEY),
+    headers={"X-Cohere-Api-Key": COHERE_API_KEY}
 )
 
-cohere_client = cohere.Client("LEvCVeZkqZMW1aLYjxDqlstCzWi4Cvlt9PiysqT8")
-
-def load_pdf(file):
-    reader = PyPDF2.PdfReader(file)
-    text = ''
-    for page in reader.pages:
-        text += page.extract_text()
-    return text
+cohere_client = cohere.Client(COHERE_API_KEY)
 
+# Load sentence-transformer model
 tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
 model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
 
+def load_pdf(file):
+    """Extract text from PDF file."""
+    reader = PyPDF2.PdfReader(file)
+    return ''.join([page.extract_text() for page in reader.pages if page.extract_text()])
+
 def get_embeddings(text):
+    """Generate mean pooled embedding for the input text."""
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
     with torch.no_grad():
         embeddings = model(**inputs).last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
     return embeddings
 
 def upload_document_chunks(chunks):
+    """Insert document chunks into Weaviate collection with embeddings."""
     doc_collection = client.collections.get("Document")
     for chunk in chunks:
         embedding = get_embeddings(chunk)
@@ -44,22 +48,25 @@ def upload_document_chunks(chunks):
         )
 
 def query_answer(query):
+    """Search for top relevant document chunks based on query embedding."""
     query_embedding = get_embeddings(query)
-    response = client.collections.get("Document").query.near_vector(
+    results = client.collections.get("Document").query.near_vector(
         near_vector=query_embedding.tolist(),
         limit=3
     )
-    return response.objects
+    return results.objects
 
 def generate_response(context, query):
+    """Generate answer using Cohere model based on context and query."""
     response = cohere_client.generate(
         model='command',
-        prompt=f"Context: {context}\n\nQuestion: {query}?\nAnswer:",
+        prompt=f"Context: {context}\n\nQuestion: {query}\nAnswer:",
         max_tokens=100
     )
     return response.generations[0].text.strip()
 
 def qa_pipeline(pdf_file, query):
+    """Main pipeline for QA: parse PDF, embed chunks, query Weaviate, and generate answer."""
     document_text = load_pdf(pdf_file)
     document_chunks = [document_text[i:i+500] for i in range(0, len(document_text), 500)]
 
@@ -71,6 +78,7 @@ def qa_pipeline(pdf_file, query):
 
     return context, answer
 
+# Gradio UI
 with gr.Blocks(theme="compact") as demo:
     gr.Markdown(
         """
@@ -133,4 +141,4 @@ with gr.Blocks(theme="compact") as demo:
         """
     )
 
-demo.launch(share = True)
+demo.launch()