Update app.py

app.py

@@ -7,57 +7,68 @@ from transformers import AutoTokenizer, AutoModel
 from weaviate.classes.init import Auth
 import cohere
 
-# Load credentials from environment variables or hardcode them temporarily
+# --- Configuration ---
 WEAVIATE_URL = "vgwhgmrlqrqqgnlb1avjaa.c0.us-west3.gcp.weaviate.cloud"
 WEAVIATE_API_KEY = "7VoeYTjkOS4aHINuhllGpH4JPgE2QquFmSMn"
 COHERE_API_KEY = "LEvCVeZkqZMW1aLYjxDqlstCzWi4Cvlt9PiysqT8"
 
-# Connect to Weaviate
+# --- Initialize Clients ---
 client = weaviate.connect_to_weaviate_cloud(
     cluster_url=WEAVIATE_URL,
     auth_credentials=Auth.api_key(WEAVIATE_API_KEY),
     headers={"X-Cohere-Api-Key": COHERE_API_KEY}
 )
-
 cohere_client = cohere.Client(COHERE_API_KEY)
 
-# Load sentence-transformer model
+# --- Load Sentence Transformer ---
 tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
 model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
 
+# --- Utility Functions ---
 def load_pdf(file):
-    """Extract text from PDF file."""
+    """Extract text from a 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."""
+    """Compute mean-pooled embeddings using a transformer."""
     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")
+    """Insert document chunks into Weaviate."""
+    try:
+        doc_collection = client.collections.get("Document")
+    except Exception as e:
+        raise RuntimeError("❌ Collection 'Document' not found. Make sure it's defined in your Weaviate schema.") from e
+
     for chunk in chunks:
-        embedding = get_embeddings(chunk)
-        doc_collection.data.insert(
-            properties={"content": chunk},
-            vector=embedding.tolist()
-        )
+        try:
+            embedding = get_embeddings(chunk)
+            doc_collection.data.insert(
+                properties={"content": chunk},
+                vector=embedding.tolist()
+            )
+        except Exception as e:
+            print(f"⚠️ Skipped chunk due to error: {e}")
 
 def query_answer(query):
-    """Search for top relevant document chunks based on query embedding."""
+    """Query Weaviate for top relevant document chunks."""
     query_embedding = get_embeddings(query)
-    results = client.collections.get("Document").query.near_vector(
-        near_vector=query_embedding.tolist(),
-        limit=3
-    )
-    return results.objects
+    try:
+        results = client.collections.get("Document").query.near_vector(
+            near_vector=query_embedding.tolist(),
+            limit=3
+        )
+        return results.objects
+    except Exception as e:
+        print(f"⚠️ Query error: {e}")
+        return []
 
 def generate_response(context, query):
-    """Generate answer using Cohere model based on context and query."""
+    """Generate a natural language response using Cohere."""
     response = cohere_client.generate(
         model='command',
         prompt=f"Context: {context}\n\nQuestion: {query}\nAnswer:",
@@ -66,25 +77,23 @@ def generate_response(context, query):
     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."""
+    """Main QA pipeline."""
     try:
         document_text = load_pdf(pdf_file)
         document_chunks = [document_text[i:i+500] for i in range(0, len(document_text), 500)]
-
+        
         upload_document_chunks(document_chunks)
         top_docs = query_answer(query)
-
-        context = ' '.join([doc.properties['content'] for doc in top_docs])
+        context = ' '.join([doc.properties['content'] for doc in top_docs if 'content' in doc.properties])
         answer = generate_response(context, query)
 
         return str(context), str(answer)
     finally:
-        client.close()  # ✅ Properly close client to avoid memory leaks
+        client.close()
 
-# Gradio UI
+# --- Gradio UI ---
 with gr.Blocks(theme="compact") as demo:
-    gr.Markdown(
-        """
+    gr.Markdown("""
         <div style="text-align: center; font-size: 28px; font-weight: bold; margin-bottom: 20px; color: #2D3748;">
             📄 Interactive QA Bot 🔍
         </div>
@@ -92,8 +101,7 @@ with gr.Blocks(theme="compact") as demo:
             Upload a PDF document, ask questions, and receive answers based on the document content.
         </p>
         <hr style="border: 1px solid #CBD5E0; margin: 20px 0;">
-        """
-    )
+    """)
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -111,8 +119,7 @@ with gr.Blocks(theme="compact") as demo:
         outputs=[doc_segments_output, answer_output]
     )
 
-    gr.Markdown(
-        """
+    gr.Markdown("""
         <style>
             body {
                 background-color: #EDF2F7;
@@ -141,7 +148,6 @@ with gr.Blocks(theme="compact") as demo:
                 background-color: #FAFAFA;
             }
         </style>
-        """
-    )
+    """)
 
-demo.launch(share=True)  # ✅ Required for Hugging Face Spaces
+demo.launch(share=True)