ivodz/bert-lv-complexity

1import gradio as gr2import torch3from transformers import BertTokenizer, BertForSequenceClassification4 5label_map = {0: "viegls", 1: "vidējs", 2: "sarežģīts"}6 7# Load tokenizer and model8tokenizer = BertTokenizer.from_pretrained("ivodz/bert-lv-complexity")9model = BertForSequenceClassification.from_pretrained("ivodz/bert-lv-complexity")10model.eval()11 12def classify(text):13    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)14    with torch.no_grad():15        logits = model(**inputs).logits16        probs = torch.nn.functional.softmax(logits, dim=1).squeeze()17    pred = torch.argmax(probs).item()18    return {19        "Viegls": round(probs[0].item(), 3),20        "Vidējs": round(probs[1].item(), 3),21        "Sarežģīts": round(probs[2].item(), 3)22    }, label_map[pred]23 24gr.Interface(25    fn=classify,26    inputs=gr.Textbox(label="Ievadiet tekstu"),27    outputs=[28        gr.Label(num_top_classes=3, label="Varbūtības"),29        gr.Textbox(label="Sarežģītības klase")30    ],31    title="Latviešu teksta sarežģītība ar mBERT").launch()32