File size: 3,393 Bytes
da49047 | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import os
import torch
import torch.nn.functional as F
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_bet import BETConfig
from .bet_model import BETConfig as CoreConfig, SparkBET, uniform_steps
class BETPreTrainedModel(PreTrainedModel):
config_class=BETConfig
base_model_prefix="core"
supports_gradient_checkpointing=False
_no_split_modules=["PlainBlock","LoopedBlock"]
class BETForCausalLM(BETPreTrainedModel,GenerationMixin):
def __init__(self,config):
super().__init__(config)
core_cfg=CoreConfig(
vocab_size=config.vocab_size,
hidden_size=config.hidden_size,
intermediate_size=config.intermediate_size,
prelude_layers=config.prelude_layers,
body_blocks=config.body_blocks,
coda_layers=config.coda_layers,
num_heads=config.num_attention_heads,
num_kv_heads=config.num_key_value_heads,
head_dim=config.head_dim,
lora_rank=config.lora_rank,
hyper_lanes=config.hyper_lanes,
max_seq_len=config.max_position_embeddings,
max_loops=config.max_loops,
rope_theta=config.rope_theta,
rms_eps=config.rms_norm_eps,
ddl_beta_init=config.ddl_beta_init,
ddl_k_eps=config.ddl_k_eps,
ddl_v_sigmoid_scale=config.ddl_v_sigmoid_scale,
)
self.core=SparkBET(core_cfg)
def get_input_embeddings(self):return self.core.embed
def set_input_embeddings(self,value):self.core.embed=value
def get_output_embeddings(self):return None
def set_output_embeddings(self,value):
if value is not None:raise ValueError("SparkBET uses tied input/output embeddings")
def _cycles(self,cycles=None):
if cycles is None:
cycles=int(os.environ.get("BET_EVAL_CYCLES",self.config.refinement_cycles))
cycles=int(cycles)
if not 1<=cycles<=self.config.max_loops:
raise ValueError(f"refinement cycles must be in [1,{self.config.max_loops}]")
return cycles
def forward(
self,input_ids=None,attention_mask=None,labels=None,cycles=None,
past_key_values=None,use_cache=None,return_dict=True,**kwargs,
):
if input_ids is None:raise ValueError("input_ids is required")
if past_key_values is not None:raise ValueError("SparkBET does not implement a KV cache")
logits=self.core(input_ids,uniform_steps(self._cycles(cycles)),attention_mask=attention_mask)
loss=None
if labels is not None:
shift_logits=logits[:,:-1].contiguous().float();shift_labels=labels[:,1:].contiguous()
loss=F.cross_entropy(shift_logits.view(-1,shift_logits.size(-1)),shift_labels.view(-1),ignore_index=-100)
if not return_dict:return tuple(v for v in (loss,logits) if v is not None)
return CausalLMOutput(loss=loss,logits=logits)
def prepare_inputs_for_generation(self,input_ids,attention_mask=None,**kwargs):
max_len=self.config.max_position_embeddings
if input_ids.shape[1]>max_len:
input_ids=input_ids[:,-max_len:]
if attention_mask is not None:attention_mask=attention_mask[:,-max_len:]
return {"input_ids":input_ids,"attention_mask":attention_mask,"use_cache":False}
|