erictransformer 0.0.1__py3-none-any.whl

erictransformer/utils/load_from_repo_or_path.py

@@ -0,0 +1,14 @@
+import pathlib
+
+import huggingface_hub
+
+
+def load_from_repo_or_path(path: str, path_in_repo: str) -> str:
+    # Resolve a path that is either local or a repo into a local path.
+    if pathlib.Path(path).exists():
+        return path
+    else:
+        local_repo_dir = huggingface_hub.snapshot_download(
+            repo_id=path, repo_type="dataset", token=True
+        )
+        return str(pathlib.Path(local_repo_dir) / path_in_repo)

erictransformer/utils/test/__init__.py

@@ -0,0 +1 @@
+from .debug_hook import DebugHook

erictransformer/utils/test/debug_hook.py

@@ -0,0 +1,20 @@
+import contextlib
+import typing
+
+class DebugHook:
+    hook: typing.Callable | None
+
+    def __init__(self) -> None:
+        self.hook = None
+
+    def __call__(self, *a, **k) -> typing.Any:
+        if self.hook:
+            return self.hook(*a,**k)
+
+    @contextlib.contextmanager
+    def set(self, hook):
+        self.hook = hook
+        try:
+            yield
+        finally:
+            self.hook = None

erictransformer/utils/timer/__init__.py

@@ -0,0 +1 @@
+from erictransformer.utils.timer.eric_timer import EricTimer

erictransformer/utils/timer/eric_timer.py

@@ -0,0 +1,145 @@
+import csv
+import os
+import time
+from contextlib import contextmanager
+
+
+class EricTimer:
+    def __init__(self, out_dir: str, enabled: bool = True):
+        self.enabled = enabled  # if false self.report() does nothing.
+        self._stats = {}
+        self._active = {}
+        self._out_dir = out_dir
+        self._category_dir = os.path.join(out_dir, "categories")
+        self._ensure_dirs()
+        self.start_time = time.perf_counter()
+
+    @contextmanager
+    def section(self, category: str, label: str):
+        self.start(category, label)
+        try:
+            yield
+        finally:
+            self.stop(category, label)
+
+    def start(self, category: str, label: str):
+        if self._active:
+            raise ValueError(f"Eric Timer is already running! with {self._active}")
+        if label == "misc_time":
+            raise ValueError("misc_time is an invalid label")
+
+        if category not in self._active:
+            self._active[category] = {}
+
+        if label in self._active[category]:
+            raise ValueError(f"{category!r}/{label!r} is active. Call stop() first.")
+
+        self._active[category][label] = time.perf_counter()
+
+    def stop(self, category: str, label: str):
+        if category not in self._active or label not in self._active[category]:
+            raise RuntimeError(
+                f"stop() called without calling start for {category!r}/{label!r} first."
+            )
+
+        t0 = self._active[category][label]
+        del self._active[category][label]
+        if not self._active[category]:
+            del self._active[category]
+
+        dt = time.perf_counter() - t0
+
+        if category not in self._stats:
+            self._stats[category] = {}
+
+        if label not in self._stats[category]:
+            self._stats[category][label] = {"total": 0.0, "count": 0}
+
+        self._stats[category][label]["total"] += dt
+        self._stats[category][label]["count"] += 1
+
+        return dt
+
+    def _ensure_dirs(self):
+        if self.enabled:
+            if self._out_dir:
+                os.makedirs(self._out_dir, exist_ok=True)
+            if self._category_dir:
+                os.makedirs(self._category_dir, exist_ok=True)
+
+    def _category_path(self, category: str):
+        fname = f"{category}.csv"
+        return os.path.join(self._category_dir, fname)
+
+    def _write_category_json(self, category: str):
+        if category in self._stats:
+            labels = self._stats[category]
+        else:
+            labels = {}
+
+        rows = []
+        for label, s in labels.items():
+            total = float(s["total"])
+            count = int(s["count"])
+            mean = (total / count) if count else 0.0
+            rows.append(
+                {
+                    "label": label,
+                    "total": round(total, 4),
+                    "mean": round(mean, 4),
+                    "count": count,
+                }
+            )
+        rows.sort(key=lambda r: r["total"], reverse=True)
+
+        with open(
+            self._category_path(category), "w", newline="", encoding="utf-8"
+        ) as f:
+            writer = csv.DictWriter(f, fieldnames=["label", "total", "mean", "count"])
+            writer.writeheader()
+            for r in rows:
+                writer.writerow(r)
+
+    def _write_categories_summary_json(self):
+
+        if self.enabled:
+            path = os.path.join(self._out_dir, "summary.csv")
+            rows = []
+            for category, labels in self._stats.items():
+                total = 0.0
+                for s in labels.values():
+                    total += float(s["total"])
+                rows.append({"category": category, "total": round(total, 4)})
+            rows.sort(key=lambda r: r["total"], reverse=True)
+
+            with open(path, "w", newline="", encoding="utf-8") as f:
+                writer = csv.DictWriter(f, fieldnames=["category", "total"])
+                writer.writeheader()
+                for r in rows:
+                    writer.writerow(r)
+
+    def report(self):
+        if self.enabled:
+            accounted_time = 0.0
+            for labels in self._stats.values():
+                for s in labels.values():
+                    accounted_time += float(s["total"])
+
+            misc_time = time.perf_counter() - self.start_time - accounted_time
+            if misc_time < 0:
+                misc_time = 0.0
+
+            misc_cat = "misc"
+            if misc_cat not in self._stats:
+                self._stats[misc_cat] = {}
+            self._stats[misc_cat]["misc_time"] = {"total": misc_time, "count": 1}
+
+            for category in list(self._stats.keys()):
+                self._write_category_json(category)
+
+            self._write_categories_summary_json()
+
+    def reset(self):
+        self._stats.clear()
+        self._active.clear()
+        self.start_time = time.perf_counter()

