#!/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 = ["reject", "low_value", "keep", "high_value"] LABEL2ID = {label: i for i, label in enumerate(LABELS)} ID2LABEL = {i: label for label, i in LABEL2ID.items()} def json_safe(value): if isinstance(value, Path): return str(value) if isinstance(value, dict): return {key: json_safe(item) for key, item in value.items()} if isinstance(value, (list, tuple)): return [json_safe(item) for item in value] return value 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): sample_weight = inputs.pop("sample_weight", None) labels = inputs.pop("labels") outputs = model(**inputs) logits = outputs.logits class_weights = self.class_weights.to(logits.device) if self.class_weights is not None else None loss = F.cross_entropy(logits, labels, weight=class_weights, reduction="none") if sample_weight is not None: sample_weight = sample_weight.to(logits.device).float() loss = loss * sample_weight loss = loss.sum() / sample_weight.sum().clamp_min(1e-6) else: loss = loss.mean() 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-large") parser.add_argument("--data-dir", type=Path, default=Path("/home/clouduser/hiennm/ja_filter_classifier/data_4class")) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--max-length", type=int, default=512) parser.add_argument("--learning-rate", type=float, default=1e-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=24) 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("--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("--max-train-samples", type=int, default=0) parser.add_argument("--max-eval-samples", type=int, default=0) parser.add_argument("--class-weight", choices=["none", "sqrt_balanced", "balanced"], default="sqrt_balanced") parser.add_argument("--label-smoothing", type=float, default=0.0) parser.add_argument("--gradient-checkpointing", action="store_true") return parser.parse_args() def class_weights(counts: Counter, mode: str) -> torch.Tensor | None: if mode == "none": return None total = sum(counts.values()) values = [] for label in LABELS: value = total / (len(LABELS) * max(1, counts[label])) if mode == "sqrt_balanced": value = math.sqrt(value) values.append(value) arr = np.array(values, dtype=np.float32) arr = arr / arr.mean() return torch.tensor(arr, dtype=torch.float32) def compute_metrics(eval_pred): logits, labels = eval_pred preds = logits.argmax(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), "reject_f1": f1_score(labels == LABEL2ID["reject"], preds == LABEL2ID["reject"], zero_division=0), "keep_plus_f1": f1_score(labels >= LABEL2ID["keep"], preds >= LABEL2ID["keep"], 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"])))) counts = Counter(raw["train"]["label"]) weights = class_weights(counts, args.class_weight) tokenizer = AutoTokenizer.from_pretrained(args.model_name) def encode_labels(batch): return {"labels": [LABEL2ID[label] for label in batch["label"]]} def tokenize(batch): return tokenizer(batch["text"], max_length=args.max_length, truncation=True) tokenized = raw.map(tokenize, batched=True, num_proc=args.num_proc, desc="tokenize") tokenized = tokenized.map(encode_labels, batched=True, desc="labels") remove_cols = [c for c in ["text", "label", "source_d", "raw_w", "url", "line_no"] if c in tokenized["train"].column_names] tokenized = tokenized.remove_columns(remove_cols) model = AutoModelForSequenceClassification.from_pretrained( args.model_name, num_labels=len(LABELS), id2label=ID2LABEL, label2id=LABEL2ID, ) if args.gradient_checkpointing: model.gradient_checkpointing_enable() train_args = TrainingArguments( output_dir=str(args.output_dir), learning_rate=args.learning_rate, weight_decay=args.weight_decay, 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, 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, load_best_model_at_end=True, metric_for_best_model="macro_f1", greater_is_better=True, logging_steps=50, bf16=True, fp16=False, dataloader_num_workers=4, report_to=[], seed=args.seed, ddp_find_unused_parameters=False, label_smoothing_factor=args.label_smoothing, remove_unused_columns=False, ) trainer = WeightedTrainer( model=model, args=train_args, train_dataset=tokenized["train"], eval_dataset=tokenized["test"], data_collator=DataCollatorWithPadding(tokenizer=tokenizer), compute_metrics=compute_metrics, class_weights=weights, ) train_result = trainer.train() eval_metrics = trainer.evaluate() pred = trainer.predict(tokenized["test"]) y_true = pred.label_ids y_pred = pred.predictions.argmax(axis=-1) report = classification_report( y_true, y_pred, labels=list(range(len(LABELS))), target_names=LABELS, output_dict=True, zero_division=0, ) result = { "labels": LABELS, "model_name": args.model_name, "hyperparameters": json_safe(vars(args)), "train_counts": dict(counts), "class_weights": weights.tolist() if weights is not None else None, "train_metrics": train_result.metrics, "eval_metrics": eval_metrics, "classification_report": report, "confusion_matrix": confusion_matrix(y_true, y_pred, labels=list(range(len(LABELS)))).tolist(), } is_main = trainer.is_world_process_zero() if is_main: (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")) if is_main: tokenizer.save_pretrained(str(args.output_dir / "best_model")) print(json.dumps(eval_metrics, ensure_ascii=False, indent=2)) if torch.distributed.is_available() and torch.distributed.is_initialized(): torch.distributed.barrier() torch.distributed.destroy_process_group() if __name__ == "__main__": main()