added gpt-2 fine tune training script

Author: paramkpr

notebooks/test.ipynb

@@ -11,6 +11,7 @@
 
 from src.models import build_teacher
 from src.data import ClassificationDataModule
+from utils.wandb_setup import setup_wandb
 
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
@@ -38,16 +39,7 @@ def main(config_path: Path = type.Argument(..., help="Path to YAML config")):
     cfg = yaml.safe_load(config_path.read_text())
 
     # --- SETUP W&B ---
-    run_name = cfg['training'].get("run_name", f"teacher_train_{Path(cfg['training']['output_dir']).name}")
-    report_to = cfg['training'].get("report_to", "none") # Default to no reporting
-    if report_to == "wandb":
-        project_name = cfg['training'].get("wandb_project", "senti_synth_teacher")
-        os.environ.pop("WANDB_DISABLED", None) # Ensure it's enabled if requested
-        os.environ["WANDB_PROJECT"] = project_name
-        logger.info(f"Reporting to W&B project: {project_name}")
-    else:
-        os.environ["WANDB_DISABLED"] = "true" # Explicitly disable
-        logger.info("W&B reporting disabled.")
+    run_name, report_to = setup_wandb(cfg)
 
     # --- BUILD MODEL ---
     model, tokenizer = build_teacher(cfg['model'])

src/cli/01_train_teacher.py

@@ -0,0 +1,114 @@
+import typer
+import yaml
+from pathlib import Path
+import logging
+
+import torch
+from transformers import DataCollatorForLanguageModeling, Trainer, TrainingArguments, IntervalStrategy
+
+from utils.wandb_setup import setup_wandb
+from utils.metrics import compute_metrics
+from models import build_generator
+from data import GeneratorDataModule
+
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+app = typer.Typer()
+
+
+@app.command()
+def main(config_path: Path = type.Argument(..., help="Path to YAML config")):
+    cfg = yaml.safe_load(config_path.read_text())
+
+    # --- SETUP W&B ---
+    run_name, report_to = setup_wandb(cfg)
+
+    # --- BUILD MODEL ---
+    model, tokenizer = build_generator(cfg['model'])
+
+    # --- SETUP DATA ---
+    data_module = GeneratorDataModule(cfg['data'], tokenizer)
+    data_module.setup()
+
+    train_dataset = data_module.get_train_dataset()
+    eval_dataset = data_module.get_eval_dataset()
+
+    # --- SETUP TRAINER ---
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer,
+        mlm=False,
+    )
+
+    training_args_dict = {
+        "output_dir": cfg['training']['output_dir'],
+        "overwrite_output_dir": cfg['training'].get("overwrite_output_dir", True),
+        "do_train": True,
+        "do_eval": eval_dataset is not None,
+        "per_device_train_batch_size": cfg['training'].get("per_device_train_batch_size", 8),
+        "per_device_eval_batch_size": cfg['training'].get("per_device_eval_batch_size", 16),
+        "gradient_accumulation_steps": cfg['training'].get("gradient_accumulation_steps", 1),
+        "num_train_epochs": cfg['training'].get("num_train_epochs", 3),
+        "learning_rate": cfg['training'].get("learning_rate", 5e-5),
+        "warmup_ratio": cfg['training'].get("warmup_ratio", 0.1),
+        "fp16": cfg['training'].get("fp16", torch.cuda.is_available()),
+        "logging_dir": cfg['training'].get("logging_dir", f"{cfg['training']['output_dir']}/logs"),
+        "logging_steps": cfg['training'].get("logging_steps", 100),
+        "eval_strategy": IntervalStrategy.STEPS if eval_dataset is not None else IntervalStrategy.NO,
+        "eval_steps": cfg['training'].get("eval_steps", 500),
+        "save_strategy": IntervalStrategy.STEPS,
+        "save_steps": cfg['training'].get("save_steps", 500),
+        "save_total_limit": cfg['training'].get("save_total_limit", 2),
+        "load_best_model_at_end": cfg['training'].get("load_best_model_at_end", eval_dataset is not None),
+        "metric_for_best_model": cfg['training'].get("metric_for_best_model", "eval_loss" if eval_dataset else None),
+        "greater_is_better": cfg['training'].get("greater_is_better", False),
+        "report_to": [report_to] if report_to != "none" else [],
+        "run_name": run_name,
+        "remove_unused_columns": False,
+        "ddp_find_unused_parameters": cfg['training'].get("ddp_find_unused_parameters", False),
+    }
+
+    training_args = TrainingArguments(**training_args_dict)
+    logger.info(f"Training arguments: {training_args}. FP16 Enabled: {training_args.fp16}")
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics if eval_dataset is not None else None,
+    )
+
+    # --- TRAIN ---
+    logger.info("Training model...")
+    train_result = trainer.train()
+    logger.info(f"Training results: {train_result}")
+
+    # Save final model & metrics
+    logger.info(f"Saving best model to {training_args.output_dir}")
+    trainer.save_model() # Saves the best model due to load_best_model_at_end=True
+    trainer.save_state()
+
+    # Log final metrics
+    metrics = train_result.metrics
+    trainer.log_metrics("train", metrics)
+    trainer.save_metrics("train", metrics)
+
+    # Evaluate on test set if available
+    test_dataset = data_module.get_test_dataset()
+    if test_dataset and cfg['training'].get("do_test_eval", True):
+        logger.info("Evaluating on test set...")
+        test_metrics = trainer.evaluate(eval_dataset=test_dataset, metric_key_prefix="test")
+        trainer.log_metrics("test", test_metrics)
+        trainer.save_metrics("test", test_metrics)
+        logger.info(f"Test set evaluation complete: {test_metrics}")
+
+
+    logger.info("Script finished successfully.")
+
+
+if __name__ == "__main__":
+    app()

