File size: 704 Bytes
493108b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import os
os.environ['TRANSFORMERS_CACHE'] = '/app/cache'
app = FastAPI(
title="Text Embedding API",
description="Dùng mô hình Vietnamese_Embedding từ AITeamVN",
version="1.0"
)
# Load mô hình
model = SentenceTransformer("AITeamVN/Vietnamese_Embedding")
model.max_seq_length = 2048
# Định nghĩa schema input
class TextInput(BaseModel):
texts: list[str]
@app.post("/embed")
def embed_text(data: TextInput):
inputs = [t.strip() for t in data.texts]
embs = model.encode(inputs, convert_to_numpy=True)
return {"embeddings": [e.tolist() for e in embs]}
|