Update app.py

app.py

@@ -8,24 +8,15 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStream
 
 MODEL_ID = "Amey9766/qwen-0.6b-hospitality-housekeeping"
 
-# Load once
 tokenizer = None
 model = None
 device = None
 
 
 def load_model(hf_access_token: Optional[str] = None):
-    """
-    Load tokenizer + model once for the whole app.
-
-    Token priority:
-      1) Gradio OAuth token (LoginButton)
-      2) HF_TOKEN Space secret
-      3) None (public model)
-    """
     global tokenizer, model, device
 
-    if tokenizer is not None and model is not None:
+    if tokenizer and model:
         return
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -51,50 +42,42 @@ def load_model(hf_access_token: Optional[str] = None):
     if device == "cpu":
         model.to(device)
 
-    # 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:
-        pass
+    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:
     """
-    Builds a strict chat prompt.
-    We inject hard rules to force English and stop chain-of-thought.
+    HARD rules to stop:
+    - self questioning
+    - rule narration
+    - exam-style continuation
     """
     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"
+        "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"
     )
 
-    combined_system = hard_rules + ("\n" + system_message.strip() if system_message else "")
-
-    messages = [{"role": "system", "content": combined_system}]
+    messages = [{"role": "system", "content": hard_rules}]
     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(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True,
-            )
-        except Exception:
-            pass
-
-    # Fallback: plain transcript prompt
-    prompt = f"System: {combined_system}\n"
+        return tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+
+    prompt = hard_rules + "\n"
     for m in history:
-        role = m.get("role", "user")
-        content = m.get("content", "")
-        prompt += f"{role.capitalize()}: {content}\n"
+        prompt += f"{m['role'].capitalize()}: {m['content']}\n"
     prompt += f"User: {user_message}\nAssistant:"
     return prompt
 
@@ -108,48 +91,48 @@ def respond(
     top_p: float,
     hf_token: gr.OAuthToken,
 ):
-    # 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)
+    load_model(hf_token.token if hf_token else None)
 
-    # 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),
         do_sample=True,
-        temperature=float(temperature),
-        top_p=float(top_p),
+        temperature=0.3,          # 🔒 low temperature = less roleplay
+        top_p=0.85,
+        repetition_penalty=1.2,   # 🔒 stops looping
         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()
 
-    partial = ""
-    for token_text in streamer:
-        partial += token_text
-        yield partial
+    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",
+        ]):
+            break
+
+        output += text
+        yield output.strip()
 
 
 chatbot = gr.ChatInterface(
@@ -159,21 +142,19 @@ chatbot = gr.ChatInterface(
     description=f"Running model: `{MODEL_ID}`",
     additional_inputs=[
         gr.Textbox(
-            value="Answer like a hotel housekeeping SOP assistant. Use bullet points when helpful.",
+            value="Provide SOP-style answers for hotel housekeeping staff.",
             label="System message",
         ),
-        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"),
+        gr.Slider(1, 1024, value=256, step=1, label="Max new tokens"),
+        gr.Slider(0.1, 1.0, value=0.3, step=0.05, label="Temperature"),
+        gr.Slider(0.5, 1.0, value=0.85, step=0.05, label="Top-p"),
     ],
 )
 
 with gr.Blocks() as demo:
     with gr.Sidebar():
-        gr.Markdown("### Login (only needed if model is private/gated)")
         gr.LoginButton()
         gr.Markdown(f"**Model:** `{MODEL_ID}`")
-        gr.Markdown("If you want no login, make the model public or set Space secret `HF_TOKEN`.")
 
     chatbot.render()