0
GLUE_TASKS = ["cola", "mnli", "mnli-mm", "mrpc", "qnli", "qqp", "rte", "sst2", "stsb", "wnli"]

#useful in preprocessing, this sets what each task does. 
task_to_keys = {
    "cola": ("sentence", None),
    "mnli": ("premise", "hypothesis"),
    "mnli-mm": ("premise", "hypothesis"),
    "mrpc": ("sentence1", "sentence2"),
    "qnli": ("question", "sentence"),
    "qqp": ("question1", "question2"),
    "rte": ("sentence1", "sentence2"),
    "sst2": ("sentence", None),
    "stsb": ("sentence1", "sentence2"),
    "wnli": ("sentence1", "sentence2"),
}

def set_module_grad_status(module, flag=False):
    if isinstance(module, list):
        for m in module:
            set_module_grad_status(m, flag)
    else:
        for p in module.parameters():
            p.requires_grad = flag

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    if task != "stsb":
        predictions = np.argmax(predictions, axis=1)
    else:
        predictions = predictions[:, 0]
    return metric.compute(predictions=predictions, references=labels)

#setting some model parameters
task = "qnli"
sentence1_key, sentence2_key = task_to_keys[task]
# model_checkpoint = "test-glue/basemodel"

num_labels = 3 if task.startswith("mnli") else 1 if task=="stsb" else 2

metric_name = "pearson" if task == "stsb" else "matthews_correlation" if task == "cola" else "accuracy"
validation_key = "validation_mismatched" if task == "mnli-mm" else "validation_matched" if task == "mnli" else "validation"

model_checkpoint = "google/mobilebert-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint, num_labels=num_labels)
model = model.to(device)

args = TrainingArguments(
    "test-glue",
    evaluation_strategy = "epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=epochs,
    weight_decay=0.01,
    save_steps=0
    # load_best_model_at_end=True,
    # metric_for_best_model=metric_name,
)

#load dataset
actual_task = "mnli" if task == "mnli-mm" else task #some error checking.. 
dataset = load_dataset("glue", actual_task)
metric = load_metric('glue', actual_task)

# tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
def preprocess_function(examples):
    if sentence2_key is None:
        return tokenizer(examples[sentence1_key], truncation=True)
    return tokenizer(examples[sentence1_key], examples[sentence2_key], truncation=True)
encoded_dataset = dataset.map(preprocess_function, batched=True)


trainer = Trainer(
    model,
    args,
    train_dataset=encoded_dataset["train"],
    eval_dataset=encoded_dataset[validation_key],
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)
trainer.train()

我不太清楚为什么会收到错误“RuntimeError:张量 a (549) 的大小必须与非单维 1 处的张量 b (512) 的大小相匹配”,例如“trainer.train()”。就像不是每个样本都应该是一致的 512 大小。这只发生在 MobileBERT 模型上,而不是在 DistilBERT 模型上。我使用的代码主要基于 huggingface colab 教程https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification.ipynb#scrollTo=7k8ge1L1IrJk

4

0 回答 0