Update app.py

app.py

@@ -14,7 +14,6 @@ if not HF_TOKEN:
 print(f"Loading model: {MODEL_ID} ...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 
-# Critical: Ensure pad token is set
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
@@ -23,7 +22,7 @@ model = AutoModelForCausalLM.from_pretrained(
     token=HF_TOKEN,
     torch_dtype="auto",
     device_map="auto",
-    trust_remote_code=True,   # Often needed for Qwen models
+    trust_remote_code=True,
 )
 
 def build_messages(history: List[Tuple[str, str]], user_message: str):
@@ -39,9 +38,13 @@ def build_messages(history: List[Tuple[str, str]], user_message: str):
 def chat_fn(
     message: str,
     history: List[Tuple[str, str]],
-    temperature: float = 0.7,
-    max_new_tokens: int = 512,
+    temperature: float | None = 0.7,      # <-- default here
+    max_new_tokens: int | None = 512,     # <-- default here
 ) -> str:
+    # Handle None values (from example caching)
+    temperature = temperature if temperature is not None else 0.7
+    max_new_tokens = max_new_tokens if max_new_tokens is not None else 512
+
     messages = build_messages(history, message)
     
     prompt = tokenizer.apply_chat_template(
@@ -53,7 +56,6 @@ def chat_fn(
     inputs = tokenizer(prompt, return_tensors="pt")
     inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
-    # Stronger generation config to avoid NaN/inf in probabilities
     with torch.no_grad():
         output_ids = model.generate(
             **inputs,
@@ -65,10 +67,9 @@ def chat_fn(
             repetition_penalty=1.1,
             pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
-            renormalize_logits=True,      # Very important for stability
+            renormalize_logits=True,
         )
     
-    # Extract only the new tokens
     generated_ids = output_ids[0][inputs["input_ids"].shape[-1] :]
     response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
     
@@ -81,12 +82,14 @@ demo = gr.ChatInterface(
         gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature"),
         gr.Slider(32, 2048, value=512, step=32, label="Max New Tokens"),
     ],
+    additional_inputs_accordion=gr.Accordion("Generation Settings", open=False),
     title="Qwen Workflow Planner Chat",
     description=f"Model: {MODEL_ID}",
     examples=[
         ["Plan a simple content creation workflow"],
         ["How to automate a daily report generation process?"],
     ],
+    cache_examples=False,   # Recommended on HF Spaces with additional inputs
 )
 
 if __name__ == "__main__":