Views
No views yet
1git clone https://github.com/ZHZisZZ/dllm.git
2cd dllm
3pip install -e .1from dataclasses import dataclass
2import transformers
3import dllm
4from dllm.tools.chat import decode_trim
5from dllm.pipelines import llada
6
7''' #or log in using `huggingface-cli login`
8token= 'hf_...'
9from huggingface_hub import login
10login(token=token)
11'''
12
13# ---------------------------------------------------------
14# Load model + tokenizer
15# ---------------------------------------------------------
16
17@dataclass
18class ScriptArguments:
19 model_name_or_path: str = "lamm-mit/LLaDA-8B-Bioinspired-dLLM-Instruct-11-21-2025"
20
21 def __post_init__(self):
22 self.model_name_or_path = dllm.utils.resolve_with_base_env(
23 self.model_name_or_path, "BASE_MODELS_DIR"
24 )
25
26script_args = ScriptArguments()
27
28transformers.set_seed(42)
29
30model = dllm.utils.get_model(model_args=script_args).eval()
31tokenizer = dllm.utils.get_tokenizer(model_args=script_args)
32
33generator = llada.LLaDAGenerator(
34 model=model,
35 tokenizer=tokenizer,
36)
37
38gen_config = llada.LLaDAGeneratorConfig(
39 steps=256,
40 max_new_tokens=256,
41 block_length=32,
42 temperature=0.0,
43 remasking="low_confidence",
44)
45
46# ---------------------------------------------------------
47# Batched inference step
48# ---------------------------------------------------------
49
50messages_batch = [
51 [{"role": "user", "content": "Explain materiomics briefly."}],
52 [{"role": "user", "content": "Define mechanobiology in one paragraph."}],
53 [{"role": "user", "content": "Why is silk stronger than elastin?"}],
54]
55
56inputs = tokenizer.apply_chat_template(
57 messages_batch,
58 add_generation_prompt=True,
59 tokenize=True,
60)
61
62outputs = generator.generate(
63 inputs,
64 gen_config,
65 return_dict_in_generate=True,
66)
67
68sequences = decode_trim(tokenizer, outputs.sequences.tolist(), inputs)
69
70# ---------------------------------------------------------
71# Results
72# ---------------------------------------------------------
73
74for i, s in enumerate(sequences):
75 print("\n" + "-" * 70)
76 print(f"[Sample {i}]")
77 print("-" * 70)
78 print(s.strip())
791terminal_visualizer = dllm.core.generation.visualizer.TerminalVisualizer(
2 tokenizer=tokenizer
3 )
4terminal_visualizer.visualize(outputs.histories, rich=True)
1gen_config = llada.LLaDAGeneratorConfig(
2 steps=512,
3 max_new_tokens=512,
4 block_length=32,
5 temperature=0.2,
6 remasking="low_confidence",
7)
8masked_messages = [
9 [
10 {
11 "role": "user",
12 "content": (
13 "In spider-silk materiomics, we often optimize hierarchical structure "
14 "from amino-acid sequence to β-sheet nanocrystal arrangement. "
15 "Complete the missing reasoning steps for the following design question:\n\n"
16 f"**Design Problem:** How could one tune the fraction of β-sheet "
17 f"nanocrystals to increase toughness without compromising elasticity?\n\n"
18 f"Missing reasoning: {tokenizer.mask_token * 128}"
19 ),
20 },
21 {
22 "role": "assistant",
23 "content": (
24 f"The summary is: {tokenizer.mask_token * 20}" #
25 ),
26 },
27 ],
28
29 [
30 {
31 "role": "user",
32 "content": (
33 "In nacre-inspired composite design, we often tune the architecture of "
34 "brick-and-mortar layers to balance stiffness, strength, and toughness. "
35 "Complete the missing reasoning steps for the following design question:\n\n"
36 "**Design Problem:** How could one introduce controlled mineral platelet "
37 "misalignment to enhance toughness while preserving high stiffness?\n\n"
38 f"Missing reasoning: {tokenizer.mask_token * 128}"
39 ),
40 },
41 {
42 "role": "assistant",
43 "content": (
44 f"The design principle is: {tokenizer.mask_token * 20}"
45 ),
46 },
47]
48
49]
50
51# Tokenize input with NO generation prompt
52inputs = tokenizer.apply_chat_template(
53 masked_messages,
54 add_generation_prompt=False,
55 tokenize=True,
56)
57
58# Infilling
59outputs = generator.infill(inputs, gen_config, return_dict_in_generate=True)
60sequences = decode_trim(tokenizer, outputs.sequences.tolist(), inputs)
61
62# Print results
63for idx, (inp, filled) in enumerate(zip(inputs, sequences)):
64 print("\n" + "-" * 80)
65 print(f"[Case {idx}]")
66 print("-" * 80)
67 print("[Masked]:\n" + tokenizer.decode(inp))
68 print("\n[Filled]:\n" + (filled.strip() if filled.strip() else "<empty>"))
69
70print("\n" + "=" * 80 + "\n")
71
72terminal_visualizer.visualize(outputs.histories, rich=True)