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()