src/cli/02_fine_tune_generator.py

@@ -1,6 +1,6 @@
 import logging
 from datasets import load_from_disk, DatasetDict
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, GPT2Tokenizer
 logger = logging.getLogger(__name__)
 
 
@@ -76,3 +76,76 @@ def get_test_dataset(self):
         if not self.tokenized_datasets: self.setup()  # noqa: E701
         return self.tokenized_datasets["test"]
 
+
+
+class GeneratorDataModule:
+    """
+    Data module for generative fine-tuning tasks.
+    Handles text generation setup.
+    """
+    def __init__(self, cfg: dict, tokenizer: GPT2Tokenizer):
+        self.cfg = cfg
+        self.tokenizer = tokenizer
+        self.dataset_path = cfg.get("dataset_path", None)
+
+        self.max_len = cfg.get("max_len", 32)
+
+        self.tokenized_datasets = None
+
+        self.required_splits = ["train", "val", "sanity", "test"]
+        self.text_column = "text"
+
+    def _load_clean_dataset(self) -> DatasetDict:
+        logger.info(f"Loading dataset from: {self.dataset_path}")
+        dataset = load_from_disk(self.dataset_path)
+        
+        missing_splits = [s for s in self.required_splits if s not in dataset]
+        if missing_splits:
+            raise ValueError(f"Dataset missing splits: {missing_splits}")
+        
+        return dataset
+    
+    def _tokenize_function(self, examples):
+        """Tokenization function for map."""
+        return self.tokenizer(
+            examples[self.text_column],
+            truncation=True,
+            padding=False,
+            max_length=self.max_len
+        )
+    
+    def setup(self):
+        """Loads and tokenizes the dataset."""
+        if self.tokenized_datasets:
+            return
+        
+        raw_datasets = self._load_clean_dataset()
+
+
+        self.tokenized_datasets = raw_datasets.map(
+            self._tokenize_function,
+            batched=True,
+            remove_columns=[c for c in raw_datasets["train"].column_names if c not in
+                             ["input_ids", "attention_mask", "labels"]]
+        )
+
+        logger.info(f"Loaded and tokenized datasets with max length: {self.max_len}")
+        logger.info(f"Columns in tokenized datasets: {self.tokenized_datasets[self.required_splits[0]].column_names}")
+
+
+    def get_train_dataset(self):
+        if not self.tokenized_datasets: self.setup()  # noqa: E701
+        return self.tokenized_datasets["train"]
+    
+    def get_eval_dataset(self):
+        if not self.tokenized_datasets: self.setup()  # noqa: E701
+        return self.tokenized_datasets["val"]
+
+    def get_sanity_dataset(self):
+        if not self.tokenized_datasets: self.setup()  # noqa: E701
+        return self.tokenized_datasets["sanity"]
+
+    def get_test_dataset(self):
+        if not self.tokenized_datasets: self.setup()  # noqa: E701
+        return self.tokenized_datasets["test"]
+        

