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}
|