pri2k
Remove login(), pass token to SentenceTransformer
ec7f722
raw
history blame
574 Bytes
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import os
app = FastAPI()
hf_token = os.getenv("HF_TOKEN")
if not hf_token:
raise ValueError("HF_TOKEN not set in environment variables.")
# Pass token while loading the model
model = SentenceTransformer("mental/mental-bert-base-uncased", use_auth_token=hf_token)
class InputText(BaseModel):
text: str
@app.post("/embed")
async def embed_text(input: InputText):
embedding = model.encode(input.text).tolist()
return {"embedding": embedding}