Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_name = "Lyte/Llama-3.2-3B-Overthinker"
4tokenizer = AutoTokenizer.from_pretrained(model_name)
5model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto")
6
7def generate_response(prompt, max_tokens=16384, temperature=0.8, top_p=0.95, repeat_penalty=1.1, num_steps=3):
8 messages = [{"role": "user", "content": prompt}]
9
10 # Generate reasoning
11 reasoning_template = tokenizer.apply_chat_template(messages, tokenize=False, add_reasoning_prompt=True)
12 reasoning_inputs = tokenizer(reasoning_template, return_tensors="pt").to(model.device)
13
14 reasoning_ids = model.generate(
15 **reasoning_inputs,
16 max_new_tokens=max_tokens // 3,
17 temperature=temperature,
18 top_p=top_p,
19 repetition_penalty=repeat_penalty
20 )
21 reasoning_output = tokenizer.decode(reasoning_ids[0, reasoning_inputs.input_ids.shape[1]:], skip_special_tokens=True)
22
23 # Generate thinking (step-by-step and verifications)
24 messages.append({"role": "reasoning", "content": reasoning_output})
25 thinking_template = tokenizer.apply_chat_template(messages, tokenize=False, add_thinking_prompt=True, num_steps=num_steps)
26 thinking_inputs = tokenizer(thinking_template, return_tensors="pt").to(model.device)
27
28 thinking_ids = model.generate(
29 **thinking_inputs,
30 max_new_tokens=max_tokens // 3,
31 temperature=temperature,
32 top_p=top_p,
33 repetition_penalty=repeat_penalty
34 )
35 thinking_output = tokenizer.decode(thinking_ids[0, thinking_inputs.input_ids.shape[1]:], skip_special_tokens=True)
36
37 # Generate final answer
38 messages.append({"role": "thinking", "content": thinking_output})
39 answer_template = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
40 answer_inputs = tokenizer(answer_template, return_tensors="pt").to(model.device)
41
42 answer_ids = model.generate(
43 **answer_inputs,
44 max_new_tokens=max_tokens // 3,
45 temperature=temperature,
46 top_p=top_p,
47 repetition_penalty=repeat_penalty
48 )
49 answer_output = tokenizer.decode(answer_ids[0, answer_inputs.input_ids.shape[1]:], skip_special_tokens=True)
50 return reasoning_output, thinking_output, answer_output
51
52# Example usage:
53prompt = "Explain the process of photosynthesis."
54response = generate_response(prompt, num_steps=5)
55
56print("Response:", response)