Update app.py

app.py

@@ -7,22 +7,23 @@ SUBFOLDER = "merged-model-fp16"
 
 print(f"Loading model {MODEL_ID}/{SUBFOLDER} ...")
 
-# Tokenizer
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID,
     subfolder=SUBFOLDER,
 )
 
-# Model – fp16 on CPU, memory-friendly
+# Load model (fp16 on CPU to fit in HF Space)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     subfolder=SUBFOLDER,
-    dtype=torch.float16,
+    dtype=torch.float16,          # use dtype (torch_dtype is deprecated)
     low_cpu_mem_usage=True,
     device_map="cpu",
 )
 model.eval()
 
+# Predefined “styles” as system prompts
 STYLE_SYSTEM_PROMPTS = {
     "Default": "You are a helpful, polite assistant.",
     "Short answer": (
@@ -41,23 +42,26 @@ STYLE_SYSTEM_PROMPTS = {
 
 def build_prompt(message, history, style):
     """
-    history är i 'messages'-format:
-      [{"role": "user", "content": "..."},
-       {"role": "assistant", "content": "..."}, ...]
-    Vi lägger till en systemprompt + tidigare meddelanden + nuvarande fråga.
+    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.
     """
     messages = []
 
-    # System / style prompt (läggs bara i prompten, inte i history)
+    # Add system / style message
     system_prompt = STYLE_SYSTEM_PROMPTS.get(style, STYLE_SYSTEM_PROMPTS["Default"])
     messages.append({"role": "system", "content": system_prompt})
 
-    # Tidigare dialog (bara user/assistant)
-    messages.extend(history)
+    # 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})
 
-    # Nuvarande användarmeddelande
+    # Current user message
     messages.append({"role": "user", "content": message})
 
+    # Use chat_template from your tokenizer
     prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
@@ -67,7 +71,7 @@ def build_prompt(message, history, style):
 
 
 def chat_fn(message, history, max_new_tokens, temperature, top_p, repetition_penalty, style):
-    # history är redan i messages-format, vi gör inte om det
+    # Build full prompt including history + style
     prompt = build_prompt(message, history, style)
 
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
@@ -80,7 +84,7 @@ def chat_fn(message, history, max_new_tokens, temperature, top_p, repetition_pen
         "repetition_penalty": float(repetition_penalty),
     }
 
-    # Greedy vs sampling beroende på temperatur
+    # Deterministic if temperature == 0, otherwise sampling
     if temperature <= 0.0:
         gen_kwargs.update(
             dict(
@@ -106,14 +110,11 @@ def chat_fn(message, history, max_new_tokens, temperature, top_p, repetition_pen
         skip_special_tokens=True,
     ).strip()
 
-    # Uppdatera history i messages-format:
-    new_history = history + [
-        {"role": "user", "content": message},
-        {"role": "assistant", "content": generated},
-    ]
+    # Update history in the default (user, bot) format
+    history = history + [[message, generated]]
 
-    # Töm textboxen, uppdatera chatten
-    return "", new_history
+    # Return empty textbox + updated chat history
+    return "", history
 
 
 with gr.Blocks() as demo:
@@ -121,14 +122,14 @@ with gr.Blocks() as demo:
         "# 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`.\n\n"
-        "Use the controls on the right like a DJ board to explore how decoding "
+        "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: chat
+        # Left side: chatbot
         with gr.Column(scale=3):
-            chatbot = gr.Chatbot(label="Chat", type="messages")
+            chatbot = gr.Chatbot(label="Chat")   # no 'type' argument
             msg = gr.Textbox(
                 label="Your message",
                 placeholder="Ask the model something...",
@@ -137,7 +138,7 @@ with gr.Blocks() as demo:
             send_btn = gr.Button("Send")
             clear_btn = gr.Button("Clear chat")
 
-        # Right: generation controls
+        # Right side: generation settings (DJ board)
         with gr.Column(scale=1):
             gr.Markdown("### Generation controls")
 
@@ -181,7 +182,7 @@ with gr.Blocks() as demo:
                 label="Answer style",
             )
 
-    # Koppla knappar / enter
+    # Hook up buttons / enter key
     send_btn.click(
         chat_fn,
         inputs=[msg, chatbot, max_new_tokens, temperature, top_p, repetition_penalty, style],