File size: 842 Bytes
d99519d 8aaa3ed d99519d 8aaa3ed d99519d 8aaa3ed d99519d 21f924e d99519d 8aaa3ed 57f74cd 8aaa3ed d99519d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
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()
|