Update app.py

app.py

@@ -1,3 +1,18 @@
 import gradio as gr
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, BertForSequenceClassification
+
+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("models/AkshatSurolia/ICD-10-Code-Prediction").launch()
+demo = gr.Interface(fn=prompt,inputs= "text", outputs="text")
+demo.launch(debug=True)