from fastapi import FastAPI, Request
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import numpy as np

app = FastAPI()

# Load the Sentence Transformer model
model = SentenceTransformer("BAAI/bge-large-zh-v1.5")

class EmbeddingRequest(BaseModel):
    input: str

@app.post("/embeddings")
async def embeddings(request: EmbeddingRequest):
    input_text = request.input

    # Calculate embeddings
    embeddings = model.encode(input_text)

    # Format the embeddings in OpenAI compatible format
    data = {
        "object": "list",
        "data": [
            {
                "object": "embedding",
                "embedding": embeddings.tolist(),
                "index": 0
            }
        ],
        "model": "BAAI/bge-large-zh-v1.5",
        "usage": {
            "prompt_tokens": len(input_text),
            "total_tokens": len(input_text)
        }
    }

    return data