Upload mirothinker server.py

mirothinker server.py

@@ -69,12 +69,17 @@ def chat(req: ChatRequest):
     continue_final = bool(req.continue_final_message)
 
     # Build the prompt with the model's own chat template.
-    input_ids = tokenizer.apply_chat_template(
+    # return_dict=True gives a BatchEncoding (input_ids + attention_mask);
+    # handle it as a dict so .shape / generate work correctly.
+    enc = tokenizer.apply_chat_template(
         req.messages,
         add_generation_prompt=not continue_final,
         continue_final_message=continue_final,
         return_tensors="pt",
-    ).to(model.device)
+        return_dict=True,
+    )
+    enc = {k: v.to(model.device) for k, v in enc.items()}
+    input_ids = enc["input_ids"]
 
     prompt_tokens = input_ids.shape[-1]
     max_new = req.max_completion_tokens or req.max_tokens or 16384
@@ -90,7 +95,7 @@ def chat(req: ChatRequest):
         gen_kwargs["repetition_penalty"] = req.repetition_penalty
 
     with torch.inference_mode():
-        out = model.generate(input_ids, **gen_kwargs)
+        out = model.generate(**enc, **gen_kwargs)
 
     new_tokens = out[0][prompt_tokens:]
     text = tokenizer.decode(new_tokens, skip_special_tokens=True)