|
|
|
|
|
|
|
|
from typing import Optional, Tuple |
|
|
|
|
|
import torch |
|
|
from physicsnemo.models.transolver import Transolver as TransolverBase |
|
|
from transformers import PretrainedConfig, PreTrainedModel |
|
|
|
|
|
|
|
|
class TransolverConfig(PretrainedConfig): |
|
|
model_type = "transolver" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
functional_dim: int = 5, |
|
|
out_dim: int = 1, |
|
|
embedding_dim: Optional[int] = 3, |
|
|
n_layers: int = 4, |
|
|
n_hidden: int = 128, |
|
|
dropout: float = 0.0, |
|
|
n_head: int = 8, |
|
|
act: str = "gelu", |
|
|
mlp_ratio: int = 4, |
|
|
slice_num: int = 32, |
|
|
unified_pos: bool = False, |
|
|
ref: int = 8, |
|
|
structured_shape: Optional[Tuple[int, ...]] = None, |
|
|
use_te: bool = False, |
|
|
time_input: bool = False, |
|
|
**kwargs, |
|
|
): |
|
|
super().__init__(**kwargs) |
|
|
self.functional_dim = functional_dim |
|
|
self.out_dim = out_dim |
|
|
self.embedding_dim = embedding_dim |
|
|
self.n_layers = n_layers |
|
|
self.n_hidden = n_hidden |
|
|
self.dropout = dropout |
|
|
self.n_head = n_head |
|
|
self.act = act |
|
|
self.mlp_ratio = mlp_ratio |
|
|
self.slice_num = slice_num |
|
|
self.unified_pos = unified_pos |
|
|
self.ref = ref |
|
|
self.structured_shape = structured_shape |
|
|
self.use_te = use_te |
|
|
self.time_input = time_input |
|
|
|
|
|
|
|
|
class TransolverModel(PreTrainedModel): |
|
|
config_class = TransolverConfig |
|
|
|
|
|
def __init__(self, config: TransolverConfig): |
|
|
super().__init__(config) |
|
|
|
|
|
self.transolver = TransolverBase( |
|
|
functional_dim=config.functional_dim, |
|
|
out_dim=config.out_dim, |
|
|
embedding_dim=config.embedding_dim, |
|
|
n_layers=config.n_layers, |
|
|
n_hidden=config.n_hidden, |
|
|
dropout=config.dropout, |
|
|
n_head=config.n_head, |
|
|
act=config.act, |
|
|
mlp_ratio=config.mlp_ratio, |
|
|
slice_num=config.slice_num, |
|
|
unified_pos=config.unified_pos, |
|
|
ref=config.ref, |
|
|
structured_shape=config.structured_shape, |
|
|
use_te=config.use_te, |
|
|
time_input=config.time_input, |
|
|
) |
|
|
|
|
|
|
|
|
self.post_init() |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
fx: torch.Tensor, |
|
|
embedding: Optional[torch.Tensor] = None, |
|
|
time: Optional[torch.Tensor] = None, |
|
|
**kwargs, |
|
|
): |
|
|
"""Thin wrapper around TransolverBase.forward. |
|
|
|
|
|
Args: |
|
|
fx: [B, N, functional_dim] or [B, *structure, functional_dim] |
|
|
embedding: position / embeddings |
|
|
time: optional time tensor |
|
|
""" |
|
|
return self.transolver(fx, embedding=embedding, time=time) |
|
|
|
