File size: 1,509 Bytes
5038157 | 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 41 42 43 | import torch
from torch.nn import functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput
from .configuration_kairo import KairoGPTConfig
from .kairo_model import KairoGPT, KairoGPTConfig as _InternalCfg
class KairoGPTForCausalLM(PreTrainedModel):
config_class = KairoGPTConfig
def __init__(self, config: KairoGPTConfig):
super().__init__(config)
internal_cfg = _InternalCfg(
vocab_size=config.vocab_size,
block_size=config.block_size,
n_layer=config.n_layer,
n_head=config.n_head,
n_embd=config.n_embd,
dropout=config.dropout,
)
self.transformer = KairoGPT(internal_cfg)
def forward(self, input_ids, labels=None, **kwargs):
logits, _ = self.transformer(input_ids)
loss = None
if labels is not None:
loss = F.cross_entropy(
logits[:, :-1, :].reshape(-1, logits.size(-1)),
labels[:, 1:].reshape(-1),
)
return CausalLMOutput(loss=loss, logits=logits)
def prepare_inputs_for_generation(self, input_ids, **kwargs):
return {"input_ids": input_ids[:, -self.config.block_size:]}
@torch.no_grad()
def generate(self, input_ids, max_new_tokens=400, temperature=0.8, top_k=40, **kwargs):
return self.transformer.generate(
input_ids, max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k
)
|