ActionBert-109M

ActionBERT 109M (fine tuned from bert-base-uncased)

This is a 109M parameter model fine-tuned to detect actions vs. non-action descriptions of stock images. It outputs a tensor with two values, one for the probability of it being an action and zero for the probability of it being a non-action.

Example Inference Code

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from tabulate import tabulate
import torch
import torch.nn.functional as F

# Load the model and tokenizer from the Hugging Face Hub
model_name = "qingy2024/ActionBERT-109M"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

text = "A mushroom exploding in a forest."

inputs = tokenizer(text, truncation=True, max_length=512, padding=True, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    action_probability, no_action_probability = F.softmax(outputs.logits[0], dim = 0).tolist()

    table = [
        ['Probability of Action', action_probability],
        ['Probability of No Action', no_action_probability]
    ]
    
    print(tabulate(table, tablefmt='rounded_grid'))

Output

╭──────────────────────────┬─────────────╮
│ Probability of Action    │ 0.999654    │
├──────────────────────────┼─────────────┤
│ Probability of No Action │ 0.000346127 │
╰──────────────────────────┴─────────────╯