|
from fastapi import FastAPI, Request |
|
from pydantic import BaseModel |
|
from sentence_transformers import SentenceTransformer |
|
import numpy as np |
|
|
|
app = FastAPI() |
|
|
|
|
|
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 |
|
|
|
|
|
embeddings = model.encode(input_text) |
|
|
|
|
|
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 |
|
|