erictransformer/utils/tok_data/__init__.py

@@ -0,0 +1,8 @@
+from erictransformer.utils.tok_data.num_proc import get_procs
+from erictransformer.utils.tok_data.save_tok_data import save_json_tok_data
+from erictransformer.utils.tok_data.tok_data_to_dataset import tok_dir_to_dataset
+from erictransformer.utils.tok_data.tok_helpers import (
+    prepare_output_locations,
+    resolve_input_files,
+    write_details_file,
+)

erictransformer/utils/tok_data/num_proc.py

@@ -0,0 +1,15 @@
+import math
+import os
+from typing import Optional
+
+
+def get_procs(procs: int) -> Optional[int]:
+    if procs == -1:
+        return None
+    elif procs == 0:
+        cpu_count = os.cpu_count()
+        if cpu_count is None or cpu_count <= 1:
+            return 1
+        return math.floor(cpu_count / 2)
+
+    return procs

erictransformer/utils/tok_data/save_tok_data.py

@@ -0,0 +1,36 @@
+import json
+
+from erictransformer.exceptions import EricDatasetError, EricIOError
+
+
+def save_json_tok_data(tok_dataset, output_files, shards, formatter):
+    total = len(tok_dataset)
+    per_shard = total // shards
+    remainder = total % shards
+
+    index = 0
+
+    for shard_id, output_file in enumerate(output_files):
+        count = per_shard + (1 if shard_id < remainder else 0)
+
+        for _ in range(count):
+            if index >= total:
+                break
+            try:
+                raw = formatter(tok_dataset[index])
+            except Exception as e:
+                raise EricDatasetError(
+                    f"Failed to format dataset example at index {index}: {e}"
+                )
+
+            try:
+                example = {
+                    k: v.tolist() if hasattr(v, "tolist") else v for k, v in raw.items()
+                }
+                output_file.write(json.dumps(example) + "\n")
+            except Exception as e:
+                raise EricIOError(
+                    f"Failed to write formatted example at index {index} to file: {e}"
+                )
+
+            index += 1

erictransformer/utils/tok_data/tok_data_to_dataset.py

@@ -0,0 +1,48 @@
+import json
+import os
+
+from datasets import Dataset
+
+from erictransformer.exceptions import EricDatasetError, EricInputError, EricIOError
+
+
+def tok_dir_to_dataset(tok_dir: str):
+    # Validate the input directory
+    if not os.path.isdir(tok_dir):
+        raise EricInputError(f"Invalid tok_dir: {tok_dir} is not a directory.")
+
+    # Attempt to construct the path to eric_details.json
+    try:
+        eric_details_path = os.path.join(tok_dir, "eric_details.json")
+    except Exception as e:
+        raise EricInputError(f"Error creating path to eric_details.json: {e}")
+
+    # Read the eric_details.json file
+    try:
+        with open(eric_details_path, "r") as f:
+            eric_details = json.load(f)
+    except Exception as e:
+        raise EricIOError(f"Unable to read '{eric_details_path}': {e}")
+
+    # List of full paths to all JSONL files
+    full_paths = eric_details.get("paths", [])
+    if not full_paths:
+        raise EricDatasetError("No file paths found in eric_details.json.")
+
+    # Load all entries from each JSONL file
+    all_data = []
+    for path in full_paths:
+        try:
+            with open(path, "r") as f:
+                for line in f:
+                    all_data.append(json.loads(line))
+        except Exception as e:
+            raise EricIOError(f"Unable to read '{path}': {e}")
+
+    # Convert list of dicts to Hugging Face Dataset
+    try:
+        dataset = Dataset.from_list(all_data)
+    except Exception as e:
+        raise EricDatasetError(f"Unable to create dataset: {e}")
+
+    return dataset

