from fastapi import FastAPI
from sentence_transformers import SentenceTransformer

app = FastAPI()
model = None


def get_model():
    global model
    if model is None:
        model = SentenceTransformer(
            "nvidia/llama-3.2-nv-embedqa-1b-v2", trust_remote_code=True
        )
    return model


@app.post("/v1/embeddings")
async def embed(request: dict):
    texts = request.get("input", [])
    if isinstance(texts, str):
        texts = [texts]
    m = get_model()
    embeddings = m.encode(texts).tolist()
    return {
        "data": [{"embedding": e, "index": i} for i, e in enumerate(embeddings)],
        "model": request.get("model", ""),
    }