Views
No views yet
transformers==4.46.2 and torch==2.5.1. Here is an example to use Dream-Coder 7B:1import torch
2from transformers import AutoModel, AutoTokenizer
3
4model_path = "Dream-org/Dream-Coder-v0-Instruct-7B"
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
9messages = [
10 {"role": "user", "content": "Write a quick sort algorithm."}
11]
12inputs = tokenizer.apply_chat_template(
13 messages, return_tensors="pt", return_dict=True, add_generation_prompt=True
14)
15input_ids = inputs.input_ids.to(device="cuda")
16attention_mask = inputs.attention_mask.to(device="cuda")
17
18output = model.diffusion_generate(
19 input_ids,
20 attention_mask=attention_mask,
21 max_new_tokens=768,
22 output_history=True,
23 return_dict_in_generate=True,
24 steps=768,
25 temperature=0.1,
26 top_p=0.95,
27 alg="entropy",
28 alg_temp=0.,
29)
30generations = [
31 tokenizer.decode(g[len(p) :].tolist())
32 for p, g in zip(input_ids, output.sequences)
33]
34
35print(generations[0].split(tokenizer.eos_token)[0])