from torch import Tensor
from transformers import AutoTokenizer, AutoModel
from ctranslate2 import Translator
from typing import Union

from fastapi import FastAPI
from pydantic import BaseModel


def average_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(
        ~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]


# text-ada replacement
embeddingTokenizer = AutoTokenizer.from_pretrained(
    './multilingual-e5-base')
embeddingModel = AutoModel.from_pretrained('./multilingual-e5-base')

# chatGpt replacement
inferenceTokenizer = AutoTokenizer.from_pretrained(
    "./flan-alpaca-gpt4-xl-ct2")
inferenceTranslator = Translator(
    "./flan-alpaca-gpt4-xl-ct2", compute_type="int8", device="cpu")


class EmbeddingRequest(BaseModel):
    input: Union[str, None] = None


class TokensCountRequest(BaseModel):
    input: Union[str, None] = None


class InferenceRequest(BaseModel):
    input: Union[str, None] = None
    max_length: Union[int, None] = 0


app = FastAPI()


@app.get("/")
async def root():
    return {"message": "Hello World"}


@app.post("/text-embedding")
async def text_embedding(request: EmbeddingRequest):
    input = request.input

    # Process the input data
    batch_dict = embeddingTokenizer([input], max_length=512,
                                    padding=True, truncation=True, return_tensors='pt')
    outputs = embeddingModel(**batch_dict)
    embeddings = average_pool(outputs.last_hidden_state,
                              batch_dict['attention_mask'])

    # create response
    return {
        'embedding': embeddings[0].tolist()
    }


@app.post('/inference')
async def inference(request: InferenceRequest):
    input_text = request.input
    max_length = 256
    try:
        max_length = int(request.max_length)
        max_length = min(1024, max_length)
    except:
        pass

    # process request
    input_tokens = inferenceTokenizer.convert_ids_to_tokens(
        inferenceTokenizer.encode(input_text))

    results = inferenceTranslator.translate_batch(
        [input_tokens], beam_size=1, max_input_length=0, max_decoding_length=max_length, num_hypotheses=1, repetition_penalty=1.3, sampling_topk=40, sampling_temperature=0.7, use_vmap=False)

    output_tokens = results[0].hypotheses[0]
    output_text = inferenceTokenizer.decode(
        inferenceTokenizer.convert_tokens_to_ids(output_tokens), skip_special_tokens=True)

    # create response
    return {
        'generated_text': output_text
    }


@app.post('/tokens-count')
async def tokens_count(request: TokensCountRequest):
    input_text = request.input

    tokens = inferenceTokenizer.convert_ids_to_tokens(
        inferenceTokenizer.encode(input_text))

    # create response
    return {
        'tokens': tokens,
        'total': len(tokens)
    }