#
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
#

"""
Utils specific to GPT2 network.
"""

# torch
import torch


# from HuggingFace transformers
from transformers.generation_logits_process import (
    MinLengthLogitsProcessor,
    LogitsProcessorList,
    ForcedEOSTokenLogitsProcessor,
)
from transformers.generation_stopping_criteria import (
    MaxLengthCriteria,
    StoppingCriteriaList,
)

# TRT-HuggingFace
from NNDF.general_utils import measure_python_inference_code
from NNDF.torch_utils import use_cuda
from NNDF.tensorrt_utils import TRTNativeRunner

@use_cuda
def gpt2_inference(gpt2, input_ids, timing_profile, use_cuda=True, use_cache=False, past_key_values = None):
    gpt2_stmt = lambda: gpt2(input_ids=input_ids, use_cache=use_cache, past_key_values=past_key_values)
    gpt2_e2e_time = measure_python_inference_code(gpt2_stmt, timing_profile)
    return (gpt2_stmt(), gpt2_e2e_time)


# Code specifically for Pythonic inference measurement used across all GPT2 related scripts
@use_cuda
def full_inference(
    gpt2,
    input_ids,
    tokenizer,
    timing_profile,
    max_length,
    min_length = 0,
    use_cuda=True,
    batch_size=1,
    early_stopping=False,
    use_cache=False,
    num_beams = 1,
):

    if isinstance(gpt2, TRTNativeRunner):
        gpt2.set_return_device("cuda" if use_cuda else "cpu")

    def _e2e():
        with torch.no_grad():
            output = gpt2.generate(
                input_ids,
                max_length=max_length,
                min_length=min_length,
                batch_size=batch_size,
                num_beams=num_beams,
                use_cache=use_cache,
                early_stopping=early_stopping
            )

        return output

    full_e2e_time = measure_python_inference_code(_e2e, timing_profile)
    return (_e2e(), full_e2e_time)


@use_cuda
def calculate_perplexity(gpt2, input_ids, max_seq_len=None, use_cuda=True):
    if isinstance(gpt2, TRTNativeRunner):
        gpt2.set_return_device("cuda" if use_cuda else "cpu")

    with torch.no_grad():
        if max_seq_len is not None:
            input_ids = input_ids[:, :max_seq_len]
        logits = gpt2(input_ids).logits
        # Shift logits and target ids so that probabilities generated by token < n line up with output token n.
        shifted_logits = logits[:, :-1, :]
        target_ids = input_ids[:, 1:]
        loss = torch.nn.CrossEntropyLoss()(shifted_logits.permute((0, 2, 1)), target_ids)
        return torch.exp(loss).item()
