|
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)
|
|
|
|
|
