|
from typing import Any |
|
import uvicorn |
|
import json |
|
from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Body |
|
from fastapi.responses import JSONResponse, StreamingResponse |
|
from fastapi.middleware.cors import CORSMiddleware |
|
import pandas as pd |
|
from datasets import load_from_disk |
|
from transformers import AutoTokenizer, TFAutoModel |
|
from dotenv import load_dotenv |
|
|
|
load_dotenv() |
|
|
|
app = FastAPI() |
|
app.add_middleware( |
|
CORSMiddleware, |
|
allow_origins=["*"], |
|
allow_credentials=True, |
|
allow_methods=["*"], |
|
allow_headers=["*"], |
|
) |
|
|
|
model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1" |
|
tokenizer = AutoTokenizer.from_pretrained(model_ckpt) |
|
model = TFAutoModel.from_pretrained(model_ckpt, from_pt=True) |
|
|
|
def cls_pooling(model_output): |
|
return model_output.last_hidden_state[:, 0] |
|
|
|
def get_embeddings(text_list): |
|
encoded_input = tokenizer( |
|
text_list, padding=True, truncation=True, return_tensors="tf" |
|
) |
|
encoded_input = {k: v for k, v in encoded_input.items()} |
|
model_output = model(**encoded_input) |
|
return cls_pooling(model_output) |
|
|
|
embeddings_dataset = load_from_disk("data") |
|
embeddings_dataset.add_faiss_index(column="embeddings") |
|
|
|
def recommendations(question): |
|
question_embedding = get_embeddings([question]).numpy() |
|
scores, samples = embeddings_dataset.get_nearest_examples( |
|
"embeddings", question_embedding, k=5 |
|
) |
|
samples_df = pd.DataFrame.from_dict(samples) |
|
samples_df["scores"] = scores |
|
samples_df.sort_values("scores", ascending=False, inplace=True,ignore_index=True) |
|
return samples_df[['drugName', 'review', 'scores']] |
|
|
|
@app.post("/recommend") |
|
async def upload_image(question: str = Body(...,embed=True)): |
|
custom_recommendation_result = recommendations(question) |
|
custom_recommendation_result = custom_recommendation_result.to_dict(orient='records') |
|
custom_recommendation_result = json.dumps(custom_recommendation_result) |
|
return JSONResponse({"data":json.loads(custom_recommendation_result)},status_code=200) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
uvicorn.run(app, host="127.0.0.1", port=5000) |
