tiny-random-LlamaForCausalLM

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

out = "ccmodular/tiny-random-LlamaForCausalLM"
template = "meta-llama/Llama-3.1-8B-Instruct"

print("Loading original tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(template)

# Shrink the tokenizer to a smaller vocab.
print("Shrinking tokenizer...")
tokenizer = tokenizer.train_new_from_iterator(
    text_iterator=[],
    # Effective minimum.
    # BPE requires 256 + special tokens, rounded up to the next power of 2 => 512.
    vocab_size=512,
)

# Minimize the chat template (remove tools and other context).
tokenizer.chat_template = (
    "{{- bos_token }}"
    "{%- for message in messages %}"
    "{{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\\n\\n' + message['content'] | trim + '<|eot_id|>' }}"
    "{%- endfor %}"
    "{%- if add_generation_prompt %}"
    "{{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }}"
    "{%- endif %}"
)

# Configure the model.
config = AutoConfig.from_pretrained(template)
config.num_attention_heads = 1
config.head_dim = 2  # minimum required by RoPE in MAX
config.hidden_size = config.num_attention_heads * config.head_dim
config.num_hidden_layers = 1
config.intermediate_size = 1
config.max_position_embeddings = 32
config.num_key_value_heads = 1
config.vocab_size = len(tokenizer)

# Create the model.
model = AutoModelForCausalLM.from_config(config)

# Print some stats.
print(f"Model parameters: {model.num_parameters():,}")
print(f"Vocab size: {config.vocab_size}")

# Push to HF hub (requires auth).
print(f"Pushing to {out}...")
model.push_to_hub(out, private=False)
tokenizer.push_to_hub(out, private=False)
config.push_to_hub(out, private=False)
print(f"Successfully pushed minimal model to {out}")