bart-mnli.py

import logging

from transformers.modeling_bart import BartForSequenceClassification
from transformers.pipelines import TextClassificationPipeline
from transformers.tokenization_bart import BartTokenizer


logging.basicConfig(level=logging.INFO)





model = BartForSequenceClassification.from_pretrained("bart-large-mnli")
tokenizer = BartTokenizer.from_pretrained("bart-large-mnli")

pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer)

x = ["BART is a seq2seq model.", "BART is not sequence to sequence."]
y = ["BART is denoising autoencoder.", "BART is version of autoencoder."]

print(pipeline([x]))
# [{'label': 'contradiction', 'score': 0.9971421}]

print(pipeline([y]))
# [{'label': 'entailment', 'score': 0.93996036}]