File size: 1,848 Bytes
eed2432 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import numpy as np
# Load the model and tokenizer
model_name = "AventIQ-AI/distilbert-disease-specialist-recommendation"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# Define the labels
labels = ["Cardiology", "Neurology", "Orthopedics", "Dermatology"]
# Define the prediction function
def recommend_specialist(symptoms):
inputs = tokenizer(symptoms, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
outputs = model(**inputs)
probabilities = torch.nn.functional.softmax(outputs.logits, dim=1)[0]
confidence, predicted_class = torch.max(probabilities, dim=0)
recommended_specialist = labels[predicted_class]
return f"Recommended Specialist: {recommended_specialist} (Confidence: {confidence.item():.2f})"
# Create the Gradio interface
iface = gr.Interface(
fn=recommend_specialist,
inputs=gr.Textbox(label="📝 Describe Your Symptoms", placeholder="e.g., experiencing chest pain and shortness of breath...", lines=3),
outputs=gr.Textbox(label="🔍 Specialist Recommendation", interactive=True),
title="🏥 Medical Specialist Recommender",
description="Enter your symptoms to receive a recommendation for the appropriate medical specialist.",
theme="compact",
allow_flagging="never",
examples=[
["I have persistent headaches and dizziness."],
["My skin has developed a red, itchy rash."],
["I'm experiencing joint pain and stiffness in the mornings."]
],
)
if __name__ == "__main__":
iface.launch()
|