Update app.py

app.py

@@ -1,70 +1,234 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
+from sentence_transformers import SentenceTransformer
+import pandas as pd
+import numpy as np
+import zipfile
+import os
+import re
+import torch
 
+###############################################################################
+# 1) LOAD MISTRAL IN 4-BIT (MUCH FASTER)
+###############################################################################
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+)
 
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    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
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+model_name = "mistralai/Mistral-7B-Instruct-v0.2"
 
-    messages = [{"role": "system", "content": system_message}]
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=bnb_config,
+    device_map="auto"
+)
 
-    messages.extend(history)
+llm = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=200,
+    temperature=0.4,
+)
 
-    messages.append({"role": "user", "content": message})
+###############################################################################
+# 2) LOAD EMBEDDINGS
+###############################################################################
+embedding_model = SentenceTransformer("nlpaueb/legal-bert-base-uncased")
 
-    response = ""
+###############################################################################
+# 3) EXTRACT ZIP + PARSE PROVINCE FILES
+###############################################################################
+zip_path = "/app/provinces.zip"
+extract_folder = "/app/provinces_texts"
 
-    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
+if os.path.exists(extract_folder):
+    import shutil
+    shutil.rmtree(extract_folder)
 
-        response += token
-        yield response
+with zipfile.ZipFile(zip_path, "r") as zip_ref:
+    zip_ref.extractall(extract_folder)
 
+date_regex = re.compile(r"(\d{4}[-_]\d{2}[-_]\d{2})")
 
+def parse_metadata_and_content(raw):
+    if "CONTENT:" not in raw:
+        raise ValueError("Missing CONTENT: block.")
+    header, content = raw.split("CONTENT:", 1)
+
+    metadata = {}
+    pdfs = []
+    for line in header.split("\n"):
+        if ":" in line and not line.strip().startswith("-"):
+            key, value = line.split(":", 1)
+            metadata[key.strip().upper()] = value.strip()
+        elif line.strip().startswith("-"):
+            pdfs.append(line.strip())
+
+    if pdfs:
+        metadata["PDF_LINKS"] = "\n".join(pdfs)
+
+    return metadata, content.strip()
+
+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)
+        try:
+            with open(filepath, "r", encoding="latin-1") as f:
+                raw = f.read()
+            metadata, content = parse_metadata_and_content(raw)
+            for p in [x.strip() for x in content.split("\n\n") if x.strip()]:
+                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
+                })
+        except Exception as e:
+            print("Skipping:", filepath, str(e))
+
+###############################################################################
+# 4) EMBEDDINGS + DATAFRAME
+###############################################################################
+texts = [d["text"] for d in documents]
+embs = embedding_model.encode(texts).astype("float16")
+
+df = pd.DataFrame(documents)
+df["Embedding"] = list(embs)
+
+###############################################################################
+# 5) RAG RETRIEVAL
+###############################################################################
+def retrieve_with_pandas(query, province=None, top_k=2):
+    q_emb = embedding_model.encode([query])[0]
+
+    subset = df if province is None else df[df["province"] == province].copy()
+
+    subset["Similarity"] = subset["Embedding"].apply(
+        lambda x: np.dot(q_emb, x) /
+                  (np.linalg.norm(q_emb) * np.linalg.norm(x))
+    )
+
+    return subset.sort_values("Similarity", ascending=False).head(top_k)
+
+###############################################################################
+# 6) Province detection
+###############################################################################
+def detect_province(query):
+    provinces = {
+        "yukon": "Yukon",
+        "alberta": "Alberta",
+        "bc": "British Columbia",
+        "british columbia": "British Columbia",
+        "manitoba": "Manitoba",
+        "newfoundland": "Newfoundland and Labrador",
+        "labrador": "Newfoundland and Labrador",
+        "sask": "Saskatchewan",
+        "saskatchewan": "Saskatchewan",
+        "ontario": "Ontario",
+        "pei": "Prince Edward Island",
+        "prince edward island": "Prince Edward Island",
+        "quebec": "Quebec",
+        "new brunswick": "New Brunswick",
+        "nb": "New Brunswick",
+        "nova scotia": "Nova Scotia",
+        "nunavut": "Nunavut",
+        "nwt": "Northwest Territories",
+        "northwest territories": "Northwest Territories",
+    }
+    q = query.lower()
+    for k, p in provinces.items():
+        if k in q:
+            return p
+    return None
+
+###############################################################################
+# 7) Guardrails
+###############################################################################
+def is_disallowed(q):
+    banned = ["kill", "suicide", "harm yourself", "bomb", "weapon"]
+    return any(b in q.lower() for b in banned)
+
+def is_off_topic(q):
+    keys = [
+        "tenant","landlord","rent","evict","lease",
+        "deposit","tenancy","rental","apartment",
+        "unit","heating","notice","repair","pets"
+    ]
+    return not any(k in q.lower() for k in keys)
+
+###############################################################################
+# 8) MAIN RAG PIPELINE
+###############################################################################
+def generate_with_rag(query):
+    if is_disallowed(query):
+        return "Sorry — I can’t help with harmful or dangerous topics."
+    if is_off_topic(query):
+        return "Sorry — I can only answer questions about Canadian tenancy and housing law."
+
+    province = detect_province(query)
+    top_docs = retrieve_with_pandas(query, province)
+
+    context = " ".join(top_docs["text"].tolist())
+
+    prompt = f"""
+Use ONLY the context below to answer.
+If the context does not contain the answer, say so.
+Answer in a simple, conversational way.
+
+Context:
+{context}
+
+Question: {query}
+Answer:
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        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)",
-        ),
-    ],
+
+    out = llm(prompt)[0]["generated_text"]
+    answer = out.split("Answer:", 1)[-1].strip()
+
+    # metadata section
+    meta = ""
+    for _, r in top_docs.iterrows():
+        meta += (
+            f"- **Province:** {r['province']}\n"
+            f"  Source: {r['source_title']} (Updated {r['last_updated']})\n"
+            f"  URL: {r['url']}\n"
+        )
+
+    return f"{answer}\n\n**Sources Used:**\n{meta}"
+
+###############################################################################
+# 9) GRADIO CHAT — INTRO ONLY ONCE
+###############################################################################
+INTRO = (
+    "👋 **Welcome!** I'm a Canadian rental housing assistant.\n\n"
+    "I can help you find and explain information from tenancy laws across all provinces.\n"
+    "I am **not a lawyer** — this is not legal advice.\n\n"
+    "What would you like to know?"
 )
 
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
+def start_chat():
+    return [(None, INTRO)]
+
+def respond(message, history):
+    answer = generate_with_rag(message)
+    history.append((message, answer))
+    return history, history
 
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot(value=start_chat())
+    msg = gr.Textbox(label="Ask your question")
+    msg.submit(respond, [msg, chatbot], [chatbot, chatbot])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)