erictransformer/utils/tok_data/tok_helpers.py

@@ -0,0 +1,79 @@
+import glob
+import json
+import os
+from datetime import datetime
+from typing import List, Tuple, Union
+
+from datasets import Dataset
+
+from erictransformer.exceptions import EricInputError, EricIOError
+
+EXT_TO_TYPE_MAP = {
+    ".json": "json",
+    ".jsonl": "json",  # jsonl is loaded as json
+}
+
+
+def resolve_input_files(input_data: Union[str, Dataset]) -> List[tuple]:
+    try:
+        if isinstance(input_data, str):
+            if os.path.isdir(input_data):
+                return collect_files_from_dir(input_data)
+            elif os.path.isfile(input_data):
+                ext = os.path.splitext(input_data)[1].lower()
+                if ext not in EXT_TO_TYPE_MAP:
+                    raise EricInputError(f"Unsupported file extension: {ext}")
+                return [(input_data, EXT_TO_TYPE_MAP[ext])]
+            else:
+                raise EricInputError(f"Path does not exist: {input_data}")
+        else:
+            raise EricInputError("Input must be a file path or directory string.")
+    except Exception as e:
+        raise EricInputError(f"Failed to resolve input files: {e}")
+
+
+def collect_files_from_dir(dir_path: str) -> List[Tuple[str, str]]:
+    try:
+        jsonl = [(f, "json") for f in glob.glob(os.path.join(dir_path, "*.jsonl"))]
+        # text = [(f, "text") for f in glob.glob(os.path.join(dir_path, "*.txt"))]
+        # we might support .txt files for EricGeneration models one day
+        all_files = jsonl
+        if not all_files:
+            raise EricInputError(
+                f"No valid jsonl files found in directory: {dir_path}"
+            )
+        return all_files
+    except Exception as e:
+        raise EricIOError(f"Failed to collect files from directory '{dir_path}': {e}")
+
+
+def prepare_output_locations(out_dir: str, shards: int):
+    try:
+        os.makedirs(out_dir, exist_ok=True)
+        paths = [
+            os.path.join(out_dir, f"tok_shard_{i + 1}.jsonl")
+            for i in range(shards)
+        ]
+        handles = []
+        for p in paths:
+            try:
+                handles.append(open(p, "w"))
+            except OSError as e:
+                raise EricIOError(f"Failed to open file for writing: {p} - {e}")
+        return paths, handles
+    except Exception as e:
+        raise EricIOError(f"Failed to prepare output locations in '{out_dir}': {e}")
+
+
+def write_details_file(out_dir: str, num_cases: int, output_paths: List[str]):
+    try:
+        detail = {
+            "num_cases": num_cases,
+            "timestamp": datetime.now().isoformat(),
+            "paths": output_paths,
+        }
+        details_path = os.path.join(out_dir, "eric_details.json")
+        with open(details_path, "w") as f:
+            json.dump(detail, f, indent=2)
+    except Exception as e:
+        raise EricIOError(f"Failed to write details file in '{out_dir}': {e}")

erictransformer/utils/train/__init__.py

@@ -0,0 +1,6 @@
+from erictransformer.utils.train.confirm_optimizer import get_optim
+from erictransformer.utils.train.create_dir import create_tracker_dir, make_dir
+from erictransformer.utils.train.get_num_training_steps import get_num_training_steps
+from erictransformer.utils.train.get_precision import get_precision
+from erictransformer.utils.train.get_tok_data import get_tok_data
+from erictransformer.utils.train.resume import resume_training

erictransformer/utils/train/confirm_optimizer.py

@@ -0,0 +1,18 @@
+from torch.optim import SGD, AdamW
+
+from erictransformer.args import EricTrainArgs
+
+
+def get_optim(args: EricTrainArgs, model, logger):
+    if args.optim == "adamw":
+        logger.debug("Using PyTorch's AdamW optimizer.")
+        return AdamW(model.parameters(), lr=args.lr, weight_decay=0.05)
+
+    elif args.optim == "sgd":
+        logger.debug("Using PyTorch's SGD optimizer")
+        return SGD(
+            model.parameters(), lr=args.lr, momentum=0, weight_decay=1e-4
+        )
+
+    else:
+        raise ValueError(f"Unsupported optimizer: {args.optim}")

