| import gradio as gr |
| import torch |
| import json |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification |
|
|
| MODEL_PATH = "./model" |
|
|
| def load_model(): |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=False) |
| model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH) |
| model.eval() |
|
|
| with open(f"{MODEL_PATH}/label_mapping.json") as f: |
| label_mapping = json.load(f) |
|
|
| return tokenizer, model, label_mapping |
|
|
| tokenizer, model, label_mapping = load_model() |
|
|
| def predict(text: str): |
| if not text.strip(): |
| return {} |
|
|
| inputs = tokenizer( |
| text, |
| return_tensors="pt", |
| padding="max_length", |
| truncation=True, |
| max_length=128, |
| ) |
| with torch.no_grad(): |
| logits = model(**inputs).logits |
|
|
| probs = torch.softmax(logits, dim=-1).squeeze() |
| scores = {label_mapping[str(i)]: float(probs[i]) for i in range(len(probs))} |
| return scores |
|
|
| demo = gr.Interface( |
| fn=predict, |
| inputs=gr.Textbox( |
| lines=5, |
| placeholder="Write something", |
| label="Example Tweet" |
| ), |
| outputs=gr.Label( |
| num_top_classes=len(label_mapping), |
| label="Predicted Label" |
| ), |
| title="Radical Speech Detector", |
| description="Enter a text and the model will classify it into one of the radical speech categories.\n\n\nBy Sofia Maldonado, Viviana Toledo & Oscar Rocha", |
| ) |
|
|
| demo.launch() |