erictransformer 0.0.1__py3-none-any.whl

erictransformer/loops/eval_loop.py

@@ -0,0 +1,111 @@
+import math
+import sys
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+from torch.amp import autocast
+from torch.nn import Module
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+
+from erictransformer.eval_models import EvalModel
+from erictransformer.exceptions import EricEvalError
+from erictransformer.utils import get_precision
+
+
+@dataclass(kw_only=True)
+class EvalResult:
+    loss: float
+    metrics: Optional[dict] = None
+
+
+def eval_loop(
+    model: Module,
+    dataloader: DataLoader,
+    eval_tok_data_cases: int,
+    eval_bs: int,
+    eval_models: List[EvalModel],
+    precision: str = "auto",
+) -> EvalResult:
+    total_loss = 0.0
+    total_examples = 0
+    pbar = tqdm(
+        total=math.ceil(eval_tok_data_cases / eval_bs),
+        desc="Model Evaluating",
+        position=1,
+        leave=False,
+        disable=not sys.stdout.isatty(),
+    )
+
+    for eval_model in eval_models:
+        eval_model.reset()
+
+    model.eval()
+
+    device = next(model.parameters()).device
+    precision_type = get_precision(device=device)
+
+    with torch.no_grad():
+        for batch in dataloader:
+            try:
+                try:
+                    input_ids = batch["input_ids"].to(device, non_blocking=True)
+                    attention_mask = batch["attention_mask"].to(
+                        device, non_blocking=True
+                    )
+                    labels = batch["labels"].to(device, non_blocking=True)
+                except Exception as e:
+                    raise EricEvalError(f"Failed to extract inputs from batch: {e}")
+
+                if device.type == "cuda":
+                    dtype = (
+                        torch.bfloat16 if precision_type == "bf16" else
+                        torch.float16 if precision_type == "fp16" else
+                        torch.float32
+                    )
+                    with autocast(device_type="cuda", enabled=(dtype != torch.float32), dtype=dtype):
+                        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+                else:
+                    outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+
+                bs = batch["input_ids"].size(0)
+                result = outputs.loss.item() * bs
+            except Exception as e:
+                raise EricEvalError(f"Error processing batch: {e}")
+
+            if not math.isnan(result):
+                total_loss += result
+                total_examples += bs
+            if eval_models:
+                batch["input_ids"] = batch["input_ids"].to(device)
+                batch["attention_mask"] = batch["attention_mask"].to(device)
+                batch["labels"] = batch["labels"].to(device)
+
+                for eval_model in eval_models:
+                    eval_model(batch, outputs)
+
+            pbar.update(1)
+
+    pbar.close()
+
+    custom_metrics = {}
+    for eval_model in eval_models:
+        name = eval_model.name
+        if name in custom_metrics:
+            for i in range(0, 1000):
+                new_name = f"{eval_model.name}-{i}"
+                if new_name not in custom_metrics:
+                    name = new_name
+                    break
+            else:
+                raise EricEvalError(
+                    f"Couldn't find unique name for eval model named {eval_model.name}"
+                )
+
+        custom_metrics[name] = eval_model.result()
+        eval_model.reset()
+
+    average_loss = total_loss / total_examples if total_examples else 0.0
+
+    return EvalResult(loss=average_loss, metrics=custom_metrics)

erictransformer/loops/train_loop.py

@@ -0,0 +1,310 @@
+import json
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+from torch.nn import Module
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+from transformers import PretrainedConfig, PreTrainedTokenizerBase
+
+from erictransformer.args.eric_args import EricTrainArgs
+from erictransformer.eval_models import EvalModel
+from erictransformer.exceptions import EricTrainError
+from erictransformer.eric_tracker.eric_tracker import EricTracker
+from erictransformer.loops.eval_loop import eval_loop
+from erictransformer.utils import EricTimer
+from erictransformer.utils.test import DebugHook
+
+
+@dataclass
+class TrainResult:
+    final_train_loss: float
+    final_eval_loss: Optional[float] = None
+    best_eval_loss: Optional[float] = None
+
+
+def train_loop(
+    args: EricTrainArgs,
+    train_dataloader: DataLoader,
+    eric_tracker: EricTracker,
+    train_steps: int,
+    model: Module,
+    optim: Optimizer,
+    lr_sched: LRScheduler,
+    eval_cases: int,
+    eval_dataloader: DataLoader,
+    tokenizer: PreTrainedTokenizerBase,
+    config: PretrainedConfig,
+    skip_cases: int,
+    starting_epoch: int,
+    eval_models: List[EvalModel],
+    eric_timer: EricTimer,
+    precision_type: str,
+) -> TrainResult:
+    best_tokenizer_config_saved = False
+    checkpoint_tokenizer_config_saved = False
+    first_epoch_train_iter = iter(train_dataloader)
+
+    if skip_cases:
+        try:
+            skip_batches = max(1, int(skip_cases / args.bs))
+            for _ in tqdm(range(skip_batches)):
+                next(first_epoch_train_iter, None)
+        except Exception as e:
+            raise EricTrainError(f"Failed to skip initial training cases: {e}")
+
+    device = next(model.parameters()).device
+
+    gas = max(1, args.gas)
+
+    start_step = eric_tracker.state.current_step
+
+    for epoch in range(starting_epoch - 1, int(args.epochs)):
+        epoch_iter = (
+            first_epoch_train_iter
+            if epoch == starting_epoch - 1
+            else iter(train_dataloader)
+        )
+
+        batch_idx = 0
+        while True:
+            if eric_tracker.state.current_step > train_steps:
+                break
+
+            with eric_timer.section("training_core", "iter"):
+                try:
+                    batch = next(epoch_iter)
+                except StopIteration:
+                    break
+
+            try:
+                with eric_timer.section("training_core", "to_device"):
+                    input_ids = batch["input_ids"].to(device, non_blocking=True)
+                    attention_mask = batch["attention_mask"].to(
+                        device, non_blocking=True
+                    )
+                    labels = batch["labels"].to(device, non_blocking=True)
+            except Exception as e:
+                raise EricTrainError(f"Failed to extract inputs from batch: {e}")
+
+            try:
+                with eric_timer.section("training_core", "forward"):
+                    if device.type == "cuda":
+                        with torch.autocast(
+                            device_type="cuda",
+                            enabled=precision_type != "fp32",
+                            dtype=(
+                                torch.float16
+                                if precision_type == "fp16"
+                                else torch.bfloat16
+                            ),
+                        ):
+                            outputs = model(
+                                input_ids=input_ids,
+                                attention_mask=attention_mask,
+                                labels=labels,
+                            )
+                    else:
+                        outputs = model(
+                            input_ids=input_ids,
+                            attention_mask=attention_mask,
+                            labels=labels,
+                        )
+
+                with eric_timer.section("training_core", "get loss"):
+                    raw_loss = outputs.loss.detach()
+                    loss = outputs.loss / gas
+
+            except Exception as e:
+                raise EricTrainError(f"Model forward pass failed: {e}")
+
+            try:
+                with eric_timer.section("training_core", "backwards"):
+                    loss.backward()
+            except Exception as e:
+                raise EricTrainError(f"Backward pass failed: {e}")
+
+            try:
+                do_step = ((batch_idx + 1) % gas == 0) or (
+                    (batch_idx + 1) == len(train_dataloader)
+                )
+
+                if do_step:
+                    if precision_type == "fp16":
+                        # for now we don't clip gradients when using bf16 or fp32 since this
+                        # operation takes a noticeable amount of time and they're stable enough.
+                        with eric_timer.section("training_core", "clip gradients"):
+                            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+
+                    with eric_timer.section("training_core", "optim.step()"):
+                        optim.step()
+                    with eric_timer.section("training_core", "lr_sched.step()"):
+                        lr_sched.step()
+                    with eric_timer.section("training_core", "optim.zero_grad()"):
+                        optim.zero_grad()
+                    eric_tracker.optim_step()
+
+            except Exception as e:
+                raise EricTrainError(f"optim or LR scheduler step failed: {e}")
+
+            try:
+                do_eval = bool(eric_tracker.time_to_eval() and eval_cases)
+
+                if do_eval:
+                    with eric_timer.section("training_extra", "eval loop"):
+                        model.eval()
+                        eval_result = eval_loop(
+                            model,
+                            eval_dataloader,
+                            eval_cases,
+                            args.eval_bs
+                            if args.eval_bs
+                            else args.bs,
+                            eval_models=eval_models,
+                        )
+                        model.train()
+
+                    global_eval_loss = float(eval_result.loss)
+
+                    is_best_model = eric_tracker.set_eval_loss(global_eval_loss)
+                    eric_tracker.set_metrics(eval_result.metrics)
+
+                    if args.save_best and is_best_model:
+                        with eric_timer.section("training_extra", "save best model"):
+                            try:
+                                path = eric_tracker.tracker_paths.best_model_path
+                                model.save_pretrained(path)
+                            except Exception as e:
+                                path = eric_tracker.tracker_paths.best_model_path
+                                raise EricTrainError(
+                                    f"Failed to save best model to: {path} |  {e}"
+                                )
+
+                            try:
+                                state_path = (
+                                    eric_tracker.tracker_paths.best_model_state_path
+                                )
+                                with open(state_path, "w") as f:
+                                    json.dump(
+                                        eric_tracker.state.to_dict(), f, indent=2
+                                    )
+                            except Exception as e:
+                                state_path = (
+                                    eric_tracker.tracker_paths.best_model_state_path
+                                )
+                                raise EricTrainError(
+                                    f"Failed to save best model state to: {state_path} | {e}"
+                                )
+
+                            if not best_tokenizer_config_saved:
+                                try:
+                                    tokenizer.save_pretrained(
+                                        eric_tracker.tracker_paths.best_model_path
+                                    )
+                                except Exception as e:
+                                    path = eric_tracker.tracker_paths.best_model_path
+                                    raise EricTrainError(
+                                        f"Failed to save tokenizer to: {path} | {e}"
+                                    )
+                                try:
+                                    config.save_pretrained(
+                                        eric_tracker.tracker_paths.best_model_path
+                                    )
+                                except Exception as e:
+                                    path = eric_tracker.tracker_paths.best_model_path
+                                    raise EricTrainError(
+                                        f"Failed to save config to: {path} | {e}"
+                                    )
+                                best_tokenizer_config_saved = True
+            except Exception as e:
+                raise EricTrainError(f"Evaluation loop failed: {e}")
+
+            try:
+                do_ckpt = eric_tracker.time_to_checkpoint()
+
+                if do_ckpt:
+                    with eric_timer.section("training_extra", "checkpoint"):
+                        checkpoint_path = eric_tracker.tracker_paths.checkpoint_path
+                        try:
+                            model.save_pretrained(checkpoint_path)
+                        except Exception as e:
+                            raise EricTrainError(
+                                f"Failed to save checkpoint model to: {checkpoint_path} | {e}"
+                            )
+
+                        if not checkpoint_tokenizer_config_saved:
+                            try:
+                                tokenizer.save_pretrained(checkpoint_path)
+                            except Exception as e:
+                                raise EricTrainError(
+                                    f"Failed to save tokenizer to: {checkpoint_path} |{e}"
+                                )
+                            try:
+                                config.save_pretrained(checkpoint_path)
+                            except Exception as e:
+                                raise EricTrainError(
+                                    f"Failed to save config to: {checkpoint_path} | {e}"
+                                )
+                            checkpoint_tokenizer_config_saved = True
+
+                        try:
+                            # Reloading the optimizer was causing problems. So for now we restart it.
+                            #torch.save(
+                            #     optim.state_dict(),
+                            #     eric_tracker.tracker_paths.optim_path,
+                            #  )
+
+                            torch.save(
+                                lr_sched.state_dict(),
+                                eric_tracker.tracker_paths.lr_sched_path,
+                            )
+
+                        except Exception as e:
+                            raise EricTrainError(
+                                f"Failed to save lr scheduler state to: "
+                                f"{eric_tracker.tracker_paths.optim_path} |  {e}"
+                            )
+            except Exception as e:
+                raise EricTrainError(str(e))
+
+            try:
+                with eric_timer.section("training_extra", "logging"):
+                    if (
+                        eric_tracker.time_to_log()
+                        or eric_tracker.time_to_eval()
+                        or eric_tracker.time_to_checkpoint()
+                    ):
+                        eric_timer.report()
+
+                    num_tokens = int(attention_mask.sum().item())
+                    eric_tracker.step(
+                        raw_loss, lr_sched.get_last_lr()[0], num_tokens=num_tokens
+                    )
+            except Exception as e:
+                raise EricTrainError(f"Tracker step failed: {e}")
+
+            debug_hook_post_checkpoint()
+
+            batch_idx += 1
+
+        eric_tracker.mark_epoch()
+
+        if eric_tracker.state.current_step >= train_steps:
+            break
+
+    debug_hook_steps({
+        "start_step": start_step,
+        "total_steps": eric_tracker.state.current_step
+    })
+
+    return TrainResult(
+        final_train_loss=eric_tracker.state.train_loss,
+        final_eval_loss=eric_tracker.state.eval_loss if eval_cases else None,
+        best_eval_loss=eric_tracker.state.best_eval_loss if eval_cases else None,
+    )
+
+debug_hook_post_checkpoint = DebugHook()
+debug_hook_steps = DebugHook()

erictransformer/utils/__init__.py

@@ -0,0 +1,21 @@
+from erictransformer.utils.init import get_model_components, et_retrieve_tokenizer
+from erictransformer.utils.init.get_device import et_get_device
+from erictransformer.utils.init.get_logger import et_get_logger
+from erictransformer.utils.timer import EricTimer
+from erictransformer.utils.tok_data import (
+    get_procs,
+    prepare_output_locations,
+    resolve_input_files,
+    tok_dir_to_dataset,
+    write_details_file,
+)
+from erictransformer.utils.tok_data.save_tok_data import save_json_tok_data
+from erictransformer.utils.train import (
+    create_tracker_dir,
+    get_num_training_steps,
+    get_optim,
+    get_precision,
+    get_tok_data,
+    make_dir,
+    resume_training,
+)

erictransformer/utils/init/__init__.py

@@ -0,0 +1,5 @@
+from erictransformer.utils.init.get_components import (
+    get_model_components,
+    get_model_components_tc,
+    et_retrieve_tokenizer,
+)

erictransformer/utils/init/get_components.py

@@ -0,0 +1,204 @@
+from typing import List, Optional, Union
+
+import torch
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+)
+
+from erictransformer.exceptions import EricLoadModelError, EricLoadTokenizerError
+
+
+def _get_torch_dtype(torch_dtype: str) -> Union[torch.dtype, str]:
+    if torch_dtype == "fp32":
+        return torch.float32
+    elif torch_dtype == "fp16":
+        return torch.float16
+    elif torch_dtype == "bf16":
+        return torch.bfloat16
+    else:
+        raise ValueError(
+            f"Invalid torch_dtype {torch_dtype}. Provide one of auto, fp32, fp16, or bf16."
+        )
+
+
+def get_model(
+    model_name_path: Union[str, PreTrainedModel],
+    model_class: AutoModel,
+    trust_remote_code: bool,
+    precision,
+) -> PreTrainedModel:
+    try:
+        if isinstance(model_name_path, PreTrainedModel):
+            return model_name_path
+
+        loaded_config = AutoConfig.from_pretrained(
+            model_name_path, trust_remote_code=trust_remote_code
+        )
+        torch_dtype = _get_torch_dtype(precision)
+
+        model = model_class.from_pretrained(
+            model_name_path,
+            config=loaded_config,
+            trust_remote_code=trust_remote_code,
+            dtype=torch_dtype,
+        )
+
+        return model
+
+    except Exception as e:
+        raise EricLoadModelError(f"Failed to load model from '{model_name_path}': {e}")
+
+
+def et_retrieve_tokenizer(
+    tokenizer_path: str, trust_remote_code: bool
+):
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_path,
+            trust_remote_code=trust_remote_code,
+        )
+
+        return tokenizer
+    except Exception as e:
+        raise EricLoadTokenizerError(
+            f"Failed to load tokenizer from '{tokenizer_path}': {e}"
+        )
+
+
+def get_tokenizer(
+    model_name_path: str,
+    tokenizer_path: Union[str, PreTrainedTokenizerBase],
+    trust_remote_code: bool,
+) -> PreTrainedTokenizerBase:
+    if (
+        tokenizer_path is None
+    ):  # by default we use the same value provided to model_name_path
+        tokenizer = et_retrieve_tokenizer(
+            tokenizer_path=model_name_path,
+            trust_remote_code=trust_remote_code,
+        )
+
+    elif isinstance(tokenizer_path, str):
+        tokenizer = et_retrieve_tokenizer(
+            tokenizer_path=tokenizer_path,
+            trust_remote_code=trust_remote_code,
+        )
+
+    elif isinstance(tokenizer_path, PreTrainedTokenizerBase):
+        # tokenizer provided directly
+        tokenizer = tokenizer_path
+        # maybe remove the following code since if an advanced user provides their own tokenizer they might want full control
+        if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+    else:
+        raise ValueError(
+            "Invalid tokenizer_path. Provide a None, a string or AutoTokenizer"
+        )
+
+    if tokenizer.pad_token is None:
+        if tokenizer.eos_token is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+        else:
+            tokenizer.add_special_tokens({"pad_token": "<pad>"})
+
+    return tokenizer
+
+
+def get_model_components(
+    model_name_path: Union[str, PreTrainedModel, None],
+    trust_remote_code: bool,
+    model_class: AutoModel,
+    tokenizer_path: Union[str, PreTrainedTokenizerBase],
+    precision: str,
+):
+    tokenizer = get_tokenizer(
+        model_name_path=model_name_path,
+        tokenizer_path=tokenizer_path,
+        trust_remote_code=trust_remote_code,
+    )
+
+    if model_name_path is not None:
+        model = get_model(
+            model_name_path=model_name_path,
+            model_class=model_class,
+            trust_remote_code=trust_remote_code,
+            precision=precision,
+        )
+        config = model.config
+    else:
+        model = None
+        config = None
+
+    return config, tokenizer, model
+
+
+def get_model_components_tc(
+    model_name_path: Union[str, PreTrainedModel, None],
+    trust_remote_code: bool,
+    model_class: AutoModel,
+    tokenizer_path: Union[str, PreTrainedTokenizerBase],
+    precision: str,
+    labels: Optional[List[str]] = None,
+):
+
+    config = AutoConfig.from_pretrained(
+        model_name_path, trust_remote_code=trust_remote_code
+    )
+    reset_labels =  False  # set to True if we should reset the labels to the size of num_labels
+
+    config_id2label = getattr(config, "id2label", None)
+    config_label2id = getattr(config, "label2id", None)
+
+
+    if not config_id2label or not config_label2id:
+        reset_labels = True
+    elif labels is not None:
+        reset_labels = True
+    else:
+        id2label_length = len(config_id2label)
+        label2id_length = len(config_label2id)
+        if id2label_length != label2id_length:
+            reset_labels = True
+
+
+    if reset_labels:
+        if labels is None:
+            labels = ["LABEL_0", "LABEL_1"]
+
+        if labels:
+            config.id2label = {i: labels[i] for i in range(len(labels))}
+        else:
+            num_labels = 2
+            config.id2label = {i: f"LABEL_{i}" for i in range(num_labels)}
+
+        config.label2id = {v: k for k, v in config.id2label.items()}
+
+    config.num_labels = len(config.id2label)
+
+    torch_dtype = _get_torch_dtype(precision)
+
+    model = model_class.from_pretrained(
+        model_name_path,
+        config=config,
+        dtype=torch_dtype,
+        trust_remote_code=trust_remote_code,
+        ignore_mismatched_sizes=True,
+    )
+
+    tokenizer = get_tokenizer(
+        model_name_path=model_name_path,
+        tokenizer_path=tokenizer_path,
+        trust_remote_code=trust_remote_code,
+    )
+
+    if tokenizer.pad_token is None:
+        if tokenizer.eos_token is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+        else:
+            tokenizer.add_special_tokens({"pad_token": "<pad>"})
+
+    return config, tokenizer, model

erictransformer/utils/init/get_device.py

@@ -0,0 +1,22 @@
+import torch
+from torch import device
+
+from erictransformer.exceptions import EricDeviceError
+
+
+def et_get_device() -> device:
+    try:
+        d = None
+        if torch.backends.mps.is_available():
+            if torch.backends.mps.is_built():
+                d = torch.device("mps")
+
+        if torch.cuda.is_available():
+            d = torch.device("cuda:0")
+
+        if not d:
+            d = torch.device("cpu")
+
+        return d
+    except Exception as e:
+        raise EricDeviceError(f"Device selection failed: {e}")

erictransformer/utils/init/get_logger.py

@@ -0,0 +1,15 @@
+import logging
+from logging import Logger
+
+
+def et_get_logger() -> Logger:
+    logger = logging.getLogger("erictransformer")
+    handler = logging.StreamHandler()
+    handler.addFilter(logging.Filter("erictransformer"))
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+        handlers=[handler],
+    )
+    return logger