src/data.py

@@ -3,6 +3,7 @@
 """
 import logging
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
 logger = logging.getLogger(__name__)
 
@@ -31,4 +32,25 @@ def build_teacher(cfg: dict):
     logger.info(f"Loading tokenizer for: {model_name}")
     tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=use_fast_tokenizer)
 
+    return model, tokenizer
+
+
+def build_generator(cfg: dict):
+    """
+    Builds and returns the generator model and tokenizer using Hugging Face.
+
+    Args:
+        cfg (dict): Configuration dictionary for the model, expecting keys like:
+
+    Returns:
+        tuple: (model, tokenizer)
+    """
+    model_name = cfg.get("model_name", "gpt2")
+    use_fast_tokenizer = cfg.get("use_fast_tokenizer", True)
+
+    logger.info(f"Loading generator model: {model_name}")
+    model = GPT2LMHeadModel.from_pretrained(model_name)
+    tokenizer = GPT2Tokenizer.from_pretrained(model_name, use_fast=use_fast_tokenizer)
+    tokenizer.pad_token = tokenizer.eos_token
+
     return model, tokenizer

src/models.py

@@ -0,0 +1,16 @@
+import numpy as np
+from sklearn.metrics import precision_recall_fscore_support, accuracy_score
+
+
+def compute_metrics(p):
+    """Computes metrics for HF Trainer."""
+    preds = np.argmax(p.predictions, axis=1)
+    labels = p.label_ids
+    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='binary') # Assuming binary
+    acc = accuracy_score(labels, preds)
+    return {
+        'accuracy': acc,
+        'f1': f1,
+        'precision': precision,
+        'recall': recall
+    }

src/utils/metrics.py

@@ -1,12 +1,19 @@
 import os
 import logging
-
+from pathlib import Path
 logger = logging.getLogger(__name__)
 
 
 def setup_wandb(cfg: dict):
     """Setup WANDB for logging."""
-    project_name = cfg.get("project_name", "senti-synth-teacher")
-    os.environ.pop("WANDB_DISABLED", None)
-    os.environ["WANDB_PROJECT"] = project_name
-    logger.info(f"Reporting to W&B project: {project_name}")
+    run_name = cfg['training'].get("run_name", f"teacher_train_{Path(cfg['training']['output_dir']).name}")
+    report_to = cfg['training'].get("report_to", "none") # Default to no reporting
+    if report_to == "wandb":
+        project_name = cfg['training'].get("wandb_project", "senti_synth_teacher")
+        os.environ.pop("WANDB_DISABLED", None) # Ensure it's enabled if requested
+        os.environ["WANDB_PROJECT"] = project_name
+        logger.info(f"Reporting to W&B project: {project_name}")
+    else:
+        os.environ["WANDB_DISABLED"] = "true" # Explicitly disable
+        logger.info("W&B reporting disabled.")
+    return run_name, report_to