| import torch |
| import torch.nn as nn |
| from transformers import PretrainedConfig, PreTrainedModel |
| from diffusionLM.model.diffusionLM import LLaDAModel |
|
|
| class DiffusionConfig(PretrainedConfig): |
| """Configuration class for Diffusion-LLM model.""" |
| model_type = "diffusionLM" |
| |
| def __init__( |
| self, |
| vocab_size: int = 50257, |
| hidden_size: int = 768, |
| num_hidden_layers: int = 12, |
| num_attention_heads: int = 12, |
| intermediate_size: int = 3072, |
| hidden_dropout_prob: float = 0.1, |
| attention_probs_dropout_prob: float = 0.1, |
| max_position_embeddings: int = 1024, |
| initializer_range: float = 0.02, |
| layer_norm_eps: float = 1e-12, |
| pad_token_id: int = 0, |
| mask_token_id: int = 50256, |
| eos_token_id: int = 50256, |
| num_timesteps: int = 100, |
| time_embed_dim: int = 128, |
| **kwargs |
| ): |
| super().__init__(pad_token_id=pad_token_id, **kwargs) |
| self.vocab_size = vocab_size |
| self.hidden_size = hidden_size |
| self.num_hidden_layers = num_hidden_layers |
| self.num_attention_heads = num_attention_heads |
| self.intermediate_size = intermediate_size |
| self.hidden_dropout_prob = hidden_dropout_prob |
| self.attention_probs_dropout_prob = attention_probs_dropout_prob |
| self.max_position_embeddings = max_position_embeddings |
| self.initializer_range = initializer_range |
| self.layer_norm_eps = layer_norm_eps |
| self.mask_token_id = mask_token_id |
| self.eos_token_id = eos_token_id |
| self.num_timesteps = num_timesteps |
| self.time_embed_dim = time_embed_dim |
|
|
| class DiffusionLLM(PreTrainedModel): |
| """Main Diffusion-LLM model class""" |
| config_class = DiffusionConfig |
| base_model_prefix = "diffusionLM" |
|
|
| def __init__(self, config: DiffusionConfig): |
| super().__init__(config) |
| self.model = LLaDAModel(config) |
| self.init_weights() |
|
|
| def forward( |
| self, |
| input_ids=None, |
| attention_mask=None, |
| timesteps=None, |
| labels=None, |
| return_dict=True, |
| ): |
| outputs = self.model( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| timesteps=timesteps, |
| labels=labels, |
| ) |
| |
| return outputs |
|
|
| def generate( |
| self, |
| prompt=None, |
| max_length=100, |
| num_inference_steps=50, |
| temperature=1.0, |
| strategy='random', |
| top_p=0.9, |
| top_k=50, |
| num_beams=5, |
| return_scores=False, |
| use_streaming=False, |
| callback_fn=None |
| ): |
| """Unified generation interface""" |
| if use_streaming: |
| return self.generate_stream( |
| prompt=prompt, |
| max_length=max_length, |
| num_inference_steps=num_inference_steps, |
| temperature=temperature, |
| strategy=strategy, |
| top_p=top_p, |
| top_k=top_k, |
| num_beams=num_beams, |
| callback_fn=callback_fn |
| ) |
| else: |
| return self.model.generate( |
| prompt=prompt, |
| max_length=max_length, |
| num_inference_steps=num_inference_steps, |
| temperature=temperature, |
| strategy=strategy, |
| top_p=top_p, |
| top_k=top_k, |
| num_beams=num_beams, |
| return_scores=return_scores |
| ) |
|
|
| def generate_stream(self, **kwargs): |
| """Streaming generation wrapper""" |
| return self.model.generate_stream(**kwargs) |
|
|
| def prepare_inputs_for_generation(self, input_ids, **kwargs): |
| """Prepare inputs for generation compatibility""" |
| return { |
| "input_ids": input_ids, |
| "attention_mask": kwargs.get("attention_mask", None), |
| "timesteps": kwargs.get("timesteps", None), |
| } |
|
|
| @staticmethod |
| def _reorder_cache(past, beam_idx): |
| """Reorder cache for beam search compatibility""" |
| return past |
|
|
