wow we had bad inference code

README.md

@@ -43,21 +43,27 @@ messages = [
 ]
 
 # Generate reasoning
-reasoning_template = tokenizer.apply_chat_template(messages, tokenize=False, add_reasoning_prompt=True)
-reasoning_inputs = tokenizer(reasoning_template, return_tensors="pt").to(model.device)
-reasoning_ids = model.generate(**reasoning_inputs, max_new_tokens=MAX_REASONING_TOKENS)
-reasoning_output = tokenizer.decode(reasoning_ids[0, reasoning_inputs.input_ids.shape[1]:], skip_special_tokens=True)
+input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_reasoning_prompt=True, return_tensors="pt")
+output = model.generate(
+    input_ids.to("cuda"),
+    eos_token_id=tokenizer.eos_token_id,
+    max_new_tokens=MAX_REASONING_TOKENS,
+    do_sample=False,
+)
 
-print("REASONING: " + reasoning_output)
+print("REASONING: " + tokenizer.decode(output[0]))
 
 # Generate answer
 messages.append({"role": "reasoning", "content": reasoning_output})
-response_template = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-response_inputs = tokenizer(response_template, return_tensors="pt").to(model.device)
-response_ids = model.generate(**response_inputs, max_new_tokens=MAX_RESPONSE_TOKENS)
-response_output = tokenizer.decode(response_ids[0, response_inputs.input_ids.shape[1]:], skip_special_tokens=True)
+input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, return_tensors="pt")
+output = model.generate(
+    input_ids.to("cuda"),
+    eos_token_id=tokenizer.eos_token_id,
+    max_new_tokens=MAX_RESPONSE_TOKENS,
+    do_sample=False,
+)
 
-print("ANSWER: " + response_output)
+print("REASONING: " + tokenizer.decode(output[0]))
 ```
 
 - **Trained by:** [Piotr Zalewski](https://huggingface.co/lunahr)