Views
No views yet
1"""Create a tiny random Gemma 3 model (5 layers) and upload to HuggingFace.
2
3One-off script. The model is text-only with random weights, intended for
4fast smoke tests with tunix/JAX
5
6Usage: uv run python scripts/create_tiny_gemma3.py
7"""
8
9import torch
10from transformers import (
11 AutoTokenizer,
12 GenerationConfig,
13 set_seed,
14)
15from transformers.models.gemma3 import Gemma3ForCausalLM, Gemma3TextConfig
16from huggingface_hub import HfApi
17
18source_model_id = "google/gemma-3-1b-it"
19repo_id = "wassname/gemma3-5lyr-tiny-random"
20save_folder = "/tmp/tiny-random/gemma3-5lyr"
21
22# Tokenizer from source (same vocab)
23tokenizer = AutoTokenizer.from_pretrained(source_model_id)
24tokenizer.save_pretrained(save_folder)
25
26# Tiny text-only config matching tunix ModelConfig in model.py
27config = Gemma3TextConfig(
28 vocab_size=262144,
29 hidden_size=64,
30 intermediate_size=128,
31 num_hidden_layers=5,
32 num_attention_heads=2,
33 head_dim=32,
34 num_key_value_heads=1,
35 sliding_window=512,
36 tie_word_embeddings=True,
37)
38config._name_or_path = source_model_id
39
40model = Gemma3ForCausalLM(config).to(torch.bfloat16)
41
42# Random init
43set_seed(42)
44with torch.no_grad():
45 for name, p in sorted(model.named_parameters()):
46 torch.nn.init.normal_(p, 0, 0.5)
47 print(name, p.shape)
48
49model.generation_config = GenerationConfig.from_pretrained(source_model_id)
50model.save_pretrained(save_folder)
51
52# Upload
53api = HfApi()
54api.create_repo(repo_id, exist_ok=True)
55api.upload_folder(folder_path=save_folder, repo_id=repo_id)
56print(f"Uploaded to https://huggingface.co/{repo_id}")
57