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