Update app.py

app.py

@@ -19,7 +19,7 @@ model = AutoModelForCausalLM.from_pretrained(
     subfolder=SUBFOLDER,
     dtype=torch.float16,
     low_cpu_mem_usage=True,
-    device_map="cpu",   # explicit: everything on CPU
+    device_map="cpu",
 )
 model.eval()
 
@@ -41,23 +41,19 @@ STYLE_SYSTEM_PROMPTS = {
 
 def build_prompt(message, history, style):
     """
-    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.
+    history är i 'messages'-format:
+      [{"role": "user", "content": "..."},
+       {"role": "assistant", "content": "..."}, ...]
+    Vi lägger till en systemprompt + tidigare meddelanden + nuvarande fråga.
     """
     messages = []
 
-    # System / style prompt
+    # System / style prompt (läggs bara i prompten, inte i history)
     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})
+    # Tidigare dialog (bara user/assistant)
+    messages.extend(history)
 
     # Nuvarande användarmeddelande
     messages.append({"role": "user", "content": message})
@@ -71,7 +67,7 @@ def build_prompt(message, history, style):
 
 
 def chat_fn(message, history, max_new_tokens, temperature, top_p, repetition_penalty, style):
-    # Bygg prompt med historik + style
+    # history är redan i messages-format, vi gör inte om det
     prompt = build_prompt(message, history, style)
 
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
@@ -110,9 +106,14 @@ def chat_fn(message, history, max_new_tokens, temperature, top_p, repetition_pen
         skip_special_tokens=True,
     ).strip()
 
-    # Lägg till ny rad i historiken och töm textboxen
-    history = history + [[message, generated]]
-    return "", history
+    # Uppdatera history i messages-format:
+    new_history = history + [
+        {"role": "user", "content": message},
+        {"role": "assistant", "content": generated},
+    ]
+
+    # Töm textboxen, uppdatera chatten
+    return "", new_history
 
 
 with gr.Blocks() as demo:
@@ -127,7 +128,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         # Left: chat
         with gr.Column(scale=3):
-            chatbot = gr.Chatbot(label="Chat")
+            chatbot = gr.Chatbot(label="Chat", type="messages")
             msg = gr.Textbox(
                 label="Your message",
                 placeholder="Ask the model something...",
@@ -180,7 +181,7 @@ with gr.Blocks() as demo:
                 label="Answer style",
             )
 
-    # Koppla knapparna
+    # Koppla knappar / enter
     send_btn.click(
         chat_fn,
         inputs=[msg, chatbot, max_new_tokens, temperature, top_p, repetition_penalty, style],