Upload app.py with huggingface_hub

app.py

@@ -8,13 +8,17 @@ from pathlib import Path
 from sentence_transformers import SentenceTransformer
 from huggingface_hub import CommitScheduler
 from openai import OpenAI
+from chromadb.errors import NotFoundError
 
 # Load embedding model
 embed_model = SentenceTransformer("BAAI/bge-small-en-v1.5")
 
-# Load ChromaDB client and collection
+# Load ChromaDB client and try to get the collection
 chroma_client = chromadb.PersistentClient(path="./clause_index")
-collection = chroma_client.get_collection("legal_clauses")
+try:
+    collection = chroma_client.get_collection("legal_clauses")
+except NotFoundError:
+    collection = None
 
 # Setup OpenAI client
 client = OpenAI(
@@ -23,8 +27,8 @@ client = OpenAI(
 )
 
 # Prompt template
-system_message = """You are a legal AI assistant. You answer legal questions using contract clauses retrieved from the database. 
-If no specific context is found, use your legal knowledge to infer the most likely answer based on standard contract practices."""
+system_message = """You are a legal AI assistant trained on contract clause examples from the CUAD dataset.
+If no clauses are retrieved from the database, infer the answer using your understanding of common contractual standards."""
 user_template = """
 ### Context:
 {context}
@@ -48,25 +52,27 @@ scheduler = CommitScheduler(
 def predict(question):
     try:
         query_embedding = embed_model.encode([question], normalize_embeddings=True)[0]
-
-        # Retrieve documents from vector DB
-        results = collection.query(
-            query_embeddings=[query_embedding.tolist()],
-            n_results=3
-        )
-
-        documents = results["documents"][0]
-        metadatas = results["metadatas"][0]
-
-        # If nothing is found, simulate default context
-        if not documents:
-            context = "No directly relevant clauses were retrieved from the contract database."
-        else:
-            context_parts = [
-                f"[Clause Type: {m['clause_type']}] {doc}"
-                for doc, m in zip(documents, metadatas)
-            ]
-            context = "\n\n".join(context_parts)
+        context = "No relevant clauses were found in the database. Please answer using your legal understanding from the CUAD dataset."
+
+        if collection:
+            try:
+                results = collection.query(
+                    query_embeddings=[query_embedding.tolist()],
+                    n_results=3
+                )
+                documents = results["documents"][0]
+                metadatas = results["metadatas"][0]
+
+                if documents:
+                    context_parts = [
+                        f"[Clause Type: {m['clause_type']}] {doc}"
+                        for doc, m in zip(documents, metadatas)
+                    ]
+                    context = "\n\n".join(context_parts)
+
+            except Exception as e:
+                # Log internal error, but let LLM proceed with generic context
+                context = "Due to an internal retrieval issue, please answer based on your legal knowledge from CUAD dataset."
 
         prompt = [
             {"role": "system", "content": system_message},
@@ -106,7 +112,7 @@ demo = gr.Interface(
     inputs=gr.Textbox(label="Enter your legal question:", lines=4),
     outputs=gr.Textbox(label="Answer"),
     title="⚖️ GL_LegalMind",
-    description="Ask legal contract-related questions. Answers are grounded in ChromaDB or inferred using legal AI knowledge."
+    description="Ask legal contract-related questions. Answers are based on ChromaDB if available or inferred using CUAD-based legal knowledge."
 )
 
 demo.queue()