TinyWord-v2 is a revamped and retrained version of v1. In v1, we noticed that it didn't use weight-tying, which ate up half of its parameters. This was misleading as it was effectively the same size as MicroWord.
Anyway, this version achives much better performace compared to v1.
TinyWord-v2 was trained on 753,232 unique words (entries), 3,225,398 tokens, and 7,022,310 characters. ~660k of those words are English, while ~90k of them are Spanish.
TinyWord-v2 was trained on a NVIDA RTX 2060 6GB for 6 epochs with a batch size of 32.
1# =============================================================================
2# Inference
3# =============================================================================
4
5MODEL_DIR = "Harley-ml/TinyWord2-128k" # path
6TOKENIZER_PATH = "Harley-ml/TinyWord2-128k"
7
8# --- Generation settings ---
9PROMPT = "w" # prompt
10MAX_NEW_TOKENS = 32
11TEMPERATURE = 1.2
12TOP_P = 0.95
13TOP_K = 50
14REPETITION_PENALTY = 1.1
15DO_SAMPLE = True
16
17# =============================================================================
18
19import torch
20from pathlib import Path
21from transformers import (
22 AutoModelForCausalLM,
23 PreTrainedTokenizerFast,
24 AddedToken,
25)
26
27# ---------------------------------------------------------------------------
28# Device
29# ---------------------------------------------------------------------------
30
31device = (
32 "cuda" if torch.cuda.is_available() else
33 "mps" if torch.backends.mps.is_available() else
34 "cpu"
35)
36print(f"Device : {device}")
37
38# ---------------------------------------------------------------------------
39# Tokenizer (mirrors training setup)
40# ---------------------------------------------------------------------------
41
42def load_tokenizer(path: str):
43 p = Path(path).resolve()
44 if not p.exists():
45 raise FileNotFoundError(f"Tokenizer not found: {p}")
46 tok = PreTrainedTokenizerFast(tokenizer_file=str(p))
47 specials = {}
48 if tok.bos_token is None: specials["bos_token"] = AddedToken("<|bos|>", special=True)
49 if tok.eos_token is None: specials["eos_token"] = AddedToken("<|eos|>", special=True)
50 if tok.unk_token is None: specials["unk_token"] = AddedToken("<|unk|>", special=True)
51 if tok.pad_token is None:
52 if tok.eos_token is not None:
53 tok.pad_token = tok.eos_token
54 else:
55 specials["pad_token"] = AddedToken("<|pad|>", special=True)
56 if specials:
57 tok.add_special_tokens(specials)
58 tok.padding_side = "left" # left-pad for batched generation
59 return tok
60
61print("Loading tokenizer...")
62tokenizer = load_tokenizer(TOKENIZER_PATH)
63print(f" Vocab size : {tokenizer.vocab_size}")
64print(f" BOS : {tokenizer.bos_token!r}")
65print(f" EOS : {tokenizer.eos_token!r}")
66print(f" PAD : {tokenizer.pad_token!r} (id={tokenizer.pad_token_id})")
67
68# ---------------------------------------------------------------------------
69# Model
70# ---------------------------------------------------------------------------
71
72print(f"\nLoading model from {MODEL_DIR} ...")
73model = AutoModelForCausalLM.from_pretrained(
74 MODEL_DIR,
75 dtype=torch.float16 if device == "cuda" else torch.float32,
76 low_cpu_mem_usage=True,
77)
78model.eval()
79model.to(device)
80
81total_params = sum(p.numel() for p in model.parameters())
82print(f" Parameters : {total_params:,}")
83
84# ---------------------------------------------------------------------------
85# Generation helper
86# ---------------------------------------------------------------------------
87
88def generate(
89 prompt: str = PROMPT,
90 max_new_tokens: int = MAX_NEW_TOKENS,
91 temperature: float = TEMPERATURE,
92 top_p: float = TOP_P,
93 top_k: int = TOP_K,
94 repetition_penalty: float = REPETITION_PENALTY,
95 do_sample: bool = DO_SAMPLE,
96) -> str:
97
98 bos = tokenizer.bos_token or ""
99 full_prompt = bos + prompt
100
101 inputs = tokenizer(
102 full_prompt,
103 return_tensors="pt",
104 add_special_tokens=False,
105 ).to(device)
106 inputs.pop("token_type_ids", None) # Qwen3 doesn't use this
107
108 gen_kwargs = dict(
109 max_new_tokens = max_new_tokens,
110 do_sample = do_sample,
111 repetition_penalty = repetition_penalty,
112 eos_token_id = tokenizer.eos_token_id,
113 pad_token_id = tokenizer.pad_token_id,
114 )
115 if do_sample:
116 gen_kwargs["temperature"] = temperature
117 gen_kwargs["top_p"] = top_p
118 gen_kwargs["top_k"] = top_k
119
120 with torch.inference_mode():
121 output_ids = model.generate(**inputs, **gen_kwargs)
122
123 # Strip the prompt tokens so we only return what was generated
124 prompt_len = inputs["input_ids"].shape[-1]
125 new_ids = output_ids[0][prompt_len:]
126 return tokenizer.decode(new_ids, skip_special_tokens=True)
127
128
129# ---------------------------------------------------------------------------
130# Run
131# ---------------------------------------------------------------------------
132
133if __name__ == "__main__":
134 print(f"\nPrompt : {PROMPT!r}")
135 print("-" * 60)
136
137 output = generate(PROMPT)
138
139 print("Generated:")
140 print(output)
1@misc{tinyword2-128k,
2 title = {TinyWord-134k: A Test of Morphological Compression in TLMs},
3 author = {Harley-ml},
4 year = {2026},
5 url = {https://huggingface.co/Harley-ml/TinyWord2-128k}
6}
