HF_compatibility_testv3-ForCausalLM / modeling_lightningtransformer.py
Jumpr's picture
Upload model
c7d303a verified
Raw
History Blame Contribute Delete
605 Bytes
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from .configuration_lightningtransformer import LightningTransformerModelConfig
from .lightningtransformer import LightningTransformer
class LightningTransformerModelForCausalLM(PreTrainedModel, GenerationMixin):
config_class = LightningTransformerModelConfig
def __init__(self, config):
super().__init__(config)
self.model = LightningTransformer(**config.cfg)
self.use_cache = False
self.post_init()
def forward(self, input_ids, **kwargs):
return self.model.forward(input_ids)