import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch 
tokenizer = AutoTokenizer.from_pretrained("alibidaran/Symptom2disease")
model = AutoModelForSequenceClassification.from_pretrained("alibidaran/Symptom2disease")
label_2id={'Psoriasis': 0,
'Varicose Veins': 1,
'Typhoid': 2,
'Chicken pox': 3,
'Impetigo': 4,
'Dengue': 5,
'Fungal infection': 6,
'Common Cold': 7,
'Pneumonia': 8,
'Dimorphic Hemorrhoids': 9,
'Arthritis': 10,
'Acne': 11,
'Bronchial Asthma': 12,
'Hypertension': 13,
'Migraine': 14,
'Cervical spondylosis': 15,
'Jaundice': 16,
'Malaria': 17,
'urinary tract infection': 18,
'allergy': 19,
'gastroesophageal reflux disease': 20,
'drug reaction': 21,
'peptic ulcer disease': 22,
'diabetes': 23}
id2_label={i:v for v,i in label_2id.items()}

def detect_symptom(symptoms):
  inputs_id=tokenizer(symptoms,padding=True,truncation=True,return_tensors="pt")
  output=model(inputs_id['input_ids'])
  preds=torch.nn.functional.softmax(output.logits,-1).topk(5)
  results={id2_label[preds.indices[0][i].item()]:preds.values[0][i].item() for i in range(5)} 
  return results

demo=gr.Interface(fn=detect_symptom,inputs='text',outputs=gr.Label(num_top_classes=5),
                  examples=["I can't stop sneezing and I feel really tired and crummy. My throat is really sore",
                            "I have been experiencing a severe headache that is accompanied by pain behind my eyes.",
                            "There are small red spots all over my body that I can't explain. It's worrying me.",
                            "I've been having a really hard time going to the bathroom lately. It's really painful"])
demo.launch()