Views
No views yet
1tokenizer = AutoTokenizer.from_pretrained("bert-base_uncased")
2
3# Function to tokenize data with
4def tokenize_function(batch):
5 return tokenizer(
6 batch['text'],
7 truncation=True,
8 max_length=512 # Can't be greater than model max length
9 )
10
11# Tokenize Data
12train_data = dataset['train'].map(tokenize_function, batched=True)
13test_data = dataset['test'].map(tokenize_function, batched=True)
14valid_data = dataset['valid'].map(tokenize_function, batched=True)
15
16# Convert lists to tensors
17train_data.set_format("torch", columns=['input_ids', "attention_mask", "label"])
18test_data.set_format("torch", columns=['input_ids', "attention_mask", "label"])
19valid_data.set_format("torch", columns=['input_ids', "attention_mask", "label"])
20
21model = AutoModelForSequenceClassification.from_pretrained(
22 MODEL_ID,
23 num_labels=5, # adjust this based on number of labels you're training on
24 device_map='cuda',
25 dtype='auto',
26 label2id=label2id,
27 id2label=id2label
28)
29
30# Metric function for evaluation in Trainer
31def compute_metrics(eval_pred):
32 predictions, labels = eval_pred
33 predictions = np.argmax(predictions, axis=1)
34
35 return {
36 'accuracy': accuracy_score(labels, predictions),
37 'f1_macro': f1_score(labels, predictions, average='macro'),
38 'f1_weighted': f1_score(labels, predictions, average='weighted')
39 }
40
41# Data collator to handle padding dynamically per batch
42data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
43
44training_args = TrainingArguments(
45 output_dir='./bert-comment',
46 num_train_epochs=2,
47 per_device_train_batch_size=32,
48 per_device_eval_batch_size=64,
49 learning_rate=2e-5,
50 weight_decay=0.01,
51 warmup_steps=300,
52
53 # Evaluation & saving
54 eval_strategy='epoch',
55 save_strategy='epoch',
56 load_best_model_at_end=True,
57 metric_for_best_model='f1_macro',
58
59 # Logging
60 logging_steps=100,
61 report_to='tensorboard',
62
63 # Other
64 seed=42,
65 fp16=torch.cuda.is_available(), # Mixed precision if GPU available
66)
67
68# Set up Trainer
69trainer = Trainer(
70 model=model,
71 args=training_args,
72 train_dataset=train_data,
73 eval_dataset=valid_data,
74 processing_class=tokenizer,
75 data_collator=data_collator,
76 compute_metrics=compute_metrics
77)
78
79# Train!
80trainer.train()
81
82# Evaluate
83eval_results = trainer.evaluate()
84print(eval_results)f1_macro| Training Loss | Epoch | Step | Validation Loss | Accuracy | F1 Macro | F1 Weighted |
|---|---|---|---|---|---|---|
| 0.6645 | 1.0 | 1540 | 0.6703 | 0.7275 | 0.7134 | 0.7292 |
| 0.5152 | 2.0 | 3080 | 0.6531 | 0.7444 | 0.7295 | 0.7451 |