Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3model_id = "lurepaper/LURE_5.1"
4
5model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto")
6tokenizer = AutoTokenizer.from_pretrained(model_id)
7
8prompt = "You are a Lua programming language expert. Please generate generate concise Lua code that produces the Lua 5.1 opcode OP_ADD. Use a print function call at the end to show the execution result of the opcode."
9
10messages = [
11 {"role": "user", "content": prompt},
12 ]
13
14text = tokenizer.apply_chat_template(
15 messages,
16 tokenize=False,
17 add_generation_prompt=True,
18 enable_thinking=True,
19)
20
21model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
22
23generated_ids = model.generate(
24 **model_inputs,
25 max_new_tokens=32768,
26)
27
28output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
29
30# parsing thinking content
31try:
32 # rindex finding 151668 (</think>)
33 index = len(output_ids) - output_ids[::-1].index(151668)
34except ValueError:
35 index = 0
36
37thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
38content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
39
40print("Thinking content:")
41print(thinking_content)
42
43print("Generated LuaGadget:")
44print(content)