moon
/
spc-pleroma
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

B FastembedServer: Switch to OpenAI api, support changing models

This commit is contained in:
Lain Soykaf 2024-05-19 12:33:49 +04:00
parent 72ec261a69
commit b9af017a4c
1 changed files with 7 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
python/fastembed-server.py
View File

@ -2,18 +2,20 @@ from fastembed import TextEmbedding
from fastapi import FastAPI from fastapi import FastAPI
from pydantic import BaseModel from pydantic import BaseModel
model = TextEmbedding("snowflake/snowflake-arctic-embed-xs") models = {}
app = FastAPI() app = FastAPI()
class EmbeddingRequest(BaseModel): class EmbeddingRequest(BaseModel):
model: str model: str
prompt: str input: str
@app.post("/api/embeddings") @app.post("/v1/embeddings")
def embeddings(request: EmbeddingRequest): def embeddings(request: EmbeddingRequest):
embeddings = next(model.embed(request.prompt)).tolist() model = models.get(request.model) or TextEmbedding(request.model)
return {"embedding": embeddings} models[request.model] = model
embeddings = next(model.embed(request.input)).tolist()
return {"data": [{"embedding": embeddings}]}
if __name__ == "__main__": if __name__ == "__main__":
import uvicorn import uvicorn