File size: 1,017 Bytes
c6bc350 b5b1e4e c6bc350 7ceff4a c6bc350 860bfc4 b5b1e4e 860bfc4 753f724 860bfc4 b5b1e4e 860bfc4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import tensorflow as tf
from transformers.modeling_tf_utils import unpack_inputs
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)
@property
def dummy_inputs(self):
hidden = tf.zeros(shape=(1, self.config.hidden_dim))
dummy_inputs = {"hidden": hidden}
return dummy_inputs
@unpack_inputs
def call(
self,
hidden,
output_attentions=False,
output_hidden_states=False,
return_dict=False,
):
breakpoint()
self.linear(hidden)
|