"""Training script for Myanmar Ghost sentiment model.""" import argparse import logging import sys from pathlib import Path from typing import Any, Dict, List, Optional import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, Dataset from tqdm import tqdm # Add parent directory to path sys.path.insert(0, str(Path(__file__).parent.parent.parent)) from src.utils.logger import setup_logger from src.utils.metrics import compute_metrics, MetricsTracker logger = setup_logger("train", log_dir="outputs/logs") class SentimentDataset(Dataset): """Dataset for sentiment classification.""" def __init__( self, data: List[Dict], tokenizer, max_length: int = 512, label_mapping: Dict[str, int] = None, ): self.data = data self.tokenizer = tokenizer self.max_length = max_length self.label_mapping = label_mapping or { "negative": 0, "neutral": 1, "positive": 2, "sarcastic": 3, } def __len__(self) -> int: return len(self.data) def __getitem__(self, idx: int) -> tuple: item = self.data[idx] encoding = self.tokenizer( item["text"], truncation=True, max_length=self.max_length, padding="max_length", return_tensors="pt", ) label = self.label_mapping.get(item.get("label", "neutral"), 1) return ( encoding["input_ids"].squeeze(0), encoding["attention_mask"].squeeze(0), torch.tensor(label, dtype=torch.long), ) def train_epoch( model: nn.Module, dataloader: DataLoader, criterion: nn.Module, optimizer: optim.Optimizer, device: torch.device, scheduler: Optional[Any] = None, ) -> Dict[str, float]: """Train for one epoch.""" model.train() total_loss = 0.0 all_predictions = [] all_labels = [] progress_bar = tqdm(dataloader, desc="Training") for batch_idx, (input_ids, attention_mask, labels) in enumerate(progress_bar): input_ids = input_ids.to(device) attention_mask = attention_mask.to(device) labels = labels.to(device) optimizer.zero_grad() outputs = model(input_ids, attention_mask) loss = criterion(outputs, labels) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() if scheduler: scheduler.step() total_loss += loss.item() predictions = outputs.argmax(dim=-1).cpu().tolist() all_predictions.extend(predictions) all_labels.extend(labels.cpu().tolist()) progress_bar.set_postfix({"loss": loss.item()}) metrics = compute_metrics(all_predictions, all_labels) metrics["loss"] = total_loss / len(dataloader) return metrics def evaluate( model: nn.Module, dataloader: DataLoader, criterion: nn.Module, device: torch.device, ) -> Dict[str, float]: """Evaluate the model.""" model.eval() total_loss = 0.0 all_predictions = [] all_labels = [] with torch.no_grad(): for input_ids, attention_mask, labels in tqdm(dataloader, desc="Evaluating"): input_ids = input_ids.to(device) attention_mask = attention_mask.to(device) labels = labels.to(device) outputs = model(input_ids, attention_mask) loss = criterion(outputs, labels) total_loss += loss.item() predictions = outputs.argmax(dim=-1).cpu().tolist() all_predictions.extend(predictions) all_labels.extend(labels.cpu().tolist()) metrics = compute_metrics(all_predictions, all_labels) metrics["loss"] = total_loss / len(dataloader) return metrics def load_data(data_path: str) -> List[Dict]: """Load training data from JSON or JSONL file.""" import json data = [] if data_path.endswith(".jsonl"): with open(data_path, "r", encoding="utf-8") as f: for line in f: if line.strip(): data.append(json.loads(line)) elif data_path.endswith(".json"): with open(data_path, "r", encoding="utf-8") as f: data = json.load(f) else: raise ValueError(f"Unsupported file format: {data_path}") return data def main(args): """Main training function.""" logger.info("Starting training...") logger.info(f"Arguments: {vars(args)}") # Device device = torch.device( "cuda" if torch.cuda.is_available() and not args.cpu else "cpu" ) logger.info(f"Using device: {device}") # Load tokenizer logger.info(f"Loading tokenizer from {args.model_name}") from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(args.model_name) # Load data logger.info(f"Loading data from {args.train_data}") train_data = load_data(args.train_data) val_data = load_data(args.val_data) if args.val_data else [] logger.info(f"Train samples: {len(train_data)}, Val samples: {len(val_data)}") # Create datasets train_dataset = SentimentDataset(train_data, tokenizer, args.max_length) train_loader = DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=2, ) val_loader = None if val_data: val_dataset = SentimentDataset(val_data, tokenizer, args.max_length) val_loader = DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=2, ) # Create model logger.info("Creating model...") from src.models.transformer_model import TransformerSentimentModel model = TransformerSentimentModel( model_name=args.model_name, num_labels=4, dropout=args.dropout, freeze_encoder=args.freeze_encoder, ) model.to(device) logger.info(f"Model parameters: {model.get_num_parameters():,}") logger.info(f"Trainable: {model.get_num_trainable_parameters():,}") # Loss and optimizer criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW( model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, ) # Scheduler total_steps = len(train_loader) * args.num_epochs warmup_steps = int(total_steps * 0.1) scheduler = optim.lr_scheduler.LinearLR( optimizer, start_factor=0.1, total_iters=warmup_steps, ) # Training loop metrics_tracker = MetricsTracker( metrics=["loss", "accuracy", "f1_weighted"], ) best_f1 = 0.0 best_model_path = Path(args.output_dir) / "best_model.pt" for epoch in range(args.num_epochs): logger.info(f"\nEpoch {epoch + 1}/{args.num_epochs}") # Train train_metrics = train_epoch( model, train_loader, criterion, optimizer, device, scheduler ) logger.info(f"Train - Loss: {train_metrics['loss']:.4f}, " f"Acc: {train_metrics['accuracy']:.4f}, " f"F1: {train_metrics['f1_weighted']:.4f}") # Evaluate if val_loader: val_metrics = evaluate(model, val_loader, criterion, device) logger.info(f"Val - Loss: {val_metrics['loss']:.4f}, " f"Acc: {val_metrics['accuracy']:.4f}, " f"F1: {val_metrics['f1_weighted']:.4f}") metrics = {"train_" + k: v for k, v in train_metrics.items()} metrics.update({"val_" + k: v for k, v in val_metrics.items()}) else: metrics = {"train_" + k: v for k, v in train_metrics.items()} metrics_tracker.update(metrics, epoch) # Save best model current_f1 = train_metrics.get("f1_weighted", 0) if current_f1 > best_f1: best_f1 = current_f1 model.save(str(best_model_path)) logger.info(f"Saved best model (F1: {best_f1:.4f})") # Save final model final_path = Path(args.output_dir) / "final_model.pt" model.save(str(final_path)) logger.info(f"\nTraining complete! Best F1: {best_f1:.4f}") return best_f1 if __name__ == "__main__": parser = argparse.ArgumentParser(description="Train Myanmar Ghost model") # Data arguments parser.add_argument("--train_data", type=str, required=True, help="Training data file") parser.add_argument("--val_data", type=str, default=None, help="Validation data file") parser.add_argument("--output_dir", type=str, default="outputs/models", help="Output directory") # Model arguments parser.add_argument("--model_name", type=str, default="bert-base-multilingual-cased") parser.add_argument("--max_length", type=int, default=512) parser.add_argument("--dropout", type=float, default=0.1) parser.add_argument("--freeze_encoder", action="store_true") # Training arguments parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--num_epochs", type=int, default=10) parser.add_argument("--learning_rate", type=float, default=5e-5) parser.add_argument("--weight_decay", type=float, default=0.01) parser.add_argument("--cpu", action="store_true", help="Use CPU only") args = parser.parse_args() main(args)