| """Hugging Face model wrapper for HydrAMP.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| from transformers import PreTrainedModel |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
| from transformers.utils import ModelOutput |
|
|
| from .config import HydrAMPConfig |
| from .hydramp import HydrAMPDecoder, HydrAMPEncoder |
|
|
|
|
| @dataclass |
| class HydrAMPOutput(ModelOutput): |
| """HydrAMP forward outputs.""" |
|
|
| logits: torch.Tensor | None = None |
| mean: torch.Tensor | None = None |
| log_std: torch.Tensor | None = None |
|
|
|
|
| class HydrAMPModel(PreTrainedModel): |
| """HydrAMP model with HF `AutoModel` compatibility.""" |
|
|
| config_class = HydrAMPConfig |
| base_model_prefix = "hydramp" |
|
|
| def __init__(self, config: HydrAMPConfig) -> None: |
| super().__init__(config) |
| if len(config.default_condition) != config.condition_dim: |
| raise ValueError( |
| f"default_condition must contain {config.condition_dim} values, got {len(config.default_condition)}." |
| ) |
|
|
| self.encoder = HydrAMPEncoder( |
| vocab_size=config.vocab_size, |
| embedding_dim=config.embedding_dim, |
| latent_dim=config.latent_dim, |
| sequence_length=config.sequence_length, |
| gru_hidden_size=config.encoder_hidden_size, |
| ) |
| self.decoder = HydrAMPDecoder( |
| sequence_length=config.sequence_length, |
| latent_dim=config.latent_dim, |
| condition_dim=config.condition_dim, |
| hidden_size=config.decoder_hidden_size, |
| vocab_size=config.vocab_size, |
| ) |
| self.register_buffer( |
| "default_condition", |
| torch.tensor(config.default_condition, dtype=torch.float32), |
| persistent=False, |
| ) |
| self.post_init() |
|
|
| def forward_latent_positions( |
| self, |
| z: torch.Tensor, |
| num_steps: int | None = None, |
| condition: torch.Tensor | None = None, |
| *, |
| return_logits: bool = True, |
| ) -> CausalLMOutputWithPast: |
| """Decode latent vectors to sequence distributions (GRUVAE-style API). |
| |
| Output length is fixed to ``config.sequence_length``. If ``num_steps`` is |
| passed, it must equal that value. |
| """ |
| fixed = self.config.sequence_length |
| if num_steps is None: |
| num_steps = fixed |
| elif num_steps != fixed: |
| msg = f"HydrAMP decoder length is fixed at {fixed}; got num_steps={num_steps}." |
| raise ValueError(msg) |
|
|
| if condition is None: |
| condition = self.default_condition.unsqueeze(0).expand(z.shape[0], -1) |
| condition = condition.to(device=z.device, dtype=z.dtype) |
| decoder_input = torch.cat([z, condition], dim=-1) |
| out = self.decoder( |
| decoder_input, |
| return_logits=return_logits, |
| gumbel_temperature=self.config.temperature, |
| ) |
| return CausalLMOutputWithPast(logits=out, past_key_values=None) |
|
|
| def decode_to_token_ids( |
| self, |
| z: torch.Tensor, |
| num_steps: int | None = None, |
| condition: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """Greedy token IDs from latent ``z`` (argmax over vocabulary per position).""" |
| logits = self.forward_latent_positions( |
| z, num_steps=num_steps, condition=condition, return_logits=True |
| ).logits |
| assert logits is not None |
| return logits.argmax(dim=-1) |
|
|
| def forward(self, input_ids: torch.Tensor, **_: object) -> HydrAMPOutput: |
| """Run encode + deterministic decode for reconstruction.""" |
| mean, log_std = self.encoder.encode(input_ids) |
| logits = self.forward_latent_positions(mean, return_logits=True).logits |
| return HydrAMPOutput(logits=logits, mean=mean, log_std=log_std) |
|
|
