Update app.py

app.py

@@ -1,82 +1,118 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 
+import gradio as gr
+import os
+import zipfile
+import pandas as pd
+import numpy as np
 from transformers import pipeline
+from sentence_transformers import SentenceTransformer
 import torch
 
-import zipfile
-import os
-
-
-extract_folder = "yukon_texts"
-
-if not os.path.exists(extract_folder):
-    with zipfile.ZipFile("yukon.zip", 'r') as zip_ref:
-        zip_ref.extractall(extract_folder)
-
-
+# ----------------------------- #
+# Load Mistral model
+# ----------------------------- #
 llm = pipeline(
-    'text-generation',
-    model='mistralai/Mistral-7B-Instruct-v0.2',
+    "text-generation",
+    model="mistralai/Mistral-7B-Instruct-v0.2",
     torch_dtype=torch.float16,
     device_map="auto"
 )
 
-import gradio as gr
+embedding_model = SentenceTransformer("nlpaueb/legal-bert-base-uncased")
+
+# ----------------------------- #
+# Extract and load Yukon dataset
+# ----------------------------- #
+extract_folder = "yukon_texts"
+zip_path = "yukon.zip"
 
-def chat(query):
-    return generate_with_rag(query)
+if not os.path.exists(extract_folder):
+    with zipfile.ZipFile(zip_path, "r") as zip_ref:
+        zip_ref.extractall(extract_folder)
 
-iface = gr.Interface(
-    fn=chat,
-    inputs="text",
-    outputs="text",
-    title="Yukon Residential Tenancy Chatbot"
-)
-iface.launch()
-
-
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
+# ----------------------------- #
+# Parse files and create embeddings
+# ----------------------------- #
+def parse_metadata_and_content(raw_text):
     """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+    Replace this with your actual parsing function from Colab.
+    Should return metadata dict and content string.
     """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
-
-    messages = [{"role": "system", "content": system_message}]
-
-    messages.extend(history)
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+    metadata = {}
+    content = raw_text
+    return metadata, content
+
+documents = []
+for root, dirs, files in os.walk(extract_folder):
+    for filename in files:
+        if filename.startswith("._") or not filename.endswith(".txt"):
+            continue
+        filepath = os.path.join(root, filename)
+        with open(filepath, "r", encoding="latin-1") as f:
+            raw = f.read()
+        metadata, content = parse_metadata_and_content(raw)
+        paragraphs = [p.strip() for p in content.split("\n\n") if p.strip()]
+        for p in paragraphs:
+            documents.append({
+                "source_title": metadata.get("SOURCE_TITLE", "Unknown"),
+                "province": metadata.get("PROVINCE", "Unknown"),
+                "last_updated": metadata.get("LAST_UPDATED", "Unknown"),
+                "url": metadata.get("URL", "N/A"),
+                "pdf_links": metadata.get("PDF_LINKS", ""),
+                "text": p
+            })
+
+texts = [d["text"] for d in documents]
+embeddings = embedding_model.encode(texts).astype("float32")
+df = pd.DataFrame(documents)
+df["Embedding"] = list(embeddings)
+
+# ----------------------------- #
+# RAG Retrieval function
+# ----------------------------- #
+def retrieve_with_pandas(query, top_k=2):
+    query_emb = embedding_model.encode([query])[0]
+    df["Similarity"] = df["Embedding"].apply(
+        lambda x: np.dot(query_emb, x) / (np.linalg.norm(query_emb) * np.linalg.norm(x))
+    )
+    return df.sort_values("Similarity", ascending=False).head(top_k)
+
+def generate_with_rag(query, top_k=2):
+    top_docs = retrieve_with_pandas(query, top_k)
+    context = " ".join(top_docs["text"].tolist())
+
+    input_text = f"""
+Use ONLY the following context to answer the question briefly (2–3 sentences).
+Do NOT guess. Do NOT add external information.
+
+Context:
+{context}
+
+Question: {query}
 """
+    response = llm(input_text, max_new_tokens=200, num_return_sequences=1)[0]["generated_text"]
+
+    meta = []
+    for _, row in top_docs.iterrows():
+        meta.append(
+            f"- Province: {row['province']}\n"
+            f"  Source: {row['source_title']}\n"
+            f"  Updated: {row['last_updated']}\n"
+            f"  URL: {row['url']}\n"
+        )
+    metadata_block = "\n".join(meta)
+    return f"{response.strip()}\n\nSources Used:\n{metadata_block}"
+
+# ----------------------------- #
+# Gradio ChatInterface
+# ----------------------------- #
+def respond(message, history: list[dict[str, str]], system_message, max_tokens, temperature, top_p, hf_token: gr.OAuthToken):
+    # We ignore the system_message, max_tokens, temperature, top_p for simplicity; adjust if needed
+    response = generate_with_rag(message)
+    yield response
+
 chatbot = gr.ChatInterface(
     respond,
     type="messages",
@@ -84,13 +120,7 @@ chatbot = gr.ChatInterface(
         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
 )
 
@@ -99,6 +129,5 @@ with gr.Blocks() as demo:
         gr.LoginButton()
     chatbot.render()
 
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()