Views
No views yet
nvidia/OpenMathInstruct-2 dataset, covering algebra, calculus, and probability.1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4# 1. Load Model
5model_id = "PrimeTJ/Emo-v1"
6tokenizer = AutoTokenizer.from_pretrained(model_id)
7model = AutoModelForCausalLM.from_pretrained(
8 model_id,
9 torch_dtype=torch.float16,
10 device_map="auto"
11)
12
13# 2. Define the Prompt
14system_prompt = "You are a helpful math assistant. Think step by step."
15user_prompt = "A bat and a ball cost $1.10 in total. The bat costs $1.00 more than the ball. How much does the ball cost?"
16
17messages = [
18 {"role": "system", "content": system_prompt},
19 {"role": "user", "content": user_prompt}
20]
21
22# 3. Generate
23text = tokenizer.apply_chat_template(
24 messages,
25 tokenize=False,
26 add_generation_prompt=True
27)
28model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
29
30generated_ids = model.generate(
31 **model_inputs,
32 max_new_tokens=1024,
33 temperature=0.6 # Low temperature for logic
34)
35
36response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
37print(response)