1127 update to latest

dllm/examples/bert/chat.py

@@ -0,0 +1,71 @@
+"""
+Interactive chat / generation script for Bert models.
+
+Examples
+--------
+# Raw multi-turn generation (default)
+python -u examples/bert/chat.py --model_name_or_path "YOUR_MODEL_PATH" --chat True
+"""
+
+import sys
+from dataclasses import dataclass
+import transformers
+
+import dllm
+from dllm.pipelines import llada
+from dllm.tools.chat import multi_turn_chat, single_turn_generate
+
+
+@dataclass
+class ScriptArguments:
+    model_name_or_path: str = "dllm-collection/ModernBERT-large-chat-v0"
+    seed: int = 42
+    chat: bool = True
+    visualize: bool = True
+
+    def __post_init__(self):
+        # same base-path resolution logic as in generate.py
+        self.model_name_or_path = dllm.utils.resolve_with_base_env(
+            self.model_name_or_path, "BASE_MODELS_DIR"
+        )
+
+
+@dataclass
+class GeneratorConfig(llada.LLaDAGeneratorConfig):
+    steps: int = 128
+    max_new_tokens: int = 128
+    block_length: int = 32
+    temperature: float = 0.0
+    remasking: str = "low_confidence"
+
+
+def main():
+    parser = transformers.HfArgumentParser((ScriptArguments, GeneratorConfig))
+    script_args, gen_config = parser.parse_args_into_dataclasses()
+    transformers.set_seed(script_args.seed)
+
+    model = dllm.utils.get_model(model_args=script_args).eval()
+    tokenizer = dllm.utils.get_tokenizer(model_args=script_args)
+    generator = llada.LLaDAGenerator(model=model, tokenizer=tokenizer)
+
+    if script_args.chat:
+        multi_turn_chat(
+            generator=generator,
+            gen_config=gen_config,
+            visualize=script_args.visualize,
+        )
+    else:
+        print("\nSingle-turn generation (no chat template).")
+        single_turn_generate(
+            generator=generator,
+            gen_config=gen_config,
+            visualize=script_args.visualize,
+        )
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\nInterrupted. Bye!")
+        sys.exit(0)