| import gradio as gr | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig | |
| import json | |
| with open("tag_map.json") as tag_map_file: | |
| tag_map = json.load(tag_map_file) | |
| reverse_map = {j: i for i, j in tag_map.items()} | |
| model_name_or_path = "gpucce/ProSolAdv_full_train" | |
| config = AutoConfig.from_pretrained(model_name_or_path) | |
| config.num_classes = len(tag_map) | |
| model = AutoModelForSequenceClassification.from_pretrained( | |
| model_name_or_path, config=config | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) | |
| def classify(text): | |
| return ( | |
| reverse_map[ | |
| model(**tokenizer(text, return_tensors="pt")).logits.argmax(-1).item() | |
| ] | |
| .replace("_", " ") | |
| .capitalize() | |
| ) | |
| iface = gr.Interface(fn=classify, inputs="text", outputs="text") | |
| iface.launch() | |