Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import re
 import threading
 from typing import List, Dict, Optional
 
@@ -16,7 +17,7 @@ device = None
 def load_model(hf_access_token: Optional[str] = None):
     global tokenizer, model, device
 
-    if tokenizer and model:
+    if tokenizer is not None and model is not None:
         return
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -45,43 +46,63 @@ def load_model(hf_access_token: Optional[str] = None):
     print("✅ Loaded model:", getattr(model.config, "_name_or_path", "unknown"))
 
 
-def build_prompt(system_message: str, history: List[Dict[str, str]], user_message: str) -> str:
+def build_plain_prompt(system_message: str, history: List[Dict[str, str]], user_message: str) -> str:
     """
-    HARD rules to stop:
-    - self questioning
-    - rule narration
-    - exam-style continuation
+    Universal prompt builder that does NOT require tokenizer.chat_template.
+    This works with any CausalLM.
     """
     hard_rules = (
         "You are a professional hotel housekeeping assistant.\n"
         "STRICT RULES:\n"
-        "- Answer ONLY the user's question.\n"
-        "- Do NOT generate follow-up questions.\n"
-        "- Do NOT mention rules, instructions, or reasoning.\n"
-        "- Do NOT narrate your thinking.\n"
-        "- Do NOT continue the conversation on your own.\n"
-        "- Respond in English only.\n"
-        "- Output ONLY the final answer.\n"
+        "1) Respond in English only.\n"
+        "2) Answer ONLY the user's last question.\n"
+        "3) Do NOT generate follow-up questions.\n"
+        "4) Do NOT mention rules, instructions, or your reasoning.\n"
+        "5) Provide only the final answer.\n"
     )
 
-    messages = [{"role": "system", "content": hard_rules}]
-    messages.extend(history)
-    messages.append({"role": "user", "content": user_message})
+    sys = (system_message or "").strip()
+    prompt = f"{hard_rules}\nSYSTEM NOTE: {sys}\n\n"
 
-    if hasattr(tokenizer, "apply_chat_template"):
-        return tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True,
-        )
-
-    prompt = hard_rules + "\n"
+    # Convert Gradio "messages" history into a readable transcript
     for m in history:
-        prompt += f"{m['role'].capitalize()}: {m['content']}\n"
-    prompt += f"User: {user_message}\nAssistant:"
+        role = (m.get("role") or "user").lower()
+        content = (m.get("content") or "").strip()
+        if not content:
+            continue
+        if role == "user":
+            prompt += f"User: {content}\n"
+        else:
+            prompt += f"Assistant: {content}\n"
+
+    prompt += f"User: {user_message.strip()}\nAssistant:"
     return prompt
 
 
+def clean_output(text: str) -> str:
+    """
+    Removes common fine-tune artifacts without being too aggressive.
+    """
+    # Remove leading parenthetical meta like: "(Answering in English...)"
+    text = re.sub(r"^\s*\(.*?\)\s*", "", text, flags=re.DOTALL)
+
+    # If the model starts adding "Question:" sections, cut everything after it
+    cut_markers = [
+        "\nQuestion:",
+        "\nNow, let's",
+        "\nNow let's",
+        "\nBased on the rules",
+        "\nAccording to the rules",
+    ]
+    for marker in cut_markers:
+        idx = text.lower().find(marker.lower())
+        if idx != -1:
+            text = text[:idx].strip()
+            break
+
+    return text.strip()
+
+
 def respond(
     message: str,
     history: List[Dict[str, str]],
@@ -93,7 +114,7 @@ def respond(
 ):
     load_model(hf_token.token if hf_token else None)
 
-    prompt = build_prompt(system_message, history, message)
+    prompt = build_plain_prompt(system_message, history, message)
 
     inputs = tokenizer(prompt, return_tensors="pt")
     inputs = {k: v.to(device) for k, v in inputs.items()}
@@ -104,13 +125,14 @@ def respond(
         skip_special_tokens=True,
     )
 
+    # Lower randomness to reduce “roleplay / training artifact” behavior
     gen_kwargs = dict(
         **inputs,
         max_new_tokens=int(max_tokens),
         do_sample=True,
-        temperature=0.3,          # 🔒 low temperature = less roleplay
-        top_p=0.85,
-        repetition_penalty=1.2,   # 🔒 stops looping
+        temperature=float(temperature),
+        top_p=float(top_p),
+        repetition_penalty=1.15,
         streamer=streamer,
         eos_token_id=tokenizer.eos_token_id,
         pad_token_id=tokenizer.eos_token_id,
@@ -119,20 +141,19 @@ def respond(
     thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
     thread.start()
 
-    output = ""
-    for text in streamer:
-        # HARD STOP if model tries to continue conversation
-        if any(bad in text.lower() for bad in [
-            "now let's",
-            "question:",
-            "based on the rules",
-            "according to the rules",
-            "let us",
-        ]):
+    out = ""
+    for chunk in streamer:
+        out += chunk
+
+        # Stream the cleaned output live
+        cleaned = clean_output(out)
+
+        # Hard stop if it starts self-questioning
+        if any(x in out.lower() for x in ["\nquestion:", "now, let's generate", "based on the rules"]):
+            yield cleaned
             break
 
-        output += text
-        yield output.strip()
+        yield cleaned
 
 
 chatbot = gr.ChatInterface(
@@ -142,7 +163,7 @@ chatbot = gr.ChatInterface(
     description=f"Running model: `{MODEL_ID}`",
     additional_inputs=[
         gr.Textbox(
-            value="Provide SOP-style answers for hotel housekeeping staff.",
+            value="Give SOP-style housekeeping answers. Use bullet points when helpful.",
             label="System message",
         ),
         gr.Slider(1, 1024, value=256, step=1, label="Max new tokens"),
@@ -155,7 +176,6 @@ with gr.Blocks() as demo:
     with gr.Sidebar():
         gr.LoginButton()
         gr.Markdown(f"**Model:** `{MODEL_ID}`")
-
     chatbot.render()
 
 if __name__ == "__main__":