Update app.py

app.py

@@ -9,11 +9,10 @@ from googleapiclient.discovery import build
 from openai import OpenAI
 import re
 from datasets import load_dataset, Dataset
-from transformers import AutoTokenizer, AutoModel
+from sentence_transformers import SentenceTransformer
 import torch
 import numpy as np
 import faiss
-import shutil
 import os
 os.environ["HF_HOME"] = "/data/.huggingface"
 
@@ -58,11 +57,8 @@ from huggingface_hub import hf_hub_download
 index_path = hf_hub_download(repo_id="brandonmusic/VerdictAI", filename="knn.index", repo_type="dataset")
 faiss_index = faiss.read_index(index_path)
 
-# Load LegalBERT-DPR encoder
-dpr_model_name = "jhu-clsp/LegalBERT-DPR-CLERC-ft"
-dpr_tokenizer = AutoTokenizer.from_pretrained(dpr_model_name)
-dpr_model = AutoModel.from_pretrained(dpr_model_name)
-dpr_model.eval()
+# Load the correct encoder for embeddings (BAAI/bge-base-en-v1.5)
+encoder_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
 
 # Simulated case law data (replace with Pyserini integration)
 CASE_LAW_DB = {
@@ -134,6 +130,19 @@ STATES = {
     "Other": "Other States"
 }
 
+def retrieve_legal_context(query, k=5):
+    # Encode query to vector
+    query_embedding = encoder_model.encode(query, normalize_embeddings=True)
+    query_embedding = np.array([query_embedding]).astype('float32')  # FAISS expects float32
+    # Search FAISS index
+    distances, indices = faiss_index.search(query_embedding, k)
+    # Fetch texts from dataset (assuming 'text' column; adjust if needed)
+    contexts = []
+    for idx in indices[0]:
+        if idx >= 0 and idx < len(cap_dataset):
+            contexts.append(cap_dataset[int(idx)]['text'])
+    return "\n\n".join(contexts)  # Or format as needed
+
 def route_model(prompt, task_type, files=None, search_web=False, jurisdiction="KY"):
     logger.info(f"Routing prompt: {prompt}, Task: {task_type}, Web Search: {search_web}, Jurisdiction: {jurisdiction}")
     
@@ -615,6 +624,7 @@ with gr.Blocks(css=css, theme=theme, title="VerdictAI - Legal Assistant") as app
             with gr.Accordion("📁 My Files / Cases", open=False, elem_classes=["menu-accordion"]) as files_accordion:
                 file_list = gr.State([])
                 uploaded_files = gr.File(file_count="multiple", file_types=[".pdf", "image", "text"], label="📎 Upload Files")
+                uploaded_files_list = gr.Markdown("Uploaded Files:")
                 def update_file_list(files):
                     if files:
                         file_names = "\n".join(["- " + os.path.basename(f) for f in files])
@@ -624,7 +634,6 @@ with gr.Blocks(css=css, theme=theme, title="VerdictAI - Legal Assistant") as app
                 tag_input = gr.Textbox(placeholder="Tag by case name/client", label="Tag")
                 summarize_btn = gr.Button("AI Summarize", elem_classes=["sidebar-btn"])
                 summarize_btn.click(fn=summarize_document, inputs=[uploaded_files], outputs=summary_output)
-                uploaded_files_list = gr.Markdown("Uploaded Files:")
                 summary_output = gr.Textbox(label="Summary Output", show_label=False)
 
             with gr.Accordion("💬 Saved Chats", open=False, elem_classes=["menu-accordion"]) as chat_accordion: