Update app.py

app.py

@@ -3,6 +3,7 @@ import faiss
 import numpy as np
 from sentence_transformers import SentenceTransformer
 import openai
+from datasets import load_dataset
 
 # Load FAISS index
 index = faiss.read_index("faiss_index.bin")
@@ -10,28 +11,29 @@ index = faiss.read_index("faiss_index.bin")
 # Load embedding model
 model = SentenceTransformer("all-MiniLM-L6-v2")
 
+# Load dataset (only titles for reference)
+dataset = load_dataset("macadeliccc/US-LegalKit", split="train")
+law_texts = [item['text'] for item in dataset if 'text' in item]
+
 # OpenAI API Key (store it as a secret in Hugging Face)
 openai.api_key = st.secrets["GROQ_API_KEY"]
 
-# Load the preprocessed Pile Law dataset (replace with your dataset path)
-law_data = ["Sample Legal Document 1...", "Sample Legal Document 2..."]  # Replace with actual data loading
-
 # Function to search relevant legal documents
 def search_legal_docs(query, top_k=5):
     query_embedding = model.encode([query])
     _, idxs = index.search(query_embedding, top_k)
-    return [law_data[i] for i in idxs[0]]
+    return [law_texts[i] for i in idxs[0]]  # Return matching legal documents
 
 # Streamlit UI
-st.title("🔍 Legal AI Assistant (Pile Law)")
+st.title("🔍 Legal AI Assistant (US-LegalKit)")
 
 query = st.text_input("📌 Enter your legal query:")
 
 if query:
     results = search_legal_docs(query)
     st.write("### 📄 Relevant Legal Documents:")
-    for res in results:
-        st.write(f"- {res}")
+    for i, doc in enumerate(results, 1):
+        st.write(f"**{i}.** {doc[:500]}...")  # Show preview of the document
 
     # Generate AI-based legal response
     response = openai.ChatCompletion.create(