Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
-from transformers import pipeline
-pipe=pipeline('sentiment-analysis','alibidaran/Symptom2disease')
+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,
@@ -25,11 +27,16 @@ label_2id={'Psoriasis': 0,
 'drug reaction': 21,
 'peptic ulcer disease': 22,
 'diabetes': 23}
-id2_label={f'LABEL_{i}':v for v,i in label_2id.items()}
 def detect_symptom(symptoms):
-  output=pipe(symptoms)[0]
-  label=id2_label[output['label']]
-  return f"You have {label} disease."
+  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='label')
+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()