updated

app.py

@@ -34,43 +34,58 @@ model.load_state_dict(state_dict, strict=False)
 model.eval()
 
 def predict(message, history):
-    # 1. Prepare token list
-    tokens = [tokenizer.bos_token_id]
-    for human, assistant in history:
-        if human:
-            tokens.extend([tokenizer.user_start_id] + tokenizer.encode(human) + [tokenizer.user_end_id])
-        if assistant:
-            tokens.extend([tokenizer.assistant_start_id] + tokenizer.encode(assistant) + [tokenizer.assistant_end_id])
-    
-    tokens.extend([tokenizer.user_start_id] + tokenizer.encode(message) + [tokenizer.user_end_id])
-    tokens.append(tokenizer.assistant_start_id)
-    
-    # --- THE FIX FOR ASSERTION ERROR ---
-    # The error 'assert isinstance(tokens, list)' happens here.
-    # We pass the tokens as a LIST, not a Tensor, to satisfy nanochat's requirements.
-    # -----------------------------------
-    
-    with torch.no_grad():
-        # Call generate with the LIST 'tokens'
-        output = model.generate(
-            tokens,              # Passing as list [] instead of torch.tensor([[]])
-            max_tokens=512, 
-            temperature=0.8,
-            top_k=40
-        )
+    try:
+        # 1. Prepare token list
+        tokens = [tokenizer.bos_token_id]
         
-        generated_text = ""
-        # The Traceback shows model.generate is a generator (streaming)
-        for token in output:
-            # Handle if token is an int or a single-element tensor
-            token_id = token if isinstance(token, int) else token.item()
-            char = tokenizer.decode([token_id])
+        # FIX: Robust history handling for Gradio 5/6
+        for entry in history:
+            # Handle list of dicts format: {"role": "user", "content": "..."}
+            if isinstance(entry, dict):
+                role = entry.get("role")
+                content = entry.get("content", "")
+                if role == "user":
+                    tokens.extend([tokenizer.user_start_id] + tokenizer.encode(content) + [tokenizer.user_end_id])
+                elif role == "assistant":
+                    tokens.extend([tokenizer.assistant_start_id] + tokenizer.encode(content) + [tokenizer.assistant_end_id])
             
-            if "<|assistant_end|>" in char:
-                break
+            # Handle old list of lists format: [user_msg, assistant_msg]
+            elif isinstance(entry, (list, tuple)):
+                human, assistant = entry[0], entry[1]
+                if human:
+                    tokens.extend([tokenizer.user_start_id] + tokenizer.encode(human) + [tokenizer.user_end_id])
+                if assistant:
+                    tokens.extend([tokenizer.assistant_start_id] + tokenizer.encode(assistant) + [tokenizer.assistant_end_id])
+
+        # Add current user prompt
+        tokens.extend([tokenizer.user_start_id] + tokenizer.encode(message) + [tokenizer.user_end_id])
+        tokens.append(tokenizer.assistant_start_id)
+        
+        # 2. Streaming Generation
+        with torch.no_grad():
+            # Pass as list to satisfy the nanochat assertion
+            output = model.generate(
+                tokens, 
+                max_tokens=512, 
+                temperature=0.8,
+                top_k=40
+            )
+            
+            generated_text = ""
+            for token in output:
+                token_id = token if isinstance(token, int) else token.item()
+                char = tokenizer.decode([token_id])
                 
-            generated_text += char
-            yield generated_text.strip()
+                # Stop if we hit the assistant end tag
+                if "<|assistant_end|>" in char or "<|end|>" in char:
+                    break
+                    
+                generated_text += char
+                yield generated_text.strip()
+
+    except Exception as e:
+        print(f"CRITICAL ERROR: {e}")
+        yield f"Toddler tantrum: {str(e)}"
 
 # Launching with Gradio 6.0 compatibility
 demo = gr.ChatInterface(