Update app.py

app.py

@@ -3,6 +3,16 @@ from transformers import AutoTokenizer, BertForSequenceClassification
 
 import gradio as gr
 
-gr.Interface.load("models/AkshatSurolia/ICD-10-Code-Prediction", outputs="text",
-                title="ICD 10 Code Prediction", description="type in a disease and it will predict its ICD-10 codes",).launch(debug=True)
+tokenizer = AutoTokenizer.from_pretrained("AkshatSurolia/ICD-10-Code-Prediction")
+model = BertForSequenceClassification.from_pretrained("AkshatSurolia/ICD-10-Code-Prediction")
+config = model.config
 
+def prompt(text):
+    encoded_input = tokenizer(text, return_tensors='pt')
+    output = model(**encoded_input)
+    results = output.logits.detach().cpu().numpy()[0].argsort()[::-1][:5]
+    return [ config.id2label[ids] for ids in results]
+    #return tokenizer.decode(output[0][0].argmax(dim=-1))
+
+gr.Interface.load(fn=prompt, inputs="text", outputs="text",
+                title="ICD 10 Code Prediction", description="type in a disease and it will predict its ICD-10 codes",).launch(debug=True)