erictransformer/utils/train/create_dir.py

@@ -0,0 +1,72 @@
+import datetime as dt
+import logging
+import os
+import re
+from pathlib import Path
+
+from erictransformer.exceptions import EricIOError
+
+
+def _slugify(name: str) -> str:
+    safe_chars = re.compile(r"[^A-Za-z0-9._-]+")
+    s = name.strip()
+    s = s.replace(" ", "-")
+    s = safe_chars.sub("-", s)
+    s = re.sub(r"-{2,}", "-", s)
+    s = s.strip("-._")
+    if not s:
+        raise EricIOError(f"Invalid directory name: {name}")
+    return s
+
+
+def _increment_path(base_path: Path) -> Path:
+    parent = base_path.parent
+    name = base_path.name
+
+    m = re.match(r"^(.*?)(?:-(\d+))?$", name)
+    root = m.group(1)
+    pat = re.compile(rf"^{re.escape(root)}(?:-(\d+))?$")
+
+    max_i = -1
+    if parent.exists():
+        for child in parent.iterdir():
+            m2 = pat.match(child.name)
+            if m2:
+                i = int(m2.group(1)) if m2 and m2.group(1) else 0
+                if i > max_i:
+                    max_i = i
+
+    if max_i < 0:
+        return base_path
+    next_i = max_i + 1
+    if next_i == 0:
+        return root
+    else:
+        return parent / f"{name}-{next_i}"
+
+
+def make_dir(name):
+    try:
+        os.makedirs(name, exist_ok=True)
+        logging.info(f"Directory {name} created successfully.")
+    except FileExistsError:
+        logging.warning(f"Directory {name} already exists.")
+
+
+def create_tracker_dir(out_dir: str, label: str, run_name: str):
+    timestamp = dt.datetime.now().strftime("%Y%m%d_%H%M%S")
+    if run_name:
+        slugify_run_name = _slugify(run_name)
+
+        output = Path(out_dir)
+        output.mkdir(parents=True, exist_ok=True)
+
+        dir_name = _increment_path(output / slugify_run_name)
+
+        make_dir(dir_name)
+
+    else:
+        dir_name = os.path.join(out_dir, f"{label}_{timestamp}")
+        make_dir(dir_name)
+
+    return dir_name

erictransformer/utils/train/get_num_training_steps.py

@@ -0,0 +1,15 @@
+import math
+
+
+def get_num_training_steps(
+    train_cases: int,
+    num_devices: int,
+    epochs: float,
+    gas: int,
+    bs: int,
+) -> int:
+    steps = (train_cases * epochs) // (
+        num_devices * gas * bs
+    )
+
+    return math.ceil(steps)

erictransformer/utils/train/get_precision.py

@@ -0,0 +1,22 @@
+import torch
+
+
+def get_precision(device) -> str:
+    is_cuda = device.type == "cuda"
+    bf16_ok = (
+        bool(getattr(torch.cuda, "is_bf16_supported", lambda: False)())
+        if is_cuda
+        else False
+    )
+
+    use_bf16 = is_cuda and bf16_ok
+    use_fp16 = is_cuda and not use_bf16
+
+    if use_bf16:
+        precision_type = "bf16"
+    elif use_fp16:
+        precision_type = "fp16"
+    else:
+        precision_type = "fp32"
+
+    return precision_type

erictransformer/utils/train/get_tok_data.py

