main.py

from torch import Tensor
from transformers import AutoTokenizer, AutoModel
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')


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


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()
    }