Class MossForCausalLM

models_jittor/model.py:347–397

class MossForCausalLM(Module):

    def __init__(self, config):
        super(MossForCausalLM, self).__init__()
        self.config = config
        self.transformer = MossModel(config)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size)

        # Initialize weights and apply final processing
        self.apply(partial(_init_weights, config))

    def execute(
        self,
        input_ids: Optional[jt.Var] = None,
        past_key_values: Optional[Tuple[Tuple[jt.Var]]] = None,
        attention_mask: Optional[jt.Var] = None,
        token_type_ids: Optional[jt.Var] = None,
        position_ids: Optional[jt.Var] = None,
        head_mask: Optional[jt.Var] = None,
        inputs_embeds: Optional[jt.Var] = None,
        labels: Optional[jt.Var] = None,
        use_cache: Optional[bool] = None,
    ):

        hidden_states, presents = self.transformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
        )

        lm_logits = self.lm_head(hidden_states).to('float32')

        loss = None
        if labels is not None:
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

            loss = loss.to(hidden_states.dtype)

        return dict(
            loss=loss,
            logits=lm_logits,
            past_key_values=presents
        )