@@ -0,0 +1,105 @@
+import json
+from pathlib import Path
+from typing import List
+
+from datasets import load_dataset
+from torch import device
+from torch.utils.data import DataLoader
+from transformers import DataCollator
+
+from erictransformer.exceptions import EricDatasetError, EricInputError
+
+
+def get_tok_data(
+    tok_dir: str,
+    seed: int,
+    bs: int,
+    collate_fn: DataCollator,
+    device_type: device,
+):
+    details = _load_details_file(tok_dir)
+    paths = _resolve_dataset_paths(details, tok_dir)
+    dataset = _load_streaming_dataset(paths, seed)
+    dataloader = _build_dataloader(dataset, bs, collate_fn, device_type)
+
+    return dataloader, details["num_cases"]
+
+
+def _load_details_file(tok_dir: str) -> dict:
+    details_filepath = Path(tok_dir) / "eric_details.json"
+    if not details_filepath.is_file():
+        raise FileNotFoundError(f"Missing required file: {details_filepath}")
+    try:
+        with open(details_filepath, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except Exception as e:
+        raise EricInputError(f"Failed to parse eric_details.json: {e}")
+
+
+def _resolve_dataset_paths(details: dict, tok_dir: str) -> List[Path]:
+    try:
+        paths = [Path(p) for p in details["paths"]]
+    except KeyError:
+        raise EricInputError(
+            f"'paths' field missing in eric_details.json in {tok_dir}"
+        )
+
+    if not all(p.is_file() and p.suffix == ".jsonl" for p in paths):
+        raise EricInputError("All dataset paths must exist and be `.jsonl` files")
+    return paths
+
+
+def _load_streaming_dataset(paths: List[Path], seed: int):
+    data_files = {"train": [str(p) for p in paths]}
+    try:
+        ds = load_dataset(
+            "json",
+            data_files=data_files,
+            split="train",
+            streaming=True,
+        )
+
+        # Probe for at least one example (to handle the "all shards empty" case)
+        it = iter(ds)
+        try:
+            _ = next(it)
+        except StopIteration:
+            raise EricInputError("No non-empty datasets to stream")
+
+        # Re-create the iterator since we advanced one element when probing
+        ds = load_dataset("json", data_files=data_files, split="train", streaming=True)
+
+        ds = ds.shuffle(buffer_size=10_000, seed=seed)
+        return ds
+    except EricInputError:
+        raise
+    except Exception as e:
+        raise EricDatasetError(
+            f"Failed to load streaming dataset from paths {paths}: {e}"
+        )
+
+
+def _build_dataloader(
+    dataset, bs: int, collate_fn: DataCollator, device_type: device
+):
+    try:
+        shards = getattr(dataset, "n_shards", 1)
+        base_workers = min(4, shards)
+
+        if device_type.type == "mps":
+            workers = 0
+        else:
+            workers = base_workers
+
+        return DataLoader(
+            dataset,
+            batch_size=bs,
+            collate_fn=collate_fn,
+            num_workers=workers,
+            pin_memory=(device_type.type == "cuda"),
+            persistent_workers=(workers > 0),
+            prefetch_factor=4 if workers > 0 else None,
+        )
+
+    except Exception as e:
+        raise EricInputError(f"Failed to create DataLoader: {e}")

erictransformer/utils/train/resume.py

@@ -0,0 +1,62 @@
+import json
+import os
+from pathlib import Path
+
+from erictransformer.exceptions import EricResumeError
+
+
+def resume_training(resume_path: str):
+    tracker_state = _load_tracker_state(resume_path)
+    args_dict = _load_train_args(resume_path)
+    model_tokenizer_path, lr_sched_path = _resolve_training_paths(
+        resume_path
+    )
+
+    return (
+        tracker_state,
+        args_dict,
+        model_tokenizer_path,
+        lr_sched_path,
+    )
+
+
+def _load_tracker_state(resume_path: str):
+    logs_path = Path(os.path.join(resume_path, "eric_tracker_state.json"))
+    try:
+        if not logs_path.exists():
+            raise EricResumeError(f"No log file found at {logs_path}.")
+
+        with logs_path.open(encoding="utf-8") as f:
+            tracker_state = json.load(f)
+
+        if tracker_state is None:
+            raise EricResumeError("Loaded tracker state is None.")
+
+        return tracker_state
+    except Exception as e:
+        raise EricResumeError(f"Failed to load tracker state: {e}")
+
+
+def _load_train_args(resume_path: str):
+    train_args_path = Path(os.path.join(resume_path, "train_args.json"))
+    try:
+        if not train_args_path.exists():
+            raise EricResumeError(f"No train_args.json found at {train_args_path}.")
+
+        args_dict = json.loads(train_args_path.read_text(encoding="utf-8"))
+        return args_dict
+    except Exception as e:
+        raise EricResumeError(f"Failed to load training arguments: {e}")
+
+
+def _resolve_training_paths(resume_path: str):
+    try:
+        lr_sched_path = os.path.join(resume_path, "lr_sched.pt")
+
+        if not Path(lr_sched_path).exists():
+            raise EricResumeError(f"lr_sched.pt not found at {lr_sched_path}")
+
+
+        return resume_path, lr_sched_path
+    except Exception as e:
+        raise EricResumeError(f"Failed to resolve training paths: {e}")