train.py

from datasets import load_from_disk
from transformers import (
    AutoTokenizer, 
    AutoModelForSequenceClassification,
    TrainingArguments, 
    Trainer
)
import numpy as np
from sklearn.metrics import accuracy_score, f1_score
import torch
from torch import nn
import matplotlib.pyplot as plt

# Load data
dataset = load_from_disk("./data/wisesight")
train_data = dataset['train']
test_data = dataset['test']

# Load PhayaThaiBERT tokenizer
tokenizer = AutoTokenizer.from_pretrained("clicknext/phayathaibert")

# Tokenize function
def tokenize(batch):
    return tokenizer(batch['texts'], padding='max_length', truncation=True, max_length=128)

# Tokenize datasets
train_data = train_data.map(tokenize, batched=True)
test_data = test_data.map(tokenize, batched=True)

# Rename 'category' to 'labels' for Trainer
train_data = train_data.rename_column('category', 'labels')
test_data = test_data.rename_column('category', 'labels')

# Set format
train_data.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])
test_data.set_format('torch', columns=['input_ids', 'attention_mask', 'labels'])

# Load model (3 classes now)
model = AutoModelForSequenceClassification.from_pretrained(
    "clicknext/phayathaibert",
    num_labels=3  # Changed from 4
)

# Metrics
def compute_metrics(pred):  
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    acc = accuracy_score(labels, preds)
    f1 = f1_score(labels, preds, average='weighted')
    return {'accuracy': acc, 'f1': f1}

# Auto-detect hardware capabilities
def get_device_config():
    # Use fp16 for CUDA, but do NOT enable bf16 for MPS by default (may not be supported/
    # cause Trainer/accelerate issues). Use CPU fallback otherwise.
    if torch.cuda.is_available():
        return {"fp16": True}
    else:
        return {}
    
device_config = get_device_config()
print(f"Using device config: {device_config}")

# Training config
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=5,  # More epochs
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    learning_rate=3e-5,  # Slightly higher
    weight_decay=0.01,
    warmup_steps=500,  # Add warmup
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    logging_steps=100,
    **device_config,
)

class_counts = [3866, 11795, 5491]  # pos, neu, neg
total = sum(class_counts)
weights = [total/(len(class_counts)*count) for count in class_counts]
class_weights = torch.tensor(weights)

print(f"Class weights: pos={weights[0]:.2f}, neu={weights[1]:.2f}, neg={weights[2]:.2f}")


# Custom Trainer with weighted loss
class WeightedTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        loss_fct = nn.CrossEntropyLoss(weight=class_weights.to(model.device))
        loss = loss_fct(logits, labels)
        return (loss, outputs) if return_outputs else loss

# Use WeightedTrainer instead of Trainer
trainer = WeightedTrainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    eval_dataset=test_data,
    compute_metrics=compute_metrics,
)

# Train
print(">>> Starting training...")
trainer.train()


# Plot training history
history = trainer.state.log_history
train_loss = [x['loss'] for x in history if 'loss' in x]
eval_loss = [x['eval_loss'] for x in history if 'eval_loss' in x]
eval_acc = [x['eval_accuracy'] for x in history if 'eval_accuracy' in x]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

# Loss plot
if train_loss:
    ax1.plot(train_loss, label='Train Loss')
if eval_loss:
    # align validation loss points to training-step x-axis
    import numpy as np
    eval_x = np.linspace(0, max(len(train_loss)-1, 0), num=len(eval_loss))
    ax1.plot(eval_x, eval_loss, label='Val Loss', marker='o')
ax1.set_xlabel('Steps')
ax1.set_ylabel('Loss')
ax1.legend()
ax1.set_title('Training Loss')

# Accuracy plot
if eval_acc:
    ax2.plot(eval_acc, marker='o')
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Accuracy')
ax2.set_title('Validation Accuracy')

plt.tight_layout()
plt.savefig('training_history.png', dpi=150)
print(">>> Saved training_history.png")

# Save
model.save_pretrained("./model")
tokenizer.save_pretrained("./model")
print("[yes] Model saved to ./model/")