Update app.py

app.py

@@ -13,49 +13,53 @@ tokenizer = AutoTokenizer.from_pretrained(
     subfolder=SUBFOLDER,
 )
 
-# Modell – fp16 och snålare på CPU
+# Model – fp16 on CPU, memory-friendly
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     subfolder=SUBFOLDER,
-    dtype=torch.float16,        # samma som torch_dtype men utan varningen
+    dtype=torch.float16,
     low_cpu_mem_usage=True,
-    device_map="cpu",           # var explicit, allt på CPU
+    device_map="cpu",   # explicit: everything on CPU
 )
 model.eval()
 
+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):
+def build_prompt(message, history, style):
     """
-    I Gradio 6 är history en lista av dicts:
-    [
-      {"role": "user", "content": [...]},
-      {"role": "assistant", "content": [...]},
-      ...
-    ]
-    Vi plockar ut texten och mappar till {role, content}.
+    history från gr.Chatbot är en lista av [user, bot]-par:
+      [["Hi", "Hello!"], ["Explain X", "Answer..."], ...]
+    Vi konverterar till ett messages-format och lägger på en systemprompt
+    baserat på vald style.
     """
     messages = []
 
-    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
+    # System / style prompt
+    system_prompt = STYLE_SYSTEM_PROMPTS.get(style, STYLE_SYSTEM_PROMPTS["Default"])
+    messages.append({"role": "system", "content": system_prompt})
+
+    # Tidigare dialog
+    for user_msg, bot_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            messages.append({"role": "assistant", "content": bot_msg})
+
+    # Nuvarande användarmeddelande
     messages.append({"role": "user", "content": message})
 
     prompt = tokenizer.apply_chat_template(
@@ -66,38 +70,128 @@ def build_prompt(message, history):
     return prompt
 
 
-def chat_fn(message, history):
-    prompt = build_prompt(message, history)
+def chat_fn(message, history, max_new_tokens, temperature, top_p, repetition_penalty, style):
+    # Bygg prompt med historik + style
+    prompt = build_prompt(message, history, style)
 
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    with torch.no_grad():
-        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,
+    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),
+    }
+
+    # Greedy vs sampling beroende på temperatur
+    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)
+
     generated = tokenizer.decode(
         outputs[0][inputs["input_ids"].shape[1]:],
         skip_special_tokens=True,
     ).strip()
 
-    return generated
+    # Lägg till ny rad i historiken och töm textboxen
+    history = history + [[message, generated]]
+    return "", history
 
 
-demo = gr.ChatInterface(
-    fn=chat_fn,
-    title="Lab 2 – Fine-tuned merged model (fp16)",
-    description=(
+with gr.Blocks() as demo:
+    gr.Markdown(
+        "# Lab 2 – Fine-tuned merged model (fp16)\n"
         "Chat with our fine-tuned Llama-based model, merged to fp16 and "
-        "loaded from Jeppcode/ScalableLab2/merged-model-fp16."
-    ),
-)
+        "loaded from `Jeppcode/ScalableLab2/merged-model-fp16`.\n\n"
+        "Use the controls on the right like a DJ board to explore how decoding "
+        "settings change the behaviour of the model."
+    )
+
+    with gr.Row():
+        # Left: chat
+        with gr.Column(scale=3):
+            chatbot = gr.Chatbot(label="Chat")
+            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: generation controls
+        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",
+            )
+
+    # Koppla knapparna
+    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)
 
 if __name__ == "__main__":
     demo.launch()