modified code and add files

app.py

@@ -6,10 +6,11 @@ from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
 import torch
 
-corpus = pickle.load(open("./corpus/all_embeddings.pickle", "rb"))
-label_encoder = pickle.load(open("./corpus/label_encoder.pickle", "rb"))
-model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
-df = pd.DataFrame(data={"label": pickle.load(open("./corpus/y_all.pickle", "rb"))})
+corpus = pickle.load(open("./corpus/all_embeddings_disease.pickle", "rb"))
+# label_encoder = pickle.load(open("./corpus/label_encoder.pickle", "rb"))
+# model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-mpnet-base-v2')
+df = pd.DataFrame(pickle.load(open("./corpus/y_all_disease.pickle", "rb")))
 
 app = FastAPI()
 
@@ -24,27 +25,32 @@ app.add_middleware(
 class Disease(BaseModel):
     id: int
     name: str
+    url: str
     score: float
 
 class Symptoms(BaseModel):
     query: str
 
-@app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+# @app.get("/")
+# def greet_json():
+#     return {"Hello": "World!"}
 
 @app.post("/", response_model=list[Disease])
 async def predict(symptoms: Symptoms):
     query_embedding = model.encode(symptoms.query).astype('float')
     similarity_vectors = model.similarity(query_embedding, corpus)[0]
     scores, indicies = torch.topk(similarity_vectors, k=len(corpus))
-    # print("Similarity Vector Shape: ", similarity_vectors.shape)
-    # print("Scores Shape: ", scores.shape)
-    # print("Indicies Shape: ", indicies.shape)
-    id_ = df.iloc[indicies].reset_index(drop=True)
-    id_ = id_.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)]
+    # id_ = df.iloc[indicies].reset_index(drop=True)
+    df = df.iloc[indicies]
+    # id_ = id_.drop_duplicates("label")
+    df["scores"] = scores
+    # 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],
+                      "url" : value[3],
+                      })
+                      for value in zip(df.index, df["name"], df["scores"], df["url"])]
     return diseases

corpus/all_embeddings_disease.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11a003c5bc180aaff3d06b5f64ee28034512937629b635605a2bb56edd267ff9
+size 4045987

corpus/y_all_disease.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11a003c5bc180aaff3d06b5f64ee28034512937629b635605a2bb56edd267ff9
+size 4045987