File size: 1,120 Bytes
e26d32e 5914320 a0edacc 3ab82e8 3eec3b2 00a8910 3ab82e8 3eec3b2 3f61915 5914320 a0edacc f238fcb a0edacc 5914320 3ab82e8 eb810c1 a0edacc 5914320 eb810c1 f238fcb 5914320 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
from fastapi import FastAPI, Request, Query
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
app = FastAPI()
model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
index = faiss.IndexFlatL2(384) # 384 is the dimensionality of the MiniLM model
templates = Jinja2Templates(directory=".")
class EmbedRequest(BaseModel):
texts: list[str]
class SearchRequest(BaseModel):
text: str
n: int = 5
@app.get("/")
def read_root(request: Request):
return templates.TemplateResponse("index.html", {"request": request})
@app.post("/embed")
def embed_strings(request: EmbedRequest):
embeddings = model.encode(request.texts)
index.add(np.array(embeddings))
return {"message": "Strings embedded and added to FAISS database"}
@app.post("/search")
def search_string(request: SearchRequest):
embedding = model.encode([request.text])
distances, indices = index.search(np.array(embedding), request.n)
return {"distances": distances[0].tolist(), "indices": indices[0].tolist()}
|