2024-05-18 07:55:17 +00:00
|
|
|
from fastembed import TextEmbedding
|
|
|
|
from fastapi import FastAPI
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
2024-05-19 08:33:49 +00:00
|
|
|
models = {}
|
2024-05-18 07:55:17 +00:00
|
|
|
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
|
|
class EmbeddingRequest(BaseModel):
|
|
|
|
model: str
|
2024-05-19 08:33:49 +00:00
|
|
|
input: str
|
2024-05-18 07:55:17 +00:00
|
|
|
|
2024-05-19 08:33:49 +00:00
|
|
|
@app.post("/v1/embeddings")
|
2024-05-18 07:55:17 +00:00
|
|
|
def embeddings(request: EmbeddingRequest):
|
2024-05-19 08:33:49 +00:00
|
|
|
model = models.get(request.model) or TextEmbedding(request.model)
|
|
|
|
models[request.model] = model
|
|
|
|
embeddings = next(model.embed(request.input)).tolist()
|
|
|
|
return {"data": [{"embedding": embeddings}]}
|
2024-05-18 07:55:17 +00:00
|
|
|
|
2024-05-27 10:15:04 +00:00
|
|
|
@app.get("/health")
|
|
|
|
def health():
|
|
|
|
return {"status": "ok"}
|
|
|
|
|
2024-05-18 07:55:17 +00:00
|
|
|
if __name__ == "__main__":
|
|
|
|
import uvicorn
|
|
|
|
|
|
|
|
uvicorn.run(app, host="0.0.0.0", port=11345)
|