semantic search works!!

This commit is contained in:
Tienson Qin
2023-04-10 02:58:22 +08:00
parent 5d9a3b2f52
commit 02c1d066d5
3 changed files with 113 additions and 0 deletions

17
scripts/ai/main.py Normal file
View File

@@ -0,0 +1,17 @@
from typing import Optional, List
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
class Texts(BaseModel):
texts: List
app = FastAPI()
model = SentenceTransformer("all-MiniLM-L6-v2")
@app.post("/embedding/")
async def embedding(texts: Texts):
data = model.encode(list(texts), convert_to_numpy=True).tolist()
return {"embedding": data}