640510702phithak's picture
Update app.py
edfce6b verified
from fastapi import FastAPI
from transformers import AutoTokenizer, AutoModel
import torch
# เปลี่ยน cache directory เป็น /tmp
MODEL_NAME = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir="/tmp")
model = AutoModel.from_pretrained(MODEL_NAME, cache_dir="/tmp")
app = FastAPI()
def get_embedding(text):
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
with torch.no_grad():
outputs = model(**inputs)
embedding = outputs.last_hidden_state.mean(dim=1)
return embedding.squeeze().tolist()
@app.get("/")
def home():
return "หวัดดี คุณกำลังเข้าสู่บริการรับฝากหัวใจ API ใช้ post .../embed"
@app.post("/embed")
async def embed_text(data: dict):
text = data.get("input", "")
if not text:
return {"error": "No text provided"}
vector = get_embedding(text)
return {"input": text, "embedding": vector}