fastAPI / app.py
Almaatla's picture
Update app.py
ac9faef verified
raw
history blame
928 Bytes
from fastapi import FastAPI, Request, Query
from fastapi.templating import Jinja2Templates
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=".")
@app.get("/")
def read_root(request: Request):
return templates.TemplateResponse("index.html", {"request": request})
@app.get("/embed")
def embed_string(text: str):
embedding = model.encode([text])
index.add(np.array(embedding))
return {"message": "String embedded and added to FAISS database"}
@app.get("/search")
def search_string(text: str, n: int = 5):
embedding = model.encode([text])
distances, indices = index.search(np.array(embedding), n)
return {"distances": distances[0].tolist(), "indices": indices[0].tolist()}