Update app.py

app.py

@@ -10,17 +10,25 @@ from langchain.docstore.document import Document
 from langchain.embeddings import HuggingFaceEmbeddings
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# =============================
+# ===================================
 # PAGE CONFIG
-# =============================
-st.set_page_config(page_title="Kenya Legal Assistant", layout="wide")
+# ===================================
+st.set_page_config(
+    page_title="Kenya Legal Assistant",
+    layout="wide"
+)
 
-# =============================
-# CACHE EMBEDDINGS + VECTOR DB
-# =============================
-@st.cache_resource
+st.title("🇰🇪 Kenya Legal Assistant")
+st.caption("Ask questions about Kenyan court judgments (English or Swahili)")
+
+# ===================================
+# LOAD VECTOR DATABASE (CACHED)
+# ===================================
+@st.cache_resource(show_spinner=True)
 def load_vectorstore():
 
+    st.write("🔎 Loading legal knowledge base...")
+
     dataset = load_dataset(
         "Brian269/Kenyan_Judgements",
         split="train",
@@ -29,12 +37,16 @@ def load_vectorstore():
 
     documents = []
     for i, item in enumerate(dataset):
-        if i > 200:   # prevent HF timeout
+        if i > 200:   # prevents HF startup timeout
             break
+
         documents.append(
             Document(
                 page_content=item["text"],
-                metadata={"source": item["file_name"], "page": 1},
+                metadata={
+                    "source": item["file_name"],
+                    "page": 1
+                },
             )
         )
 
@@ -46,7 +58,9 @@ def load_vectorstore():
     chunks = []
     for doc in documents:
         for chunk in splitter.split_text(doc.page_content):
-            chunks.append(Document(page_content=chunk, metadata=doc.metadata))
+            chunks.append(
+                Document(page_content=chunk, metadata=doc.metadata)
+            )
 
     embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
@@ -54,25 +68,30 @@ def load_vectorstore():
 
     INDEX_PATH = "faiss_index"
 
+    # ✅ Load prebuilt FAISS index if uploaded
     if os.path.exists(INDEX_PATH):
+        st.write("✅ Loading FAISS index...")
         vectorstore = FAISS.load_local(
             INDEX_PATH,
             embeddings,
             allow_dangerous_deserialization=True
         )
     else:
+        st.warning("⚠️ FAISS index not found — building (first run only)...")
         vectorstore = FAISS.from_documents(chunks, embeddings)
         vectorstore.save_local(INDEX_PATH)
 
     return vectorstore
 
 
-# =============================
-# CACHE MODEL (LOAD ONCE)
-# =============================
-@st.cache_resource
+# ===================================
+# LOAD LANGUAGE MODEL (CACHED)
+# ===================================
+@st.cache_resource(show_spinner=True)
 def load_llm():
 
+    st.write("🧠 Loading language model...")
+
     model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
     tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -93,12 +112,13 @@ def load_llm():
     return pipe
 
 
+# Load once
 vectorstore = load_vectorstore()
 pipe = load_llm()
 
-# =============================
+# ===================================
 # HELPERS
-# =============================
+# ===================================
 def detect_language(text):
     try:
         return detect(text)
@@ -106,17 +126,17 @@ def detect_language(text):
         return "en"
 
 
-def translate(text, target):
-    return GoogleTranslator(source="auto", target=target).translate(text)
+def translate(text, target_lang):
+    return GoogleTranslator(source="auto", target=target_lang).translate(text)
 
 
 def build_prompt(question, context):
     return f"""
 You are a Kenyan legal assistant.
 
-Answer ONLY using provided context.
-Include case citations.
-Do not hallucinate.
+Answer ONLY using the provided context.
+Include proper case citations.
+Do not fabricate information.
 
 Context:
 {context}
@@ -130,34 +150,38 @@ Structured Answer:
 
 def ask_kenya_law(question):
 
-    lang = detect_language(question)
-    q_en = translate(question, "en") if lang == "sw" else question
+    language = detect_language(question)
+
+    question_en = (
+        translate(question, "en") if language == "sw" else question
+    )
 
-    docs = vectorstore.similarity_search(q_en, k=4)
-    context = "\n\n".join([d.page_content for d in docs])
+    retrieved_docs = vectorstore.similarity_search(question_en, k=4)
 
-    prompt = build_prompt(q_en, context)
+    context = "\n\n".join([doc.page_content for doc in retrieved_docs])
+
+    prompt = build_prompt(question_en, context)
 
     result = pipe(prompt)[0]["generated_text"]
 
-    if lang == "sw":
+    if language == "sw":
         result = translate(result, "sw")
 
     sources = "\n".join(
-        [f'{d.metadata["source"]} - Page {d.metadata["page"]}' for d in docs]
+        [f'{doc.metadata["source"]} - Page {doc.metadata["page"]}'
+         for doc in retrieved_docs]
     )
 
     return result, sources
 
 
-# =============================
+# ===================================
 # STREAMLIT CHAT UI
-# =============================
-st.title("🇰🇪 Kenya Legal Assistant")
-
+# ===================================
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
+# Display history
 for msg in st.session_state.messages:
     with st.chat_message(msg["role"]):
         st.markdown(msg["content"])
@@ -165,7 +189,10 @@ for msg in st.session_state.messages:
 prompt = st.chat_input("Ask a legal question...")
 
 if prompt:
-    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    st.session_state.messages.append(
+        {"role": "user", "content": prompt}
+    )
 
     with st.chat_message("user"):
         st.markdown(prompt)
@@ -177,11 +204,14 @@ if prompt:
             response = f"""
 {answer}
 
+---
+
 📚 **Sources**
 {sources}
 
 ⚠️ DISCLAIMER:
-Educational legal information only — not legal advice.
+This AI provides legal information for educational purposes only.
+It does NOT constitute legal advice.
 """
 
             st.markdown(response)