|
from fastapi import FastAPI |
|
from pydantic import BaseModel |
|
from sentence_transformers import SentenceTransformer |
|
import numpy as np |
|
import json |
|
|
|
|
|
with open("data.json") as f: |
|
data = json.load(f) |
|
|
|
|
|
model = SentenceTransformer('all-MiniLM-L6-v2') |
|
|
|
|
|
file_names = [item["file_name"] for item in data] |
|
file_embeddings = model.encode(file_names) |
|
|
|
app = FastAPI() |
|
|
|
class Query(BaseModel): |
|
text: str |
|
|
|
@app.post("/search") |
|
async def search(query: Query): |
|
|
|
query_embedding = model.encode([query.text]) |
|
|
|
|
|
similarities = np.dot(file_embeddings, query_embedding.T).flatten() |
|
|
|
|
|
best_match_idx = np.argmax(similarities) |
|
|
|
return { |
|
"best_match": data[best_match_idx]["file_name"], |
|
"similarity_score": float(similarities[best_match_idx]) |
|
} |