| import tensorflow as tf | |
| from transformers.modeling_tf_utils import TFPreTrainedModel | |
| from .configuration_my_model import MyModelConfig | |
| class TFMyModelPretrainedModel(TFPreTrainedModel): | |
| config_class = MyModelConfig | |
| class TFMyModel(TFMyModelPretrainedModel): | |
| def __init__(self, config: MyModelConfig): | |
| super().__init__(config) | |
| self.config = config | |
| self.n_layers = config.n_layers | |
| self.hidden_dim = config.hidden_dim | |
| self.linear = tf.keras.layers.Dense(units=config.n_layers) | |