Update app.py

app.py

@@ -1,18 +1,8 @@
-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
+import gradio as gr
 
-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",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)
 
-demo = gr.Interface(fn=prompt,inputs= "text", outputs="text")
-demo.launch(debug=True)