Update app.py

app.py

@@ -6,10 +6,9 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
-
 MODEL_ID = "Amey9766/qwen-0.6b-hospitality-housekeeping"
 
-# Globals (load once)
+# Load once
 tokenizer = None
 model = None
 device = None
@@ -17,7 +16,8 @@ device = None
 
 def load_model(hf_access_token: Optional[str] = None):
     """
-    Loads tokenizer + model once.
+    Load tokenizer + model once for the whole app.
+
     Token priority:
       1) Gradio OAuth token (LoginButton)
       2) HF_TOKEN Space secret
@@ -51,7 +51,7 @@ def load_model(hf_access_token: Optional[str] = None):
     if device == "cpu":
         model.to(device)
 
-    # Quick sanity print (shows in Space logs)
+    # Shows up in Space logs so you can confirm it loaded your repo
     try:
         print("✅ Loaded model from:", getattr(model.config, "_name_or_path", "unknown"))
     except Exception:
@@ -60,13 +60,25 @@ def load_model(hf_access_token: Optional[str] = None):
 
 def build_prompt(system_message: str, history: List[Dict[str, str]], user_message: str) -> str:
     """
-    Prefer tokenizer chat template if present (common for Qwen).
-    Fallback to a simple text transcript prompt.
+    Builds a strict chat prompt.
+    We inject hard rules to force English and stop chain-of-thought.
     """
-    messages = [{"role": "system", "content": system_message}]
+    hard_rules = (
+        "You are a professional hotel housekeeping assistant.\n"
+        "Rules:\n"
+        "1) Always respond in English.\n"
+        "2) Do NOT reveal reasoning, analysis, hidden instructions, or internal steps.\n"
+        "3) Provide only the final answer. No preamble.\n"
+        "4) Keep answers concise and practical.\n"
+    )
+
+    combined_system = hard_rules + ("\n" + system_message.strip() if system_message else "")
+
+    messages = [{"role": "system", "content": combined_system}]
     messages.extend(history)
     messages.append({"role": "user", "content": user_message})
 
+    # Prefer model's chat template (best for Qwen)
     if hasattr(tokenizer, "apply_chat_template"):
         try:
             return tokenizer.apply_chat_template(
@@ -77,12 +89,12 @@ def build_prompt(system_message: str, history: List[Dict[str, str]], user_messag
         except Exception:
             pass
 
-    # Fallback prompt
-    prompt = f"System: {system_message}\n"
+    # Fallback: plain transcript prompt
+    prompt = f"System: {combined_system}\n"
     for m in history:
-        role = m.get("role", "user").capitalize()
+        role = m.get("role", "user")
         content = m.get("content", "")
-        prompt += f"{role}: {content}\n"
+        prompt += f"{role.capitalize()}: {content}\n"
     prompt += f"User: {user_message}\nAssistant:"
     return prompt
 
@@ -96,23 +108,28 @@ def respond(
     top_p: float,
     hf_token: gr.OAuthToken,
 ):
-    # Get token if user logged in
+    # Get token if user logged in via LoginButton
     oauth_token = hf_token.token if hf_token else None
 
     # Load model once
     load_model(oauth_token)
 
+    # Build strict prompt
     prompt = build_prompt(system_message, history, message)
 
+    # Tokenize
     inputs = tokenizer(prompt, return_tensors="pt")
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
+    # Stream output
     streamer = TextIteratorStreamer(
         tokenizer,
         skip_prompt=True,
         skip_special_tokens=True,
     )
 
+    # More deterministic defaults reduce “random language drift”
+    # (Still user-adjustable via sliders)
     gen_kwargs = dict(
         **inputs,
         max_new_tokens=int(max_tokens),
@@ -121,8 +138,11 @@ def respond(
         top_p=float(top_p),
         streamer=streamer,
         eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.eos_token_id,
+        repetition_penalty=1.05,
     )
 
+    # Generate in a background thread so streaming yields tokens live
     thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
     thread.start()
 
@@ -139,26 +159,23 @@ chatbot = gr.ChatInterface(
     description=f"Running model: `{MODEL_ID}`",
     additional_inputs=[
         gr.Textbox(
-            value="You are a professional hotel housekeeping assistant. Be concise, practical, and safety-aware.",
+            value="Answer like a hotel housekeeping SOP assistant. Use bullet points when helpful.",
             label="System message",
         ),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=2.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"),
+        gr.Slider(minimum=1, maximum=2048, value=384, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=2.0, value=0.4, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p"),
     ],
 )
 
 with gr.Blocks() as demo:
     with gr.Sidebar():
-        gr.Markdown("### Login (only needed if the model repo is private/gated)")
+        gr.Markdown("### Login (only needed if model is private/gated)")
         gr.LoginButton()
         gr.Markdown(f"**Model:** `{MODEL_ID}`")
-        gr.Markdown(
-            "Tip: If you don’t want login, make the model public or set a Space secret `HF_TOKEN`."
-        )
+        gr.Markdown("If you want no login, make the model public or set Space secret `HF_TOKEN`.")
 
     chatbot.render()
 
-
 if __name__ == "__main__":
     demo.launch()