Update app.py

app.py

@@ -1,70 +1,80 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 MODEL_NAME = "basmala12/smollm_finetuning5"
 
-# Load model & tokenizer
+# Load tokenizer & model once at startup (on CPU)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+model.eval()
 
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-)
 
 def respond(message, history, system_message, max_tokens, temperature, top_p):
+    """
+    ChatInterface (type="messages") passes:
+      - message: current user message (str)
+      - history: list of dicts: [{"role": "...", "content": "..."}, ...]
+      - system_message, max_tokens, temperature, top_p: from additional_inputs
+    We return a single string: the assistant reply.
+    """
 
-    # Build chat messages correctly
-    msgs = [{"role": "system", "content": system_message}]
+    # Build full conversation for the chat template
+    messages = [{"role": "system", "content": system_message}]
 
-    # History as proper chat
-    for user_msg, bot_msg in history:
-        msgs.append({"role": "user", "content": user_msg})
-        msgs.append({"role": "assistant", "content": bot_msg})
+    # history is a list of {"role": "user"/"assistant", "content": str}
+    # We append it as-is to preserve previous turns
+    messages.extend(history)
 
-    # Add the new user message
-    msgs.append({"role": "user", "content": message})
+    # Add the new user question
+    messages.append({"role": "user", "content": message})
 
-    # Apply chat template
+    # Turn into model prompt using the tokenizer's chat template
     prompt = tokenizer.apply_chat_template(
-        msgs,
+        messages,
         tokenize=False,
         add_generation_prompt=True,
     )
 
-    # Generate output
-    out = pipe(
-        prompt,
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        do_sample=True,
-    )[0]["generated_text"]
+    # Tokenize
+    inputs = tokenizer(prompt, return_tensors="pt")
 
-    # Extract only the assistant answer
-    if "<|im_start|>assistant" in out:
-        out = out.split("<|im_start|>assistant", 1)[-1]
-    out = out.replace("<|im_end|>", "").strip()
+    # Generate continuation (new assistant answer only)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=True,
+            temperature=float(temperature),
+            top_p=float(top_p),
+        )
 
-    # Enforce short answer + brief reasoning
-    # (additional safety)
-    if len(out.split()) > 45:
-        out = " ".join(out.split()[:45]) + " ..."
+    # Slice off the prompt tokens, keep only new tokens
+    generated_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+    answer = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
 
-    return out
+    # Optional: enforce "short answer + brief reasoning"
+    words = answer.split()
+    if len(words) > 60:
+        answer = " ".join(words[:60]) + " ..."
 
+    return answer
 
 
 chatbot = gr.ChatInterface(
     fn=respond,
     type="messages",
     additional_inputs=[
-        gr.Textbox("Give short answers with brief logical reasoning.", label="System message"),
+        gr.Textbox(
+            value="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"),
     ],
+    title="SmolLM2 – Short Reasoning Chatbot",
+    description="Fine-tuned SmolLM2 (basmala12/smollm_finetuning5) that gives short answers with brief logical reasoning.",
 )
 
 if __name__ == "__main__":