Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4model_name = "MoritzLaurer/MiniLM-L6-mnli"
5tokenizer = AutoTokenizer.from_pretrained(model_name)
6model = AutoModelForSequenceClassification.from_pretrained(model_name)
7
8premise = "I liked the movie"
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", "neutral", "contradiction"]
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
)