update code and requirements

app.py

@@ -1,7 +1,36 @@
+from sentence_transformers import SentenceTransformer
 from fastapi import FastAPI
+import pickle
+import pandas as pd
+from pydantic import BaseModel
+
+corpus = pickle.load(open("./corpus/all_embeddings.pickle", "rb"))
+label_encoder = pickle.load("./corpus/label_encoder.pickle")
+model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+df = pd.DataFrame(data={"label": pickle.load(open("./corpus/y_all.pickle"))})
 
 app = FastAPI()
 
+class Disease(BaseModel):
+    id: int
+    name: str
+    score: float
+
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}
+
+@app.post("/", response_model=list[Disease])
+async def predict(query: str):
+    query_embedding = model.encode(query).astype('float')
+    similarity_vectors = model.similarity(q, all_embeddings)
+    scores, indicies = torch.topk(similarity_vectors, k=len(all_embeddings))
+    id = df.iloc[indicies]
+    id = df.drop_duplicates("label")
+    scores = scores[id.index]
+    diseases = label_encoder.inverse_transform(id.label.values)
+    id = id.label.values
+    diseases = [dict("id": value[0], "name": value[1], "score" : value[2]) for value in zip(id, diseases, scores)]
+    return diseases
+
+

requirements.txt

@@ -1,2 +1,5 @@
 fastapi
 uvicorn[standard]
+pandas
+sentence-transformers
+pydantic