# Copyright (C) 2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import os
from typing import List

import lm_eval.api.registry
import torch
from docarray import BaseDoc
from evals.evaluation.lm_evaluation_harness.lm_eval.models.huggingface import HFLM, GaudiHFModelAdapter

from comps import CustomLogger, ServiceType, opea_microservices, opea_telemetry, register_microservice

logger = CustomLogger("self_hosted_hf")
logflag = os.getenv("LOGFLAG", False)

lm_eval.api.registry.MODEL_REGISTRY["hf"] = HFLM
lm_eval.api.registry.MODEL_REGISTRY["gaudi-hf"] = GaudiHFModelAdapter


class LLMCompletionDoc(BaseDoc):
    batched_inputs: List
    logprobs: int = 10
    max_new_tokens: int = 0
    temperature: float = 0.0


model = os.getenv("MODEL", "")
model_args = os.getenv("MODEL_ARGS", "")
device = os.getenv("DEVICE", "")

llm = lm_eval.api.registry.get_model(model).create_from_arg_string(
    model_args,
    {
        "batch_size": 1,  # dummy
        "max_batch_size": None,
        "device": device,
    },
)


@register_microservice(
    name="opea_service@self_hosted_hf",
    service_type=ServiceType.LLM,
    endpoint="/v1/completions",
    host="0.0.0.0",
    port=9006,
)
@opea_telemetry
def llm_generate(input: LLMCompletionDoc):
    if logflag:
        logger.info(input)
    global llm
    batched_inputs = torch.tensor(input.batched_inputs, dtype=torch.long, device=llm.device)
    with torch.no_grad():
        # TODO, use model.generate.
        logits = llm._model_call(batched_inputs)

    logits = torch.nn.functional.log_softmax(logits, dim=-1)
    # Check if per-token argmax is exactly equal to continuation
    greedy_tokens = logits.argmax(dim=-1)
    logprobs = torch.gather(logits, 2, batched_inputs[:, 1:].unsqueeze(-1)).squeeze(-1)
    result = {
        "greedy_tokens": greedy_tokens.detach().cpu().tolist(),
        "logprobs": logprobs.detach().cpu().tolist(),
        "batched_inputs": input.batched_inputs,
    }
    if logflag:
        logger.info(result)
    return result


if __name__ == "__main__":
    opea_microservices["opea_service@self_hosted_hf"].start()
