Update app.py

app.py

@@ -5,18 +5,14 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel
 
 # --- CONFIG ---
-BASE_MODEL   = "microsoft/Phi-4-mini-instruct"   # base
+BASE_MODEL = "microsoft/Phi-4-mini-instruct"
 ADAPTER_REPO = "richardprobe/phi4-mini-chris-assistant-richard-adapter"
 SYSTEM_PROMPT = "You are Richard. Be concise and casual."
-
-# Use 4-bit quantization for smaller GPU Spaces
 LOAD_4BIT = True
 
-
 def load_model():
     print("Loading tokenizer...")
     tok = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
-
     print("Loading base model...")
     kwargs = dict(device_map="auto")
     if LOAD_4BIT:
@@ -33,65 +29,95 @@ def load_model():
     base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, **kwargs)
 
     print("Loading adapter...")
+    # HF Hub auth if needed
     model = PeftModel.from_pretrained(base, ADAPTER_REPO, use_auth_token=os.getenv("HF_TOKEN"))
     model.eval()
-    return tok, model
 
+    # make sure pad token exists
+    if tok.pad_token_id is None:
+        tok.pad_token = tok.eos_token
 
-tok, model = load_model()
+    return tok, model
 
+tok, model = load_model()
 
-def chat_generate(history, temperature=0.7, top_p=0.95, max_new_tokens=256, repetition_penalty=1.1):
-    """
-    history: list[(user, assistant)] from gr.ChatInterface
-    Returns: assistant reply as a string
-    """
-    messages = []
+def _normalize_history(history):
+    """Accepts either tuples [(u,a), ...] or messages-style [{'role','content'}, ...]."""
+    msgs = []
     if SYSTEM_PROMPT:
-        messages.append({"role": "system", "content": SYSTEM_PROMPT})
-    for user, assistant in history:
-        if user:
-            messages.append({"role": "user", "content": user})
-        if assistant:
-            messages.append({"role": "assistant", "content": assistant})
+        msgs.append({"role": "system", "content": SYSTEM_PROMPT})
+
+    if not history:
+        return msgs
+
+    # messages-style
+    if isinstance(history[0], dict):
+        for m in history:
+            role = m.get("role")
+            content = m.get("content", "")
+            if isinstance(content, list):  # v5 can send [{"type":"text","text":"..."}]
+                content = "".join(
+                    c.get("text", "") if isinstance(c, dict) else str(c) for c in content
+                )
+            if role in {"user", "assistant", "system"}:
+                msgs.append({"role": role, "content": content})
+    else:
+        # tuples-style
+        for u, a in history:
+            if u:
+                msgs.append({"role": "user", "content": u})
+            if a:
+                msgs.append({"role": "assistant", "content": a})
+    return msgs
+
+def chat_generate(message, history, temperature=0.7, top_p=0.95, max_new_tokens=256, repetition_penalty=1.1):
+    # Build messages
+    messages = _normalize_history(history)
+    if message:
+        messages.append({"role": "user", "content": message})
 
     inputs = tok.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        return_tensors="pt"
+        messages, add_generation_prompt=True, return_tensors="pt"
     ).to(model.device)
 
     gen_kwargs = dict(
         max_new_tokens=int(max_new_tokens),
         temperature=float(temperature),
         top_p=float(top_p),
-        do_sample=(temperature > 0),
+        do_sample=float(temperature) > 0,
         repetition_penalty=float(repetition_penalty),
         eos_token_id=tok.eos_token_id,
-        pad_token_id=tok.eos_token_id,
+        pad_token_id=tok.pad_token_id,
     )
 
-    with torch.inference_mode(), torch.cuda.amp.autocast(enabled=torch.cuda.is_available(), dtype=torch.bfloat16):
-        output = model.generate(inputs, **gen_kwargs)
+    with torch.inference_mode():
+        with torch.cuda.amp.autocast(enabled=torch.cuda.is_available(), dtype=torch.bfloat16):
+            out = model.generate(inputs, **gen_kwargs)
 
-    gen_tokens = output[0][inputs.shape[-1]:]
+    gen_tokens = out[0][inputs.shape[-1]:]
     text = tok.decode(gen_tokens, skip_special_tokens=True, errors="ignore")
     return text.strip()
 
-
 demo = gr.ChatInterface(
     fn=chat_generate,
     title="Phi-4 Mini + LoRA Adapter (Chris style)",
     description="Base: microsoft/Phi-4-mini-instruct + your LoRA adapter. Style-tuned chat.",
+    additional_inputs=[
+        gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(0.5, 1.0, value=0.95, step=0.01, label="Top-p"),
+        gr.Slider(16, 512, value=256, step=16, label="Max new tokens"),
+        gr.Slider(1.0, 1.5, value=1.1, step=0.05, label="Repetition penalty"),
+    ],
+    # Each example is: [message, *additional_inputs]
     examples=[
         ["What are you up to?", 0.7, 0.95, 256, 1.1],
-        ["You coming?",      0.7, 0.95, 256, 1.1],
-        ["I'm on the can",        0.7, 0.95, 256, 1.1],
+        ["You coming?",         0.7, 0.95, 256, 1.1],
+        ["I'm on the can",      0.7, 0.95, 256, 1.1],
     ],
-    cache_examples=True  # (optional)
-
+    cache_examples=False,   # turn off while debugging; turn on later if you want
 )
 
-
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+    demo.queue(concurrency_count=1, max_size=8)
+    # Hide API docs to avoid the schema crash toast
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False, show_error=True)