Update app.py

app.py

@@ -1,25 +1,45 @@
-import pandas as pd
-import numpy as np
+# ----------------------------- #
+#  Imports
+# ----------------------------- #
 import re
-from sentence_transformers import SentenceTransformer
+import os
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
 import gradio as gr
-from ctransformers import AutoModelForCausalLM
+from sentence_transformers import SentenceTransformer
+
+# Mistral Inference
+from mistral_inference import MistralForCausalLM, MistralTokenizer
 
 # ----------------------------- #
-# Load Hosted Mistral 7B Q4_0
+#  Load Local Mistral Model
 # ----------------------------- #
-llm = AutoModelForCausalLM.from_pretrained(
-    "TheBloke/Mistral-7B-v0.1-Q4_0",  # hosted HF model
-    model_type="mistral",             # model type
-    gpu_layers=32                      # adjust based on GPU/VRAM
-)
+model_path = Path.home().joinpath('mistral_models', '7B-Instruct-v0.3')
+
+tokenizer = MistralTokenizer.from_pretrained(model_path)
+llm = MistralForCausalLM.from_pretrained(model_path)
 
+# ----------------------------- #
+#  Load Embedding Model
+# ----------------------------- #
 embedding_model = SentenceTransformer('nlpaueb/legal-bert-base-uncased')
 
 # ----------------------------- #
-# Parse & Prepare Your Documents
+#  Extract ZIP
+# ----------------------------- #
+zip_path = "provinces.zip"
+extract_folder = "provinces_texts"
+
+if not os.path.exists(extract_folder):
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(extract_folder)
+
+# ----------------------------- #
+#  Parse Files
 # ----------------------------- #
-# Example parsing function (from your previous code)
 date_pattern = re.compile(r"(\d{4}[-]\d{2}[-_]\d{2})")
 
 def parse_metadata_and_content(raw_text):
@@ -29,6 +49,7 @@ 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:
@@ -43,12 +64,37 @@ def parse_metadata_and_content(raw_text):
 
     return metadata, content.strip()
 
-# Load your text documents into df as before
-# df = pd.DataFrame(documents)
-# df["Embedding"] = df["text"].apply(lambda x: embedding_model.encode(x))
+documents = []
+
+for root, dirs, files in os.walk(extract_folder):
+    for filename in files:
+        if filename.startswith("._"):
+            continue
+        if filename.endswith(".txt"):
+            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)
+                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
+                    })
+            except Exception:
+                continue
+
+df = pd.DataFrame(documents)
+df["Embedding"] = df["text"].apply(lambda x: embedding_model.encode(x))
 
 # ----------------------------- #
-# Province Detection
+#  Province Detection
 # ----------------------------- #
 def detect_province(query):
     provinces = {
@@ -72,6 +118,7 @@ def detect_province(query):
         "nwt": "Northwest Territories",
         "northwest territories": "Northwest Territories"
     }
+
     q = query.lower()
     for key, prov in provinces.items():
         if key in q:
@@ -79,7 +126,7 @@ def detect_province(query):
     return None
 
 # ----------------------------- #
-# Guardrails
+#  Guardrails
 # ----------------------------- #
 def is_disallowed(query):
     banned = ["kill", "suicide", "harm yourself", "bomb", "weapon"]
@@ -101,7 +148,7 @@ INTRO_TEXT = (
 )
 
 # ----------------------------- #
-# Retrieval
+#  Retrieval
 # ----------------------------- #
 def retrieve_with_pandas(query, province=None, top_k=2):
     query_embedding = embedding_model.encode([query])[0]
@@ -120,7 +167,7 @@ def retrieve_with_pandas(query, province=None, top_k=2):
     return results
 
 # ----------------------------- #
-# RAG Generator
+#  Main RAG Generator using MistralInference
 # ----------------------------- #
 def generate_with_rag(query):
     if is_disallowed(query):
@@ -138,6 +185,7 @@ def generate_with_rag(query):
         return INTRO_TEXT + "I couldn't find relevant information."
 
     context = " ".join(top_docs_df["text"].tolist())
+
     prompt = f"""
 Use the context below to answer the question.
 CONTEXT:
@@ -147,15 +195,18 @@ QUESTION:
 ANSWER:
 """
 
-    response = llm(
-        prompt,
+    # Generate response
+    response = llm.generate(
+        tokenizer.encode(prompt, return_tensors="pt"),
         max_new_tokens=300,
         temperature=0.2
     )
-    return response[0]["generated_text"].split("ANSWER:")[-1].strip()
+
+    answer = tokenizer.decode(response[0], skip_special_tokens=True)
+    return answer.split("ANSWER:")[-1].strip()
 
 # ----------------------------- #
-# Gradio UI
+#  Gradio UI
 # ----------------------------- #
 def ui_fn(query):
     return generate_with_rag(query)
@@ -170,3 +221,4 @@ demo = gr.Interface(
 if __name__ == "__main__":
     demo.launch(share=True)
 
+