import logging

import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel, Field
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("deepvk/USER-bge-m3")

app = FastAPI(
    title="🤗 HuggingFace Embeddings API",
    description="API for inference `deepvk` HuggingFace embeddings model",
    version="0.1.0"
)


class EmbeddingRequest(BaseModel):
    texts: list[str] = Field(..., max_length=3000)


@app.post("/embeddings")
def get_embeddings(request: EmbeddingRequest) -> dict[str, list[list[float]]]:
    embeddings = model.encode_document(request.texts, normalize_embeddings=False)
    return {"embeddings": embeddings.tolist()}


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    uvicorn.run(app, host="0.0.0.0", port=7860)