Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4model_name = "MoritzLaurer/xtremedistil-l6-h256-mnli-fever-anli-ling-binary"
5tokenizer = AutoTokenizer.from_pretrained(model_name)
6model = AutoModelForSequenceClassification.from_pretrained(model_name)
7
8premise = "I first thought that I liked the movie, but upon second thought it was actually disappointing."
9hypothesis = "The movie was good."
10
11input = tokenizer(premise, hypothesis, truncation=True, return_tensors="pt")
12output = model(input["input_ids"].to(device)) # device = "cuda:0" or "cpu"
13prediction = torch.softmax(output["logits"][0], -1).tolist()
14label_names = ["entailment", "not_entailment"]
15prediction = {name: round(float(pred) * 100, 1) for pred, name in zip(prediction, label_names)}
16print(prediction)training_args = TrainingArguments(
num_train_epochs=5, # total number of training epochs
learning_rate=2e-05,
per_device_train_batch_size=32, # batch size per device during training
per_device_eval_batch_size=32, # batch size for evaluation
warmup_ratio=0.1, # number of warmup steps for learning rate scheduler
weight_decay=0.06, # strength of weight decay
fp16=True # mixed precision training
)| dataset | mnli-m-2c | mnli-mm-2c | fever-nli-2c | anli-all-2c | anli-r3-2c | lingnli-2c |
|---|---|---|---|---|---|---|
| accuracy | 0.897 | 0.898 | 0.861 | 0.607 | 0.62 | 0.827 |
| speed (text/sec, GPU Tesla P100, 128 batch) | 1490 | 1485 | 760 | 1186 | 1062 | 1791 |