Update app.py

app.py

@@ -1,18 +1,15 @@
 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 pandas as pd
+import numpy as np
+import zipfile
+import os
 import torch
 
-# ----------------------------- #
-# Load Mistral model
-# ----------------------------- #
+# -----------------------------
+# Load Mistral pipeline
+# -----------------------------
 llm = pipeline(
     "text-generation",
     model="mistralai/Mistral-7B-Instruct-v0.2",
@@ -20,58 +17,77 @@ llm = pipeline(
     device_map="auto"
 )
 
+# -----------------------------
+# Load SentenceTransformer embeddings
+# -----------------------------
 embedding_model = SentenceTransformer("nlpaueb/legal-bert-base-uncased")
 
-# ----------------------------- #
-# Extract and load Yukon dataset
-# ----------------------------- #
-extract_folder = "yukon_texts"
-zip_path = "yukon.zip"
-
-if not os.path.exists(extract_folder):
-    with zipfile.ZipFile(zip_path, "r") as zip_ref:
-        zip_ref.extractall(extract_folder)
-
-# ----------------------------- #
-# Parse files and create embeddings
-# ----------------------------- #
-def parse_metadata_and_content(raw_text):
-    """
-    Replace this with your actual parsing function from Colab.
-    Should return metadata dict and content string.
-    """
+# -----------------------------
+# Extract Yukon ZIP
+# -----------------------------
+zip_path = "/app/yukon.zip"  # make sure you uploaded here
+extract_folder = "/app/yukon_texts"
+
+# Remove old folder if exists
+if os.path.exists(extract_folder):
+    import shutil
+    shutil.rmtree(extract_folder)
+
+with zipfile.ZipFile(zip_path, "r") as zip_ref:
+    zip_ref.extractall(extract_folder)
+
+# -----------------------------
+# Parse TXT files and create dataframe
+# -----------------------------
+def parse_metadata_and_content(raw):
     metadata = {}
-    content = raw_text
+    content = raw
+
+    for line in raw.split("\n"):
+        if ":" in line:
+            key, value = line.split(":", 1)
+            metadata[key.strip().upper()] = value.strip()
+
+    content_lines = [
+        line for line in raw.split("\n") if not any(k in line.upper() for k in metadata.keys())
+    ]
+    content = "\n".join(content_lines)
     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"):
+        if filename.startswith("._"):
             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
-            })
+        if filename.endswith(".txt"):
+            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
-# ----------------------------- #
+print("Loaded documents:", len(df))
+
+# -----------------------------
+# Retrieval function
+# -----------------------------
 def retrieve_with_pandas(query, top_k=2):
     query_emb = embedding_model.encode([query])[0]
     df["Similarity"] = df["Embedding"].apply(
@@ -79,6 +95,9 @@ def retrieve_with_pandas(query, top_k=2):
     )
     return df.sort_values("Similarity", ascending=False).head(top_k)
 
+# -----------------------------
+# RAG generation
+# -----------------------------
 def generate_with_rag(query, top_k=2):
     top_docs = retrieve_with_pandas(query, top_k)
     context = " ".join(top_docs["text"].tolist())
@@ -92,7 +111,8 @@ Context:
 
 Question: {query}
 """
-    response = llm(input_text, max_new_tokens=200, num_return_sequences=1)[0]["generated_text"]
+
+    response = llm(input_text, max_new_tokens=150, num_return_sequences=1)[0]['generated_text']
 
     meta = []
     for _, row in top_docs.iterrows():
@@ -103,31 +123,22 @@ Question: {query}
             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",
-    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)"),
-    ],
-)
+    final = f"{response.strip()}\n\nSources Used:\n{metadata_block}"
+    return final
+
+# -----------------------------
+# Gradio Chat
+# -----------------------------
+def respond(message, history):
+    answer = generate_with_rag(message)
+    history.append((message, answer))
+    return history, history
 
 with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="Your question")
+    msg.submit(respond, [msg, chatbot], [chatbot, chatbot])
+    gr.Markdown("Ask questions about Yukon rental rules and landlord responsibilities.")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)