rubrix-example.py

from transformers import pipeline
from datasets import load_dataset
import rubrix as rb

model = pipeline('zero-shot-classification', model="typeform/squeezebert-mnli")
dataset = load_dataset("ag_news", split='test')
# Labels are: 'World', 'Sports', 'Business', 'Sci/Tech'
labels = dataset.features["label"].names

for example in dataset:
    prediction = model(example['text'], labels)

    record = rb.TextClassificationRecord(
        inputs=example["text"],
        prediction=list(zip(prediction['labels'], prediction['scores'])),
        annotation=labels[example["label"]],
    )
    rb.log(record, name="ag_news_zeroshot")