Update app.py

app.py

@@ -7,61 +7,57 @@ SUBFOLDER = "merged-model-fp16"
 
 print(f"Loading model {MODEL_ID}/{SUBFOLDER} ...")
 
-# Load tokenizer
+# Tokenizer
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID,
     subfolder=SUBFOLDER,
 )
 
-# Load model (fp16 on CPU to fit in HF Space)
+# Modell – fp16 och snålare på CPU
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     subfolder=SUBFOLDER,
-    dtype=torch.float16,          # use dtype (torch_dtype is deprecated)
+    dtype=torch.float16,        # samma som torch_dtype men utan varningen
     low_cpu_mem_usage=True,
-    device_map="cpu",
+    device_map="cpu",           # var explicit, allt på CPU
 )
 model.eval()
 
-# Predefined “styles” as system prompts
-STYLE_SYSTEM_PROMPTS = {
-    "Default": "You are a helpful, polite assistant.",
-    "Short answer": (
-        "You are a helpful assistant. Answer as concisely as possible, usually in 1–3 sentences."
-    ),
-    "Detailed explanation": (
-        "You are a helpful teaching assistant. Give clear, structured and detailed explanations, "
-        "often with bullet points or numbered steps when useful."
-    ),
-    "Step-by-step reasoning": (
-        "You are a careful problem solver. Think step by step and explain your reasoning clearly "
-        "before giving the final answer."
-    ),
-}
 
-
-def build_prompt(message, history, style):
+def build_prompt(message, history):
     """
-    history is a list of [user, bot] pairs (Gradio's default Chatbot format).
-    We convert it into a list of role/content messages for the chat template.
+    I Gradio 6 är history en lista av dicts:
+    [
+      {"role": "user", "content": [...]},
+      {"role": "assistant", "content": [...]},
+      ...
+    ]
+    Vi plockar ut texten och mappar till {role, content}.
     """
     messages = []
 
-    # Add system / style message
-    system_prompt = STYLE_SYSTEM_PROMPTS.get(style, STYLE_SYSTEM_PROMPTS["Default"])
-    messages.append({"role": "system", "content": system_prompt})
-
-    # Add previous conversation
-    for user_msg, bot_msg in history:
-        if user_msg is not None:
-            messages.append({"role": "user", "content": user_msg})
-        if bot_msg is not None:
-            messages.append({"role": "assistant", "content": bot_msg})
-
-    # Current user message
+    for msg in history:
+        role = msg.get("role")
+        content = msg.get("content", "")
+
+        # content kan vara en lista av blocks eller en sträng
+        if isinstance(content, list):
+            texts = []
+            for block in content:
+                if isinstance(block, dict) and block.get("type") == "text":
+                    texts.append(block.get("text", ""))
+                else:
+                    texts.append(str(block))
+            text = "\n".join(t for t in texts if t)
+        else:
+            text = str(content)
+
+        if text:
+            messages.append({"role": role, "content": text})
+
+    # nuvarande användarmeddelande
     messages.append({"role": "user", "content": message})
 
-    # Use chat_template from your tokenizer
     prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
@@ -70,130 +66,38 @@ def build_prompt(message, history, style):
     return prompt
 
 
-def chat_fn(message, history, max_new_tokens, temperature, top_p, repetition_penalty, style):
-    # Build full prompt including history + style
-    prompt = build_prompt(message, history, style)
+def chat_fn(message, history):
+    prompt = build_prompt(message, history)
 
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    gen_kwargs = {
-        **inputs,
-        "max_new_tokens": int(max_new_tokens),
-        "pad_token_id": tokenizer.eos_token_id,
-        "eos_token_id": tokenizer.eos_token_id,
-        "repetition_penalty": float(repetition_penalty),
-    }
-
-    # Deterministic if temperature == 0, otherwise sampling
-    if temperature <= 0.0:
-        gen_kwargs.update(
-            dict(
-                do_sample=False,
-                temperature=None,
-                top_p=None,
-            )
-        )
-    else:
-        gen_kwargs.update(
-            dict(
-                do_sample=True,
-                temperature=float(temperature),
-                top_p=float(top_p),
-            )
-        )
-
     with torch.no_grad():
-        outputs = model.generate(**gen_kwargs)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=64,          # kortare svar för snabbare CPU
+            do_sample=False,            # deterministiskt
+            temperature=None,
+            top_p=None,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
 
     generated = tokenizer.decode(
         outputs[0][inputs["input_ids"].shape[1]:],
         skip_special_tokens=True,
     ).strip()
 
-    # Update history in the default (user, bot) format
-    history = history + [[message, generated]]
-
-    # Return empty textbox + updated chat history
-    return "", history
+    return generated
 
 
-with gr.Blocks() as demo:
-    gr.Markdown(
-        "# Lab 2 – Fine-tuned merged model (fp16)\n"
+demo = gr.ChatInterface(
+    fn=chat_fn,
+    title="Lab 2 – Fine-tuned merged model (fp16)",
+    description=(
         "Chat with our fine-tuned Llama-based model, merged to fp16 and "
-        "loaded from `Jeppcode/ScalableLab2/merged-model-fp16`.\n\n"
-        "Use the controls on the right like a DJ board to see how decoding "
-        "settings change the behaviour of the model."
-    )
-
-    with gr.Row():
-        # Left side: chatbot
-        with gr.Column(scale=3):
-            chatbot = gr.Chatbot(label="Chat")   # no 'type' argument
-            msg = gr.Textbox(
-                label="Your message",
-                placeholder="Ask the model something...",
-                lines=3,
-            )
-            send_btn = gr.Button("Send")
-            clear_btn = gr.Button("Clear chat")
-
-        # Right side: generation settings (DJ board)
-        with gr.Column(scale=1):
-            gr.Markdown("### Generation controls")
-
-            max_new_tokens = gr.Slider(
-                minimum=16,
-                maximum=256,
-                value=64,
-                step=8,
-                label="Max new tokens (response length)",
-            )
-            temperature = gr.Slider(
-                minimum=0.0,
-                maximum=1.5,
-                value=0.0,
-                step=0.1,
-                label="Temperature (0 = deterministic, higher = more random)",
-            )
-            top_p = gr.Slider(
-                minimum=0.1,
-                maximum=1.0,
-                value=0.9,
-                step=0.05,
-                label="Top-p (nucleus sampling)",
-            )
-            repetition_penalty = gr.Slider(
-                minimum=0.8,
-                maximum=1.3,
-                value=1.0,
-                step=0.05,
-                label="Repetition penalty",
-            )
-
-            style = gr.Radio(
-                choices=[
-                    "Default",
-                    "Short answer",
-                    "Detailed explanation",
-                    "Step-by-step reasoning",
-                ],
-                value="Detailed explanation",
-                label="Answer style",
-            )
-
-    # Hook up buttons / enter key
-    send_btn.click(
-        chat_fn,
-        inputs=[msg, chatbot, max_new_tokens, temperature, top_p, repetition_penalty, style],
-        outputs=[msg, chatbot],
-    )
-    msg.submit(
-        chat_fn,
-        inputs=[msg, chatbot, max_new_tokens, temperature, top_p, repetition_penalty, style],
-        outputs=[msg, chatbot],
-    )
-    clear_btn.click(lambda: [], None, chatbot)
+        "loaded from Jeppcode/ScalableLab2/merged-model-fp16."
+    ),
+)
 
 if __name__ == "__main__":
     demo.launch()