Upload app.py

app.py

@@ -49,17 +49,27 @@ def sanitize(text: str) -> str:
 MAX_CTX = 4096
 
 def chat(raw_input, history):
-    clean = sanitize(raw_input)
-    if not clean.strip() or clean == "[redacted]":
-        return "I’m not sure what you meant—could you try phrasing it differently?", history
+    # history is a list of dicts: {"role": "user"/"assistant", "content": str}
+    # Initialise history list if None
+    history = history or []
 
-    # Build prompt from sanitized history + this turn
-    hist = ""
-    for prev_raw, prev_resp in history:
-        hist += f"User: {sanitize(prev_raw)}\nAssistant: {prev_resp}\n"
+    # Append new raw user message
+    history.append({"role": "user", "content": raw_input})
 
-    prompt = "\n".join(SYSTEM_PROMPTS) + "\n\n" + hist
-    prompt += f"User: {clean}\nAssistant:"
+    # Sanitize the new user input for the model
+    clean = sanitize(raw_input)
+    if not clean.strip() or clean == "[redacted]":
+        history.append({"role": "assistant", "content": "I’m not sure what you meant—could you try phrasing it differently?"})
+        return history, history
+
+    # Build prompt: system prompts + conversation history
+    prompt = "\n".join(SYSTEM_PROMPTS) + "\n\n"
+    for turn in history:
+        if turn["role"] == "user":
+            prompt += f"User: {sanitize(turn['content'])}\n"
+        else:
+            prompt += f"Assistant: {turn['content']}\n"
+    prompt += "Assistant:"
 
     # Tokenise & manually truncate to last MAX_CTX tokens
     all_ids = tokenizer.encode(prompt, add_special_tokens=False)
@@ -70,7 +80,7 @@ def chat(raw_input, history):
         "attention_mask": torch.tensor([[1] * len(all_ids)])
     }
 
-    # Generate
+    # Generate response
     out = model.generate(
         **inputs,
         max_new_tokens=200,
@@ -81,8 +91,9 @@ def chat(raw_input, history):
     seen = tokenizer.decode(all_ids, skip_special_tokens=True)
     resp = full[len(seen):].strip()
 
-    history.append((raw_input, resp))
-    return resp, history
+    # Append model reply
+    history.append({"role": "assistant", "content": resp})
+    return history, history
 
 # ————————————————————————————————————————————————————————————————
 # Launch Gradio App