from fastapi import FastAPI
from pydantic import BaseModel
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModel
from bs4 import BeautifulSoup
import hazm
import time

model_name="HooshvareLab/bert-base-parsbert-uncased"
tokenizer=AutoTokenizer.from_pretrained(model_name)
model=TFAutoModel.from_pretrained(model_name)

sent_tokenizer=hazm.SentenceTokenizer()
normalizer=hazm.Normalizer()


app=FastAPI()

class Input(BaseModel):
    texts: list

@app.post("/get_vectors")
def get_vecs(data: Input):
    now=time.time()
    texts=data.texts
    texts=list(map(lambda x: BeautifulSoup(x).get_text(), texts))
    texts=list(map(normalizer.normalize, texts))
    
    tokens=tokenizer(texts, return_tensors="tf", padding=True, truncation=True, max_length=512)
    outputs=model(**tokens)
    
    sentence_embedding=tf.reduce_mean(outputs.last_hidden_state, axis=1)
    vecs=sentence_embedding.numpy().tolist()
    
    return {"vectors": vecs, "duration": time.time()-now}