# Copyright 2023 PKU-Alignment Team. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from __future__ import annotations

import warnings
from typing import Any, ClassVar

import torch
import torch.nn as nn
from transformers import BloomConfig, BloomModel, BloomPreTrainedModel, PreTrainedModel
from transformers.models.bloom.modeling_bloom import (
    _CHECKPOINT_FOR_DOC,
    _CONFIG_FOR_DOC,
    BLOOM_INPUTS_DOCSTRING,
)
from transformers.utils.doc import add_code_sample_docstrings, add_start_docstrings_to_model_forward

from safe_rlhf.models.score_model import ScoreModelOutput


class BloomModelForScore(BloomPreTrainedModel):
    _keys_to_ignore_on_load_missing: ClassVar[list[str]] = [
        'h.*.self_attention.scale_mask_softmax.causal_mask',
        'lm_head.weight',
    ]

    def __init__(self, config: BloomConfig) -> None:
        super().__init__(config)
        self.transformer = BloomModel(config)

        config.score_dim = getattr(config, 'score_dim', 1)
        config.bias = getattr(config, 'bias', False)
        config.architectures = [self.__class__.__name__]
        self.score_head = nn.Linear(config.hidden_size, config.score_dim, bias=config.bias)

        # Initialize weights and apply final processing
        self.post_init()

    def get_output_embeddings(self) -> None:
        return None

    def set_decoder(self, decoder: PreTrainedModel) -> None:
        self.transformer = decoder

    def get_decoder(self) -> PreTrainedModel:
        return self.transformer

    @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING)
    @add_code_sample_docstrings(
        checkpoint=_CHECKPOINT_FOR_DOC,
        output_type=ScoreModelOutput,
        config_class=_CONFIG_FOR_DOC,
    )
    def forward(  # pylint: disable=too-many-arguments
        self,
        input_ids: torch.LongTensor | None = None,
        past_key_values: tuple[tuple[torch.Tensor, torch.Tensor], ...] | None = None,
        attention_mask: torch.Tensor | None = None,
        head_mask: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **deprecated_arguments: dict[str, Any],
    ) -> tuple[torch.Tensor, torch.Tensor] | ScoreModelOutput:
        """
        Args:

        Returns:

        Examples:

        ```python
        >>> from safe_rlhf.models.llama.modeling_llama import LlamaModelForScore
        >>> from transformers import LlamaTokenizer

        >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        # got score
        >>> outputs = model(**inputs)
        >>> scores = outputs.scores
        >>> scores
        tensor([[[0.0000]]])
        ```
        """
        assert attention_mask is not None
        if deprecated_arguments.pop('position_ids', False) is not False:
            # `position_ids` could have been `torch.Tensor` or `None`
            # so defaulting pop to `False` allows to detect if users were passing explicitly `None`
            warnings.warn(
                '`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. '
                'You can safely ignore passing `position_ids`.',
                FutureWarning,
                stacklevel=1,
            )
        if len(deprecated_arguments) > 0:
            raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        transformer_outputs = self.transformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = transformer_outputs[0]  # size = (B, L, E)
        scores = self.score_head(hidden_states)  # size = (B, L, D)

        end_scores = []
        for i in range(input_ids.size(0)):
            end_index = attention_mask[i].nonzero()[-1].item()
            end_scores.append(scores[i, end_index])  # size = (D,)
        end_scores = torch.stack(end_scores, dim=0)  # size = (B, D)

        if not return_dict:
            return scores, end_scores

        return ScoreModelOutput(
            scores=scores,  # size = (B, L, D)
            end_scores=end_scores,  # size = (B, D)
        )
