amachine.am_transformers.am_gpt2_generator

  1import json
  2from pathlib import Path
  3from typing import Literal
  4
  5from .am_models import get_tokenizer, get_text_tokenizer
  6
  7try:
  8    from transformers import GPT2Config, AutoModelForCausalLM
  9    _HAVE_TRANSFORMERS = True
 10except ImportError:
 11    _HAVE_TRANSFORMERS = False
 12
 13def _gpt2_model_name(width: int, depth: int, head_dim: int, opt: str | None = None, id_str: str | None = None) -> str:
 14    """Generates a consistent naming convention for the GPT-2 models."""
 15    parts = [f"gpt2-w{width}-d{depth}-h{head_dim}"]
 16    if opt is not None:
 17        parts.append(f"o{opt}")
 18    if id_str is not None:
 19        parts.append(f"id{id_str}")
 20    return "-".join(parts)
 21
 22def generate_gpt2_ensemble(
 23    output_dir: str,
 24    core_params: list[dict],
 25    seq_len: int = 1024,
 26    tokenizer_type: Literal["symbolic", "text"] = "symbolic",
 27    tokenizer_id : str = "ibm-granite/granite-4.0-350m",
 28    activation_function : str = "gelu_new",
 29    dropout : float = 0.1
 30) -> dict:
 31    """Generate a sweep of standard GPT-2 models with varying widths and depths."""
 32
 33    if not _HAVE_TRANSFORMERS:
 34        raise ImportError("The 'generate_gpt2_ensemble' function requires 'transformers'.")
 35
 36    output_path = Path(output_dir)
 37    output_path.mkdir(parents=True, exist_ok=True)
 38
 39    manifest = {"models": []}
 40
 41    # 1. Setup the Tokenizer dynamically just like the Granite builder
 42    if tokenizer_type == "symbolic":
 43        raw_tokenizer, vocabulary, hf_tokenizer, gen_config = get_tokenizer()
 44    elif tokenizer_type == "text":
 45        raw_tokenizer, vocabulary, hf_tokenizer, gen_config = get_text_tokenizer( tokenizer_id )
 46    else:
 47        raise ValueError(f"Unsupported tokenizer_type {tokenizer_type}. Options are symbolic or text")
 48
 49    # assert isinstance(hf_tokenizer.pad_token_id, int), "pad_token_id is not an integer!"
 50    # assert isinstance(hf_tokenizer.bos_token_id, int), "bos_token_id is not an integer!"
 51    # assert isinstance(hf_tokenizer.eos_token_id, int), "eos_token_id is not an integer!"
 52
 53    vocab_size = len(hf_tokenizer)
 54    pad_token_id = hf_tokenizer.pad_token_id
 55    bos_token_id = hf_tokenizer.bos_token_id
 56    eos_token_id = hf_tokenizer.eos_token_id
 57
 58    for idx, p in enumerate(core_params):
 59
 60        if not "width" in p or not "depth" in p or not "head_dim" in p:
 61            raise ValueError("core params must contain width, depth, and head_dim")
 62
 63        width = int(p["width"])
 64        depth = int(p["depth"])
 65        head_dim = int(p["head_dim"])
 66
 67        opt = p.get("opt", None)
 68        id_str = p.get("id", str(idx))
 69        target_seq_len = int(p.get("seq_len", seq_len))
 70
 71        if width % head_dim != 0:
 72            raise ValueError(f"Width ({width}) must be divisible by head_dim ({head_dim})")
 73        
 74        n_head = width // head_dim
 75        
 76        name = _gpt2_model_name(width=width, depth=depth, head_dim=head_dim, opt=opt, id_str=id_str)
 77        model_dir = str(output_path / name)
 78        model_path = Path(model_dir)
 79        model_path.mkdir(parents=True, exist_ok=True)
 80
 81        print(f"\n{'=' * 60}")
 82        print(f"Generating {name}")
 83        print(f"  Architecture: GPT-2 | width (n_embd)={width} | depth (n_layer)={depth} | heads={n_head}")
 84
 85        config = GPT2Config(
 86            vocab_size=vocab_size,
 87            n_positions=target_seq_len+1, # +1 for safety, huggingface shifts internally, -> 1025 model inputs/labels to predict 1024
 88            n_embd=width,
 89            n_layer=depth,
 90            n_head=n_head,
 91            bos_token_id=bos_token_id,
 92            eos_token_id=eos_token_id,
 93            pad_token_id=pad_token_id,
 94            activation_function=activation_function,
 95            resid_pdrop=dropout,
 96            embd_pdrop=dropout,
 97            attn_pdrop=dropout,
 98            use_cache=True
 99        )
100
101        model = AutoModelForCausalLM.from_config(config)
102
103        total_params = sum(param.numel() for param in model.parameters())
104        print(f"  Total parameters: {total_params / 1e6:.2f}M")
105
106        model.save_pretrained(model_dir, safe_serialization=True)
107
108        raw_tokenizer.save(str(model_path / "tokenizer.json"))
109        
110        if tokenizer_type == "symbolic":
111            with open(model_path / "vocab.json", "w", encoding="utf-8") as f:
112                json.dump(vocabulary, f, indent=2, ensure_ascii=False)
113
114        hf_tokenizer.save_pretrained(model_dir)
115        gen_config.save_pretrained(model_dir)
116
117        manifest["models"].append({
118            "name": name,
119            "architecture": "gpt2",
120            "hidden_size": width,
121            "num_hidden_layers": depth,
122            "head_dim": head_dim,
123            "num_attention_heads": n_head,
124            "total_params": total_params,
125            "seq_len": target_seq_len,
126            "tokenizer_type": tokenizer_type,
127            "estimated_params" : total_params
128        })
129
130    manifest_path = output_path / "manifest_gpt2.json"
131    with open(manifest_path, "w") as f:
132        json.dump(manifest, f, indent=2)
133    
134    print(f"\nGPT-2 Manifest written to {manifest_path}")
135
136    return manifest