class BertLMHeadModel(BertPreTrainedModel):

    _keys_to_ignore_on_load_unexpected = [r"pooler"]
    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]

    def __init__(self, config):
        super().__init__(config)

        self.bert = BertModel(config, add_pooling_layer=False)
        self.cls = BertOnlyMLMHead(config)

        self.init_weights()

    def get_output_embeddings(self):
        return self.cls.predictions.decoder

    def set_output_embeddings(self, new_embeddings):
        self.cls.predictions.decoder = new_embeddings

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        position_ids=None,
        head_mask=None,
        query_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        labels=None,
        past_key_values=None,
        use_cache=True,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        return_logits=False,
        is_decoder=True,
        reduction="mean",
    ):
       
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        if labels is not None:
            use_cache = False
        if past_key_values is not None:
            query_embeds = None

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            head_mask=head_mask,
            query_embeds=query_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            is_decoder=is_decoder,
        )

        sequence_output = outputs[0]
        if query_embeds is not None:
            sequence_output = outputs[0][:, query_embeds.shape[1] :, :]

        prediction_scores = self.cls(sequence_output)

        if return_logits:
            return prediction_scores[:, :-1, :].contiguous()

        lm_loss = None
        if labels is not None:
            # we are doing next-token prediction; shift prediction scores and input ids by one
						# raw output logits 을 한칸 씩 왼쪽으로 옮기기 위해 [:, :-1, :]  for next-token prediction
            shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
						# label token도 한칸 옮겨서 첫번째 토큰 지움 왜냐면 첫번째 토큰 앞에 오는 토큰은 없으니까
            labels = labels[:, 1:].contiguous()
						# label smoothing is regularization technique
            loss_fct = CrossEntropyLoss(reduction=reduction, label_smoothing=0.1)
            lm_loss = loss_fct(
                shifted_prediction_scores.view(-1, self.config.vocab_size),
                labels.view(-1),
            )
            if reduction == "none":
                lm_loss = lm_loss.view(prediction_scores.size(0), -1).sum(1)

        if not return_dict:
            output = (prediction_scores,) + outputs[2:]
            return ((lm_loss,) + output) if lm_loss is not None else output
				
				**# 중요!!!**
				**# 얘네가 CasualLMOutputWithCrossAttentions 이라는 함수로 처리가 되는게 아니라
				#	data class output으로 출력이 됨!
				# 어쨌든 이 BertLMHeadModel은 Q-former 가 되는 base model로 이미지 엠베딩을
				# 쿼리 토큰으로 가져와서 q-former 넣은 후 쿼리 output을 만든 후 vision_proj 함수에 넣어서
				# image feature를 만든다!
				# 이 image feature를 가지고 Image-text matching, captioning, contrastive learning 수행!**

        return CausalLMOutputWithCrossAttentions(
            loss=lm_loss, #Language Model Loss for next-token prediction
            logits=prediction_scores,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            cross_attentions=outputs.cross_attentions,
        )

    def prepare_inputs_for_generation(
        self, input_ids, query_embeds, past=None, attention_mask=None, **model_kwargs
    ):
        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
        if attention_mask is None:
            attention_mask = input_ids.new_ones(input_ids.shape)
        query_mask = input_ids.new_ones(query_embeds.shape[:-1])
        attention_mask = torch.cat([query_mask, attention_mask], dim=-1)

        # cut decoder_input_ids if past is used
        if past is not None:
            input_ids = input_ids[:, -1:]

        return {
            "input_ids": input_ids,
            "query_embeds": query_embeds,
            "attention_mask": attention_mask,
            "past_key_values": past,
            "encoder_hidden_states": model_kwargs.get("encoder_hidden_states", None),
            "encoder_attention_mask": model_kwargs.get("encoder_attention_mask", None),
            "is_decoder": True,
        }

    def _reorder_cache(self, past, beam_idx):
        reordered_past = ()
        for layer_past in past:
            reordered_past += (
                tuple(
                    past_state.index_select(0, beam_idx) for past_state in layer_past
                ),
            )
        return reordered_past

Loss 계산 방식 설명

• Input Tokens: [START], The, quick, brown, fox, jumps, over, the, lazy
• Labels: The, quick, brown, fox, jumps, over, the, lazy, dog

Loss calculation between predictions vs labels (여기선 jump, jumps 차이)

• Model Predictions: The, quick, brown, fox, jump, over, the, lazy, dog
• Correct Labels: The, quick, brown, fox, jumps, over, the, lazy, dog

CasualLMOutputWithCrossAttentions 는 transformer 내장 data class! 연산하는 함수가 아님!

self.bert 값이 여기서!