|
import torch |
|
from transformers import AutoTokenizer, BertForSequenceClassification |
|
|
|
import gradio as gr |
|
|
|
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] |
|
|
|
|
|
demo = gr.Interface(fn=prompt,inputs= "text", outputs="text") |
|
demo.launch(debug=True) |