#!/usr/bin/env python3 from __future__ import annotations import argparse import json import math from collections import Counter from pathlib import Path import numpy as np import torch import torch.nn.functional as F from datasets import load_dataset from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, f1_score from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding, Trainer, TrainingArguments, set_seed, ) LABELS = ["KH", "K", "KL", "D", "R", "X"] LABEL2ID = {label: idx for idx, label in enumerate(LABELS)} ID2LABEL = {idx: label for label, idx in LABEL2ID.items()} class WeightedTrainer(Trainer): def __init__(self, *args, class_weights: torch.Tensor | None = None, **kwargs): super().__init__(*args, **kwargs) self.class_weights = class_weights def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): labels = inputs.pop("labels") outputs = model(**inputs) logits = outputs.logits weights = self.class_weights.to(logits.device) if self.class_weights is not None else None loss = F.cross_entropy(logits, labels, weight=weights) return (loss, outputs) if return_outputs else loss def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--model-name", default="FacebookAI/xlm-roberta-base") parser.add_argument("--data-dir", type=Path, default=Path("/home/clouduser/hiennm/ja_filter_classifier/data")) parser.add_argument("--output-dir", type=Path, default=Path("/home/clouduser/hiennm/ja_filter_classifier/outputs/xlm-roberta-base")) parser.add_argument("--max-length", type=int, default=512) parser.add_argument("--learning-rate", type=float, default=2e-5) parser.add_argument("--weight-decay", type=float, default=0.01) parser.add_argument("--epochs", type=float, default=2.0) parser.add_argument("--batch-size", type=int, default=32) parser.add_argument("--gradient-accumulation-steps", type=int, default=1) parser.add_argument("--warmup-ratio", type=float, default=0.06) parser.add_argument("--seed", type=int, default=20260604) parser.add_argument("--max-train-samples", type=int, default=0) parser.add_argument("--max-eval-samples", type=int, default=0) parser.add_argument("--eval-steps", type=int, default=500) parser.add_argument("--save-steps", type=int, default=500) parser.add_argument("--num-proc", type=int, default=8) parser.add_argument("--class-weight", choices=["balanced", "sqrt_balanced", "none"], default="sqrt_balanced") return parser.parse_args() def tokenize_dataset(dataset, tokenizer, max_length: int, num_proc: int): def tokenize(batch): return tokenizer(batch["text"], truncation=True, max_length=max_length) return dataset.map( tokenize, batched=True, num_proc=num_proc, remove_columns=["text", "url", "line_no"], desc="tokenize", ) def encode_labels(dataset): def encode(batch): return {"labels": [LABEL2ID[label] for label in batch["label"]]} return dataset.map(encode, batched=True, remove_columns=["label"]) def compute_class_weights(label_counts: Counter, mode: str) -> torch.Tensor | None: if mode == "none": return None total = sum(label_counts.values()) num_classes = len(LABELS) weights = [] for label in LABELS: count = max(1, label_counts[label]) value = total / (num_classes * count) if mode == "sqrt_balanced": value = math.sqrt(value) weights.append(value) weights = np.array(weights, dtype=np.float32) weights = weights / weights.mean() return torch.tensor(weights, dtype=torch.float32) def metrics_fn(eval_pred): logits, labels = eval_pred preds = np.argmax(logits, axis=-1) return { "accuracy": accuracy_score(labels, preds), "macro_f1": f1_score(labels, preds, average="macro", zero_division=0), "weighted_f1": f1_score(labels, preds, average="weighted", zero_division=0), } def main() -> None: args = parse_args() set_seed(args.seed) args.output_dir.mkdir(parents=True, exist_ok=True) raw = load_dataset( "json", data_files={ "train": str(args.data_dir / "train.jsonl"), "test": str(args.data_dir / "test.jsonl"), }, ) if args.max_train_samples: raw["train"] = raw["train"].shuffle(seed=args.seed).select(range(min(args.max_train_samples, len(raw["train"])))) if args.max_eval_samples: raw["test"] = raw["test"].shuffle(seed=args.seed).select(range(min(args.max_eval_samples, len(raw["test"])))) train_counts = Counter(raw["train"]["label"]) class_weights = compute_class_weights(train_counts, args.class_weight) tokenizer = AutoTokenizer.from_pretrained(args.model_name) tokenized = tokenize_dataset(raw, tokenizer, args.max_length, args.num_proc) tokenized = encode_labels(tokenized) model = AutoModelForSequenceClassification.from_pretrained( args.model_name, num_labels=len(LABELS), id2label=ID2LABEL, label2id=LABEL2ID, ) train_args = TrainingArguments( output_dir=str(args.output_dir), learning_rate=args.learning_rate, per_device_train_batch_size=args.batch_size, per_device_eval_batch_size=args.batch_size * 2, gradient_accumulation_steps=args.gradient_accumulation_steps, num_train_epochs=args.epochs, weight_decay=args.weight_decay, warmup_ratio=args.warmup_ratio, eval_strategy="steps", eval_steps=args.eval_steps, save_strategy="steps", save_steps=args.save_steps, save_total_limit=2, logging_steps=50, load_best_model_at_end=True, metric_for_best_model="macro_f1", greater_is_better=True, fp16=False, bf16=True, dataloader_num_workers=4, report_to=[], seed=args.seed, ddp_find_unused_parameters=False, ) trainer = WeightedTrainer( model=model, args=train_args, train_dataset=tokenized["train"], eval_dataset=tokenized["test"], tokenizer=tokenizer, data_collator=DataCollatorWithPadding(tokenizer=tokenizer), compute_metrics=metrics_fn, class_weights=class_weights, ) train_result = trainer.train() metrics = trainer.evaluate() predictions = trainer.predict(tokenized["test"]) y_true = predictions.label_ids y_pred = np.argmax(predictions.predictions, axis=-1) report = classification_report( y_true, y_pred, labels=list(range(len(LABELS))), target_names=LABELS, output_dict=True, zero_division=0, ) matrix = confusion_matrix(y_true, y_pred, labels=list(range(len(LABELS)))).tolist() result = { "model_name": args.model_name, "output_dir": str(args.output_dir), "labels": LABELS, "hyperparameters": vars(args), "train_label_counts": dict(train_counts), "class_weights": class_weights.tolist() if class_weights is not None else None, "train_metrics": train_result.metrics, "eval_metrics": metrics, "classification_report": report, "confusion_matrix": matrix, } (args.output_dir / "final_report.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8") trainer.save_model(str(args.output_dir / "best_model")) tokenizer.save_pretrained(str(args.output_dir / "best_model")) print(json.dumps(result["eval_metrics"], ensure_ascii=False, indent=2)) if __name__ == "__main__": main()