Update app.py

app.py

@@ -1,76 +1,56 @@
 import gradio as gr
-import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# ==============================
-# CONFIG: YOUR MODEL HERE
-# ==============================
-MODEL_NAME = "smol-medical-meadow-FT"   # <--- change if needed
+# Change this to your model path / HF repo
+# e.g. "basmala12/smol-medical-meadow-FT" if it's on Hugging Face
+MODEL_NAME = "smol-medical-meadow-FT"
 
-
-# ==============================
-# LOAD MODEL + TOKENIZER
-# ==============================
+# Load model & tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float32
-)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 pipe = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    return_full_text=False,   # only return assistant's continuation
 )
 
-
-# ==============================
-# CHAT FUNCTION
-# ==============================
 def respond(message, history, system_message, max_tokens, temperature, top_p):
-
-    # Build the plain conversation (SmolLM2 style)
-    prompt = f"System: {system_message}\n\n"
-
-    for turn in history:
-        prompt += f"User: {turn['user']}\n"
-        prompt += f"Assistant: {turn['assistant']}\n"
-
-    prompt += f"User: {message}\nAssistant:"
-
-    # Run generation
-    response = pipe(
+    # Build chat-style messages
+    messages = [{"role": "system", "content": system_message}]
+    messages.extend(history)
+    messages.append({"role": "user", "content": message})
+
+    # Convert to model-specific chat template
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    # Generate
+    out = pipe(
         prompt,
         max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
         do_sample=True,
-        eos_token_id=tokenizer.eos_token_id,
     )[0]["generated_text"]
 
-    # Clean trailing text
-    # Stop if it starts generating new questions
-    for stop in ["User:", "System:", "Q:", "\n\n"]:
-        if stop in response:
-            response = response.split(stop)[0].strip()
-
-    return response.strip()
+    # Extract assistant-only text
+    if "<|im_start|>assistant" in out:
+        out = out.split("<|im_start|>assistant", 1)[-1]
+    out = out.replace("<|im_end|>", "").strip()
 
+    return out
 
-# ==============================
-# GRADIO UI
-# ==============================
 chatbot = gr.ChatInterface(
     fn=respond,
     type="messages",
     additional_inputs=[
-        gr.Textbox("You are a careful medical assistant. Answer clearly and safely.", label="System message"),
-        gr.Slider(10, 512, value=150, step=5, label="Max new tokens"),
-        gr.Slider(0.1, 2.0, value=0.7, step=0.05, label="Temperature"),
+        gr.Textbox("Give short answers with brief logical reasoning.", label="System message"),
+        gr.Slider(1, 512, value=256, step=1, label="Max new tokens"),
+        gr.Slider(0.1, 4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"),
     ],
 )