support force thinking

This commit is contained in:
liam 2025-02-12 12:43:53 +08:00
parent 6f3a39be08
commit 4385e85096
3 changed files with 7 additions and 5 deletions

View file

@ -160,7 +160,7 @@ def local_chat(
messages, add_generation_prompt=True, return_tensors="pt"
)
if force_think:
token_thinks = torch.tensor([tokenizer.encode("<think>\\n",add_special_tokens=False)])
token_thinks = torch.tensor([tokenizer.encode("<think>\\n",add_special_tokens=False)],device=input_tensor.device)
input_tensor = torch.cat(
[input_tensor, token_thinks], dim=1
)