Views
No views yet
1import torch
2from transformers import AutoModel, AutoTokenizer
3
4model_path = "apple/DiffuCoder-7B-cpGRPO"
5model = AutoModel.from_pretrained(model_path, torch_dtype=torch.bfloat16, trust_remote_code=True)
6tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
7model = model.to("cuda").eval()
8
9query = "Write a function to find the shared elements from the given two lists."
10prompt = f"""<|im_start|>system
11You are a helpful assistant.<|im_end|>
12<|im_start|>user
13{query.strip()}
14<|im_end|>
15<|im_start|>assistant
16""" ## following the template of qwen; you can also use apply_chat_template function
17
18TOKEN_PER_STEP = 1 # diffusion timesteps * TOKEN_PER_STEP = total new tokens
19
20inputs = tokenizer(prompt, return_tensors="pt")
21input_ids = inputs.input_ids.to(device="cuda")
22attention_mask = inputs.attention_mask.to(device="cuda")
23
24output = model.diffusion_generate(
25 input_ids,
26 attention_mask=attention_mask,
27 max_new_tokens=256,
28 output_history=True,
29 return_dict_in_generate=True,
30 steps=256//TOKEN_PER_STEP,
31 temperature=0.4,
32 top_p=0.95,
33 alg="entropy",
34 alg_temp=0.,
35)
36generations = [
37 tokenizer.decode(g[len(p) :].tolist())
38 for p, g in zip(input_ids, output.sequences)
39]
40
41print(generations[0].split('<|dlm_pad|>')[0])