Update app.py

app.py

@@ -1,26 +1,27 @@
 # ----------------------------- #
 #  Imports
 # ----------------------------- #
-import re
 import os
+import re
 import zipfile
 from pathlib import Path
 
-import numpy as np
 import pandas as pd
-import gradio as gr
-from sentence_transformers import SentenceTransformer
+import numpy as np
 
-# Mistral Inference
-from mistral_inference import MistralForCausalLM, MistralTokenizer
+from sentence_transformers import SentenceTransformer
+from ctransformers import AutoModelForCausalLM
+import gradio as gr
 
 # ----------------------------- #
-#  Load Local Mistral Model
+#  Load LLM (GGUF quantized Mistral)
 # ----------------------------- #
-model_path = Path.home().joinpath('mistral_models', '7B-Instruct-v0.3')
-
-tokenizer = MistralTokenizer.from_pretrained(model_path)
-llm = MistralForCausalLM.from_pretrained(model_path)
+# Make sure you have downloaded the model locally:
+# e.g., ./models/mistral-7B-v0.1.Q4_0.gguf
+llm = AutoModelForCausalLM.from_pretrained(
+    "./models/mistral-7B-v0.1.Q4_0.gguf",
+    model_type="mistral",
+)
 
 # ----------------------------- #
 #  Load Embedding Model
@@ -28,7 +29,7 @@ llm = MistralForCausalLM.from_pretrained(model_path)
 embedding_model = SentenceTransformer('nlpaueb/legal-bert-base-uncased')
 
 # ----------------------------- #
-#  Extract ZIP
+#  Extract ZIP of provincial texts
 # ----------------------------- #
 zip_path = "provinces.zip"
 extract_folder = "provinces_texts"
@@ -40,8 +41,6 @@ if not os.path.exists(extract_folder):
 # ----------------------------- #
 #  Parse Files
 # ----------------------------- #
-date_pattern = re.compile(r"(\d{4}[-]\d{2}[-_]\d{2})")
-
 def parse_metadata_and_content(raw_text):
     if "CONTENT:" not in raw_text:
         raise ValueError("File missing CONTENT: separator.")
@@ -49,7 +48,6 @@ def parse_metadata_and_content(raw_text):
     header, content = raw_text.split("CONTENT:", 1)
     metadata = {}
     lines = header.strip().split("\n")
-
     pdf_list = []
 
     for line in lines:
@@ -64,6 +62,7 @@ def parse_metadata_and_content(raw_text):
 
     return metadata, content.strip()
 
+
 documents = []
 
 for root, dirs, files in os.walk(extract_folder):
@@ -77,7 +76,6 @@ for root, dirs, files in os.walk(extract_folder):
                     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"),
@@ -90,6 +88,7 @@ for root, dirs, files in os.walk(extract_folder):
             except Exception:
                 continue
 
+# Build DataFrame and compute embeddings
 df = pd.DataFrame(documents)
 df["Embedding"] = df["text"].apply(lambda x: embedding_model.encode(x))
 
@@ -118,7 +117,6 @@ def detect_province(query):
         "nwt": "Northwest Territories",
         "northwest territories": "Northwest Territories"
     }
-
     q = query.lower()
     for key, prov in provinces.items():
         if key in q:
@@ -148,39 +146,29 @@ INTRO_TEXT = (
 )
 
 # ----------------------------- #
-#  Retrieval
+#  Retrieval Function
 # ----------------------------- #
 def retrieve_with_pandas(query, province=None, top_k=2):
     query_embedding = embedding_model.encode([query])[0]
-
-    if province:
-        filtered_df = df[df['province'] == province].copy()
-    else:
-        filtered_df = df.copy()
-
+    filtered_df = df[df['province'] == province].copy() if province else df.copy()
     filtered_df["Similarity"] = filtered_df["Embedding"].apply(
         lambda x: np.dot(query_embedding, x) /
                   (np.linalg.norm(query_embedding) * np.linalg.norm(x))
     )
-
     results = filtered_df.sort_values("Similarity", ascending=False).head(top_k)
     return results
 
 # ----------------------------- #
-#  Main RAG Generator using MistralInference
+#  Main RAG Generator
 # ----------------------------- #
 def generate_with_rag(query):
     if is_disallowed(query):
         return INTRO_TEXT + "Sorry — I can’t help with harmful topics."
-
     if is_off_topic(query):
-        return INTRO_TEXT + (
-            "Sorry — I can only answer questions about tenancy and housing law."
-        )
+        return INTRO_TEXT + "Sorry — I can only answer questions about tenancy and housing law."
 
     province = detect_province(query)
     top_docs_df = retrieve_with_pandas(query, province=province, top_k=2)
-
     if len(top_docs_df) == 0:
         return INTRO_TEXT + "I couldn't find relevant information."
 
@@ -195,15 +183,9 @@ QUESTION:
 ANSWER:
 """
 
-    # Generate response
-    response = llm.generate(
-        tokenizer.encode(prompt, return_tensors="pt"),
-        max_new_tokens=300,
-        temperature=0.2
-    )
-
-    answer = tokenizer.decode(response[0], skip_special_tokens=True)
-    return answer.split("ANSWER:")[-1].strip()
+    # Generate response with ctransformers
+    response = llm(prompt, max_new_tokens=300, temperature=0.2)
+    return response[0]["generated_text"].split("ANSWER:")[-1].strip()
 
 # ----------------------------- #
 #  Gradio UI
@@ -220,5 +202,3 @@ demo = gr.Interface(
 
 if __name__ == "__main__":
     demo.launch(share=True)
-
-