Update app.py

app.py

@@ -23,8 +23,24 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
+# Några stil-lägen som systemprompter
+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:
     [
@@ -32,10 +48,16 @@ def build_prompt(message, history):
       {"role": "assistant", "content": [...]},
       ...
     ]
-    Vi plockar ut texten och mappar till {role, content}.
+    Vi plockar ut texten och mappar till {role, content}, och lägger till en systemprompt
+    baserat på vald 'style'.
     """
     messages = []
 
+    # Lägg till system / style prompt
+    system_prompt = STYLE_SYSTEM_PROMPTS.get(style, STYLE_SYSTEM_PROMPTS["Default"])
+    messages.append({"role": "system", "content": system_prompt})
+
+    # Tidigare historik
     for msg in history:
         role = msg.get("role")
         content = msg.get("content", "")
@@ -52,7 +74,7 @@ def build_prompt(message, history):
         else:
             text = str(content)
 
-        if text:
+        if text and role in ("user", "assistant", "system"):
             messages.append({"role": role, "content": text})
 
     # nuvarande användarmeddelande
@@ -66,37 +88,107 @@ 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 + stil
+    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),
+    }
+
+    # Deterministisk om temperature == 0, annars 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)
+
     generated = tokenizer.decode(
         outputs[0][inputs["input_ids"].shape[1]:],
         skip_special_tokens=True,
     ).strip()
 
+    # ChatInterface sköter history själv, vi returnerar bara svaret
     return generated
 
 
+# DJ-reglagen (extra inputs till ChatInterface)
+max_new_tokens_slider = gr.Slider(
+    minimum=16,
+    maximum=256,
+    value=64,
+    step=8,
+    label="Max new tokens (response length)",
+)
+temperature_slider = gr.Slider(
+    minimum=0.0,
+    maximum=1.5,
+    value=0.0,
+    step=0.1,
+    label="Temperature (0 = deterministic, higher = more random)",
+)
+top_p_slider = gr.Slider(
+    minimum=0.1,
+    maximum=1.0,
+    value=0.9,
+    step=0.05,
+    label="Top-p (nucleus sampling)",
+)
+repetition_penalty_slider = gr.Slider(
+    minimum=0.8,
+    maximum=1.3,
+    value=1.0,
+    step=0.05,
+    label="Repetition penalty",
+)
+style_radio = gr.Radio(
+    choices=[
+        "Default",
+        "Short answer",
+        "Detailed explanation",
+        "Step-by-step reasoning",
+    ],
+    value="Detailed explanation",
+    label="Answer style",
+)
+
 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."
+        "loaded from Jeppcode/ScalableLab2/merged-model-fp16.\n\n"
+        "Use the controls in the accordion below like a DJ board to tweak "
+        "response length, randomness and style."
     ),
+    additional_inputs=[
+        max_new_tokens_slider,
+        temperature_slider,
+        top_p_slider,
+        repetition_penalty_slider,
+        style_radio,
+    ],
+    additional_inputs_accordion="Generation controls",
 )
 
 if __name__ == "__main__":