Update app.py

app.py

@@ -2,39 +2,63 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Hugging Face repo + subfolder där den mergade modellen ligger
 MODEL_ID = "Jeppcode/ScalableLab2"
 SUBFOLDER = "merged-model-fp16"
 
 print(f"Loading model {MODEL_ID}/{SUBFOLDER} ...")
 
-# Ladda tokenizer och modell från subfoldern
+# Tokenizer
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID,
     subfolder=SUBFOLDER,
 )
 
+# Modell – fp16 + snålare CPU-load
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     subfolder=SUBFOLDER,
-    device_map="auto",   # på HF CPU-space hamnar den på cpu
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,
+    device_map="cpu",   # var explicit – allt på CPU
 )
+model.eval()
+
 
 def build_prompt(message, history):
     """
-    Bygger upp en lista av chat-meddelanden i samma format
-    som användes vid träning, och använder sedan chat_template.
+    history (Gradio 6) är en lista av dicts:
+    [
+      {"role": "user", "content": [...]},
+      {"role": "assistant", "content": [...]},
+      ...
+    ]
+    Vi mappar det till samma roll/text-format som vid träning.
     """
     messages = []
-    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})
 
+    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:
+                    # fallback om Gradio skickar annat format
+                    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})
+
+    # Lägg till nuvarande användarmeddelande
     messages.append({"role": "user", "content": message})
 
-    # Din tokenizer har en chat_template.jinja, så apply_chat_template ska fungera
     prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
@@ -42,6 +66,7 @@ def build_prompt(message, history):
     )
     return prompt
 
+
 def chat_fn(message, history):
     prompt = build_prompt(message, history)
 
@@ -50,11 +75,12 @@ def chat_fn(message, history):
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=256,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9,
+            max_new_tokens=64,          # kortare svar = mycket snabbare
+            do_sample=False,            # deterministiskt, billigare
+            temperature=None,
+            top_p=None,
             pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
         )
 
     generated = tokenizer.decode(
@@ -64,8 +90,10 @@ def chat_fn(message, history):
 
     return generated
 
+
 demo = gr.ChatInterface(
     fn=chat_fn,
+    type="messages",   # säg tydligt att vi använder messages-formatet
     title="Lab 2 – Fine-tuned merged model (fp16)",
     description=(
         "Chat with our fine-tuned Llama-based model, merged to fp16 and "