Create app.py
Browse files
app.py
ADDED
|
@@ -0,0 +1,82 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
import pandas as pd
|
| 3 |
+
import faiss
|
| 4 |
+
import numpy as np
|
| 5 |
+
from sentence_transformers import SentenceTransformer
|
| 6 |
+
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
|
| 7 |
+
from fastapi import FastAPI
|
| 8 |
+
from pydantic import BaseModel
|
| 9 |
+
|
| 10 |
+
# 🔹 Initialize FastAPI
|
| 11 |
+
app = FastAPI()
|
| 12 |
+
|
| 13 |
+
# 🔹 Load AI Models
|
| 14 |
+
similarity_model = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")
|
| 15 |
+
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
|
| 16 |
+
summarization_model = AutoModelForSeq2SeqLM.from_pretrained("google/long-t5-tglobal-base")
|
| 17 |
+
summarization_tokenizer = AutoTokenizer.from_pretrained("google/long-t5-tglobal-base")
|
| 18 |
+
|
| 19 |
+
# 🔹 Load Datasets (Ensure files are uploaded to Hugging Face Space)
|
| 20 |
+
try:
|
| 21 |
+
recommendations_df = pd.read_csv("treatment_recommendations.csv")
|
| 22 |
+
questions_df = pd.read_csv("symptom_questions.csv")
|
| 23 |
+
except FileNotFoundError:
|
| 24 |
+
recommendations_df = pd.DataFrame(columns=["Disorder", "Treatment Recommendation"])
|
| 25 |
+
questions_df = pd.DataFrame(columns=["Questions"])
|
| 26 |
+
|
| 27 |
+
# 🔹 Create FAISS Index for Treatment Retrieval
|
| 28 |
+
if not recommendations_df.empty:
|
| 29 |
+
treatment_embeddings = similarity_model.encode(recommendations_df["Disorder"].tolist(), convert_to_numpy=True)
|
| 30 |
+
index = faiss.IndexFlatIP(treatment_embeddings.shape[1])
|
| 31 |
+
index.add(treatment_embeddings)
|
| 32 |
+
else:
|
| 33 |
+
index = None
|
| 34 |
+
|
| 35 |
+
# 🔹 Create FAISS Index for Question Retrieval
|
| 36 |
+
if not questions_df.empty:
|
| 37 |
+
question_embeddings = embedding_model.encode(questions_df["Questions"].tolist(), convert_to_numpy=True)
|
| 38 |
+
question_index = faiss.IndexFlatL2(question_embeddings.shape[1])
|
| 39 |
+
question_index.add(question_embeddings)
|
| 40 |
+
else:
|
| 41 |
+
question_index = None
|
| 42 |
+
|
| 43 |
+
# 🔹 API Request Model
|
| 44 |
+
class ChatRequest(BaseModel):
|
| 45 |
+
message: str
|
| 46 |
+
|
| 47 |
+
@app.post("/detect_disorders")
|
| 48 |
+
def detect_disorders(request: ChatRequest):
|
| 49 |
+
""" Detect psychiatric disorders from user input """
|
| 50 |
+
if index is None:
|
| 51 |
+
return {"error": "Dataset is missing or empty"}
|
| 52 |
+
|
| 53 |
+
text_embedding = similarity_model.encode([request.message], convert_to_numpy=True)
|
| 54 |
+
distances, indices = index.search(text_embedding, 3)
|
| 55 |
+
disorders = [recommendations_df["Disorder"].iloc[i] for i in indices[0]]
|
| 56 |
+
return {"disorders": disorders}
|
| 57 |
+
|
| 58 |
+
@app.post("/get_treatment")
|
| 59 |
+
def get_treatment(request: ChatRequest):
|
| 60 |
+
""" Retrieve treatment recommendations """
|
| 61 |
+
detected_disorders = detect_disorders(request)["disorders"]
|
| 62 |
+
treatments = {disorder: recommendations_df[recommendations_df["Disorder"] == disorder]["Treatment Recommendation"].values[0] for disorder in detected_disorders}
|
| 63 |
+
return {"treatments": treatments}
|
| 64 |
+
|
| 65 |
+
@app.post("/get_questions")
|
| 66 |
+
def get_recommended_questions(request: ChatRequest):
|
| 67 |
+
"""Retrieve the most relevant diagnostic questions based on patient symptoms."""
|
| 68 |
+
if question_index is None:
|
| 69 |
+
return {"error": "Questions dataset is missing or empty"}
|
| 70 |
+
|
| 71 |
+
input_embedding = embedding_model.encode([request.message], convert_to_numpy=True)
|
| 72 |
+
distances, indices = question_index.search(input_embedding, 3)
|
| 73 |
+
retrieved_questions = [questions_df["Questions"].iloc[i] for i in indices[0]]
|
| 74 |
+
return {"questions": retrieved_questions}
|
| 75 |
+
|
| 76 |
+
@app.post("/summarize_chat")
|
| 77 |
+
def summarize_chat(request: ChatRequest):
|
| 78 |
+
""" Summarize chat logs using LongT5 """
|
| 79 |
+
inputs = summarization_tokenizer("summarize: " + request.message, return_tensors="pt", max_length=4096, truncation=True)
|
| 80 |
+
summary_ids = summarization_model.generate(inputs.input_ids, max_length=500, num_beams=4, early_stopping=True)
|
| 81 |
+
summary = summarization_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
|
| 82 |
+
return {"summary": summary}
|