Duplicate from SupraLabs/Supra-Mini-v2-0.1M

Co-authored-by: LH-Tech AI <[email protected]>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,119 @@
+---
+license: apache-2.0
+datasets:
+- HuggingFaceFW/fineweb-edu
+language:
+- en
+pipeline_tag: text-generation
+library_name: transformers
+tags:
+- small
+- cpu
+- supra
+- v2
+- tiny
+- mini
+- open
+- open-source
+new_version: SupraLabs/Supra-Mini-v4-2M
+---
+
+# 🦅 Supra Mini v2 0.1M
+Supra Mini **v2** 0.1M is a very, and we mean very small base model trained on 700 million tokens of Fineweb-Edu for 3 epochs as the **second version** of our Supra Mini series.
+
+## Model Config
+
+- Parameters: 167,760 (0.1M)
+- Architecture: Llama
+- Vocab size with custom BPE tokenizer: 2048
+- Hidden Size: 48
+- Intermediate Size: 96
+- Hidden Layers: 3
+- Attention Heads: 4
+- Max Position Embeddings: 256
+- Learning rate: 6e-4
+- Weight Decay: 0.01
+
+## Final Loss
+This model reached a final train loss of **4.413**.
+
+## Benchmarks
+
+All benchmarks were executed using `lm-eval`.
+
+| Task          | Value        | Random level |
+| :------------ | :----------: | -----------: |
+| Arc_Easy      | 0.2677       | 0.25 (25%)   |
+| Wikitext      | 7.7940       | -            |
+| BLiMP         | 0.5354       | 0.5 (50%)    |
+ 
+## Examples
+**Prompt:** "Artificial intelligence is "<br>
+**Output:**: *"Artificial intelligence is irreciously, and the diet of a battery.
+These are also known as the following:
+\- Foods, in the most commonly used to be taken by the priority of the South African American Modela, which was nothing for the first time. The federal government has been a wall of the world’s two moisture and the came on the national range of the Great Department of Amazonia, and the Politary Society of the Carli. This is that the Letters were the first third of the Building S"*
+<br><br>
+**Prompt:** "The main concept of physics is "<br>
+**Output:**: *"The main concept of physics is utilized in the most commonly used to be achieved.
+Chotos, an efficient for a dietary pathogene, and that are also known as a source of the molecularly. The failure wastered in the national categories of the California, and the Modela’s brought by the Florida. In the world's this time, it is not only on the sense of the first-metrial gardens, but they can be living in the task of the Political School of History (professed"*
+<br><br>
+**Prompt:** "Once upon a time, "<br>
+**Output:**: *"Once upon a time, utilizing the fire.
+The Samboard is that the lot of the bill.
+After the day to be a widely money and in the world’s a harmful force.
+Despite the size of the Bradese, and his owner, and he was noted by the Department of Management, and the Lord of Petersonia and the Council, which had been a destroyed on the tree, but the giving a came from the Hiska. The Great Marius and Jewish Amazon's Rich"*
+
+## Usage
+To use our model, just run this code using HF Transformers to execute the model:
+```python3
+from transformers import pipeline
+import torch
+
+print("[*] Loading Supra Mini v2 0.1M model from Hugging Face Hub...")
+pipe = pipeline(
+    "text-generation", 
+    model="SupraLabs/Supra-Mini-v2-0.1M",
+    device_map="auto",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+)
+
+def generate_text(prompt, max_length=150):
+    result = pipe(
+        prompt, 
+        max_new_tokens=max_length,
+        do_sample=True,
+        temperature=0.5,
+        top_k=25,
+        top_p=0.9,
+        repetition_penalty=1.2,
+        pad_token_id=pipe.tokenizer.pad_token_id,
+        eos_token_id=pipe.tokenizer.eos_token_id
+    )
+    return result[0]['generated_text']
+
+test_prompt = "The importance of education is"
+print(f"\nPrompt: {test_prompt}")
+print("-" * 30)
+print("\nOutput:\n" + generate_text(test_prompt))
+```
+
+## Use cases
+
+1. Educational research
+2. deployment or testing/fine-tuning on edge environments
+3. Or more simply, for fun
+
+## Limitations
+
+1. Cannot reason, chat, or code
+2. Incoherent more often than not
+3. Mostly unfactual
+
+## Training guide
+We trained Supra Mini v2 0.1M on a single T4 GPU in ~2 hours for 3 epochs.<br>
+The full training code can be found in this repo as `run.sh` (easily run the complete pipeline), `train_tokenizer.py` (train costum BPE tokenizer with vocab size of 2048), `train.py` (train the model) and `inference.py` (test the model).<br>
+The model was trained on the first 700 million tokens of Sample-10BT from Fineweb-Edu using streaming tokenization.
+
+## Final thoughts
+As this is the second version of the Supra Mini series, we are very proud to release it today!<br>
+*But:* stay tuned for more models and follow us to support our open-source work! 😊

benchmarks.md

@@ -0,0 +1,79 @@
+|                           Tasks                            |Version|Filter|n-shot|    Metric     |   |  Value   |   |Stderr|
+|------------------------------------------------------------|------:|------|-----:|---------------|---|---------:|---|------|
+|blimp                                                       |      2|none  |     0|acc            |↑  |    0.5354|±  |0.0017|
+| - blimp_adjunct_island                                     |      1|none  |     0|acc            |↑  |    0.5980|±  |0.0155|
+| - blimp_anaphor_gender_agreement                           |      1|none  |     0|acc            |↑  |    0.3130|±  |0.0147|
+| - blimp_anaphor_number_agreement                           |      1|none  |     0|acc            |↑  |    0.5090|±  |0.0158|
+| - blimp_animate_subject_passive                            |      1|none  |     0|acc            |↑  |    0.5750|±  |0.0156|
+| - blimp_animate_subject_trans                              |      1|none  |     0|acc            |↑  |    0.7470|±  |0.0138|
+| - blimp_causative                                          |      1|none  |     0|acc            |↑  |    0.4810|±  |0.0158|
+| - blimp_complex_NP_island                                  |      1|none  |     0|acc            |↑  |    0.4880|±  |0.0158|
+| - blimp_coordinate_structure_constraint_complex_left_branch|      1|none  |     0|acc            |↑  |    0.1420|±  |0.0110|
+| - blimp_coordinate_structure_constraint_object_extraction  |      1|none  |     0|acc            |↑  |    0.5820|±  |0.0156|
+| - blimp_determiner_noun_agreement_1                        |      1|none  |     0|acc            |↑  |    0.6580|±  |0.0150|
+| - blimp_determiner_noun_agreement_2                        |      1|none  |     0|acc            |↑  |    0.6320|±  |0.0153|
+| - blimp_determiner_noun_agreement_irregular_1              |      1|none  |     0|acc            |↑  |    0.5150|±  |0.0158|
+| - blimp_determiner_noun_agreement_irregular_2              |      1|none  |     0|acc            |↑  |    0.6980|±  |0.0145|
+| - blimp_determiner_noun_agreement_with_adj_2               |      1|none  |     0|acc            |↑  |    0.5780|±  |0.0156|
+| - blimp_determiner_noun_agreement_with_adj_irregular_1     |      1|none  |     0|acc            |↑  |    0.4220|±  |0.0156|
+| - blimp_determiner_noun_agreement_with_adj_irregular_2     |      1|none  |     0|acc            |↑  |    0.5170|±  |0.0158|
+| - blimp_determiner_noun_agreement_with_adjective_1         |      1|none  |     0|acc            |↑  |    0.6070|±  |0.0155|
+| - blimp_distractor_agreement_relational_noun               |      1|none  |     0|acc            |↑  |    0.3120|±  |0.0147|
+| - blimp_distractor_agreement_relative_clause               |      1|none  |     0|acc            |↑  |    0.3110|±  |0.0146|
+| - blimp_drop_argument                                      |      1|none  |     0|acc            |↑  |    0.7270|±  |0.0141|
+| - blimp_ellipsis_n_bar_1                                   |      1|none  |     0|acc            |↑  |    0.2180|±  |0.0131|
+| - blimp_ellipsis_n_bar_2                                   |      1|none  |     0|acc            |↑  |    0.3480|±  |0.0151|
+| - blimp_existential_there_object_raising                   |      1|none  |     0|acc            |↑  |    0.6860|±  |0.0147|
+| - blimp_existential_there_quantifiers_1                    |      1|none  |     0|acc            |↑  |    0.8100|±  |0.0124|
+| - blimp_existential_there_quantifiers_2                    |      1|none  |     0|acc            |↑  |    0.2950|±  |0.0144|
+| - blimp_existential_there_subject_raising                  |      1|none  |     0|acc            |↑  |    0.6880|±  |0.0147|
+| - blimp_expletive_it_object_raising                        |      1|none  |     0|acc            |↑  |    0.6570|±  |0.0150|
+| - blimp_inchoative                                         |      1|none  |     0|acc            |↑  |    0.3850|±  |0.0154|
+| - blimp_intransitive                                       |      1|none  |     0|acc            |↑  |    0.5170|±  |0.0158|
+| - blimp_irregular_past_participle_adjectives               |      1|none  |     0|acc            |↑  |    0.6620|±  |0.0150|
+| - blimp_irregular_past_participle_verbs                    |      1|none  |     0|acc            |↑  |    0.5050|±  |0.0158|
+| - blimp_irregular_plural_subject_verb_agreement_1          |      1|none  |     0|acc            |↑  |    0.5880|±  |0.0156|
+| - blimp_irregular_plural_subject_verb_agreement_2          |      1|none  |     0|acc            |↑  |    0.5860|±  |0.0156|
+| - blimp_left_branch_island_echo_question                   |      1|none  |     0|acc            |↑  |    0.9020|±  |0.0094|
+| - blimp_left_branch_island_simple_question                 |      1|none  |     0|acc            |↑  |    0.2310|±  |0.0133|
+| - blimp_matrix_question_npi_licensor_present               |      1|none  |     0|acc            |↑  |    0.0380|±  |0.0060|
+| - blimp_npi_present_1                                      |      1|none  |     0|acc            |↑  |    0.6520|±  |0.0151|
+| - blimp_npi_present_2                                      |      1|none  |     0|acc            |↑  |    0.6390|±  |0.0152|
+| - blimp_only_npi_licensor_present                          |      1|none  |     0|acc            |↑  |    0.0400|±  |0.0062|
+| - blimp_only_npi_scope                                     |      1|none  |     0|acc            |↑  |    0.0020|±  |0.0014|
+| - blimp_passive_1                                          |      1|none  |     0|acc            |↑  |    0.6520|±  |0.0151|
+| - blimp_passive_2                                          |      1|none  |     0|acc            |↑  |    0.6280|±  |0.0153|
+| - blimp_principle_A_c_command                              |      1|none  |     0|acc            |↑  |    0.6890|±  |0.0146|
+| - blimp_principle_A_case_1                                 |      1|none  |     0|acc            |↑  |    0.9990|±  |0.0010|
+| - blimp_principle_A_case_2                                 |      1|none  |     0|acc            |↑  |    0.4450|±  |0.0157|
+| - blimp_principle_A_domain_1                               |      1|none  |     0|acc            |↑  |    0.8820|±  |0.0102|
+| - blimp_principle_A_domain_2                               |      1|none  |     0|acc            |↑  |    0.5450|±  |0.0158|
+| - blimp_principle_A_domain_3                               |      1|none  |     0|acc            |↑  |    0.4690|±  |0.0158|
+| - blimp_principle_A_reconstruction                         |      1|none  |     0|acc            |↑  |    0.3830|±  |0.0154|
+| - blimp_regular_plural_subject_verb_agreement_1            |      1|none  |     0|acc            |↑  |    0.6890|±  |0.0146|
+| - blimp_regular_plural_subject_verb_agreement_2            |      1|none  |     0|acc            |↑  |    0.5760|±  |0.0156|
+| - blimp_sentential_negation_npi_licensor_present           |      1|none  |     0|acc            |↑  |    0.9990|±  |0.0010|
+| - blimp_sentential_negation_npi_scope                      |      1|none  |     0|acc            |↑  |    0.4590|±  |0.0158|
+| - blimp_sentential_subject_island                          |      1|none  |     0|acc            |↑  |    0.2760|±  |0.0141|
+| - blimp_superlative_quantifiers_1                          |      1|none  |     0|acc            |↑  |    0.3040|±  |0.0146|
+| - blimp_superlative_quantifiers_2                          |      1|none  |     0|acc            |↑  |    0.3620|±  |0.0152|
+| - blimp_tough_vs_raising_1                                 |      1|none  |     0|acc            |↑  |    0.3310|±  |0.0149|
+| - blimp_tough_vs_raising_2                                 |      1|none  |     0|acc            |↑  |    0.6970|±  |0.0145|
+| - blimp_transitive                                         |      1|none  |     0|acc            |↑  |    0.6560|±  |0.0150|
+| - blimp_wh_island                                          |      1|none  |     0|acc            |↑  |    0.5110|±  |0.0158|
+| - blimp_wh_questions_object_gap                            |      1|none  |     0|acc            |↑  |    0.6180|±  |0.0154|
+| - blimp_wh_questions_subject_gap                           |      1|none  |     0|acc            |↑  |    0.9480|±  |0.0070|
+| - blimp_wh_questions_subject_gap_long_distance             |      1|none  |     0|acc            |↑  |    0.8930|±  |0.0098|
+| - blimp_wh_vs_that_no_gap                                  |      1|none  |     0|acc            |↑  |    0.9960|±  |0.0020|
+| - blimp_wh_vs_that_no_gap_long_distance                    |      1|none  |     0|acc            |↑  |    0.9910|±  |0.0030|
+| - blimp_wh_vs_that_with_gap                                |      1|none  |     0|acc            |↑  |    0.0090|±  |0.0030|
+| - blimp_wh_vs_that_with_gap_long_distance                  |      1|none  |     0|acc            |↑  |    0.0040|±  |0.0020|
+|arc_easy                                                    |      1|none  |     0|acc            |↑  |    0.2677|±  |0.0091|
+|                                                            |       |none  |     0|acc_norm       |↑  |    0.2841|±  |0.0093|
+|wikitext                                                    |      2|none  |     0|bits_per_byte  |↓  |    2.9624|±  |   N/A|
+|                                                            |       |none  |     0|byte_perplexity|↓  |    7.7940|±  |   N/A|
+|                                                            |       |none  |     0|word_perplexity|↓  |58699.2441|±  |   N/A|
+
+|Groups|Version|Filter|n-shot|Metric|   |Value |   |Stderr|
+|------|------:|------|-----:|------|---|-----:|---|-----:|
+|blimp |      2|none  |     0|acc   |↑  |0.5354|±  |0.0017|

config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "head_dim": 12,
+  "hidden_act": "silu",
+  "hidden_size": 48,
+  "initializer_range": 0.02,
+  "intermediate_size": 96,
+  "max_position_embeddings": 256,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 3,
+  "num_key_value_heads": 4,
+  "pad_token_id": 1,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.0.0",
+  "use_cache": false,
+  "vocab_size": 2048
+}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 1,
+  "transformers_version": "5.0.0",
+  "use_cache": true
+}

inference.py

@@ -0,0 +1,32 @@
+print("[*] Loading libraries...")
+import torch
+from transformers import LlamaForCausalLM, PreTrainedTokenizerFast
+
+model_path = "./Supra-Mini-v2-0.1m-FINAL"
+
+print("[*] Loading tokenizer...")
+tokenizer = PreTrainedTokenizerFast.from_pretrained(model_path)
+
+print("[*] Loading model...")
+model = LlamaForCausalLM.from_pretrained(model_path)
+model.eval()
+
+prompt = "The main concept of physics is "
+print(f"[*] Prompt: {prompt!r}")
+
+inputs = tokenizer(prompt, return_tensors="pt")
+
+with torch.no_grad():
+    outputs = model.generate(
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        max_new_tokens=150,
+        do_sample=True,
+        temperature=0.35,
+        top_p=0.85,
+        repetition_penalty=1.2,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+print("[*] Output:", tokenizer.decode(outputs[0], skip_special_tokens=True))

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2903fbf80d221058f5b6280e1f9bc8832eda5a4d5da0d5f5d54b5d5f81376164
+size 674032

run.sh

@@ -0,0 +1,8 @@
+pip install -q transformers datasets sentencepiece accelerate
+python3 train_tokenizer.py
+python3 train_model.py
+python3 inference.py
+sed -i 's/"max_position_embeddings": 256/"max_position_embeddings": 1024/' Supra-Mini-v2-0.1m-FINAL/config.json
+pip install lm-eval
+lm_eval --model hf     --model_args pretrained=./Supra-Mini-v2-0.1m-FINAL,max_length=1024     --tasks arc_easy,wikitext,blimp     --device cuda     --batch_size 16
+sed -i 's/"max_position_embeddings": 1024/"max_position_embeddings": 256/' Supra-Mini-v2-0.1m-FINAL/config.json

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}

train_model.py

@@ -0,0 +1,179 @@
+import os
+os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+print("[*] Loading libraries...")
+import torch
+import math
+import numpy as np
+from datasets import load_dataset
+from tokenizers import ByteLevelBPETokenizer
+from transformers import (
+    LlamaConfig,
+    LlamaForCausalLM,
+    PreTrainedTokenizerFast,
+    Trainer,
+    TrainingArguments,
+)
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+print("[*] Loading tokenizer...")
+fast_tokenizer = ByteLevelBPETokenizer(
+    "./custom_llama_tokenizer-vocab.json",
+    "./custom_llama_tokenizer-merges.txt"
+)
+tokenizer = PreTrainedTokenizerFast(
+    tokenizer_object=fast_tokenizer,
+    bos_token="<s>",
+    eos_token="</s>",
+    unk_token="<unk>",
+    pad_token="<pad>",
+)
+
+TOKEN_BIN     = "/kaggle/working/tokens.bin"
+TARGET_TOKENS = 700_000_000
+SEQ_LEN       = 256
+BATCH_TEXTS   = 1000
+FLUSH_EVERY   = 1_000_000
+
+
+def build_token_bin(fast_tokenizer, path=TOKEN_BIN, target_tokens=TARGET_TOKENS):
+    if os.path.exists(path) and os.path.getsize(path) >= target_tokens * 2:
+        print(f"[=] Reusing existing token file: {path}")
+        return
+
+    print(f"[*] Streaming + tokenizing {target_tokens:,} tokens → {path}")
+    mm = np.memmap(path, dtype=np.uint16, mode="w+", shape=(target_tokens,))
+
+    dataset = load_dataset(
+        "HuggingFaceFW/fineweb-edu", "sample-10BT",
+        split="train", streaming=True
+    )
+
+    written = 0
+    buf = []
+    texts = []
+    pbar = tqdm(total=target_tokens, desc="[*] Gathering tokens", unit="tok")
+
+    def flush_buf():
+        nonlocal written, buf
+        if not buf:
+            return False
+        n = min(len(buf), target_tokens - written)
+        mm[written:written + n] = np.asarray(buf[:n], dtype=np.uint16)
+        written += n
+        pbar.update(n)
+        del buf[:n]
+        return written >= target_tokens
+
+    for example in dataset:
+        texts.append(example["text"])
+        if len(texts) >= BATCH_TEXTS:
+            encs = fast_tokenizer.encode_batch(texts)
+            texts.clear()
+            for e in encs:
+                buf.extend(e.ids)
+            if len(buf) >= FLUSH_EVERY:
+                if flush_buf():
+                    break
+
+    if written < target_tokens and texts:
+        encs = fast_tokenizer.encode_batch(texts)
+        for e in encs:
+            buf.extend(e.ids)
+    if written < target_tokens:
+        flush_buf()
+
+    pbar.close()
+    mm.flush()
+    del mm
+    print(f"[+] Wrote {written:,} tokens to {path} "
+          f"({os.path.getsize(path)/1e6:.1f} MB)")
+
+
+class MemmapDataset(Dataset):
+    def __init__(self, path, total_tokens, seq_len=SEQ_LEN):
+        self.path     = path
+        self.seq_len  = seq_len
+        self.n_chunks = total_tokens // seq_len
+        self._data    = None  # lazy open (Multiprocessing-safe)
+
+    @property
+    def data(self):
+        if self._data is None:
+            self._data = np.memmap(
+                self.path, dtype=np.uint16, mode="r",
+                shape=(self.n_chunks * self.seq_len,)
+            )
+        return self._data
+
+    def __len__(self):
+        return self.n_chunks
+
+    def __getitem__(self, idx):
+        s   = idx * self.seq_len
+        arr = np.asarray(self.data[s:s + self.seq_len], dtype=np.int64)
+        ids = torch.from_numpy(arr)
+        return {"input_ids": ids, "labels": ids.clone()}
+
+
+def collate_fn(batch):
+    input_ids = torch.stack([b["input_ids"] for b in batch])
+    labels    = torch.stack([b["labels"]    for b in batch])
+    return {"input_ids": input_ids, "labels": labels}
+
+
+print(f"[*] Preparing {TARGET_TOKENS:,} tokens (streaming, memmap-backed)...")
+build_token_bin(fast_tokenizer, TOKEN_BIN, TARGET_TOKENS)
+dataset = MemmapDataset(TOKEN_BIN, TARGET_TOKENS, seq_len=SEQ_LEN)
+print(f"[+] Dataset ready: {len(dataset):,} chunks of {SEQ_LEN} tokens")
+
+print("[*] Setting up model...")
+config = LlamaConfig(
+    vocab_size=len(tokenizer.get_vocab()),
+    hidden_size=48,
+    intermediate_size=96,
+    num_hidden_layers=3,
+    num_attention_heads=4,
+    max_position_embeddings=256,
+    tie_word_embeddings=True,
+    pad_token_id=tokenizer.pad_token_id,
+    bos_token_id=tokenizer.bos_token_id,
+    eos_token_id=tokenizer.eos_token_id,
+)
+model = LlamaForCausalLM(config)
+print(f"[*] Model parameters: {model.num_parameters():,}")
+
+print("[*] Defining training arguments...")
+training_args = TrainingArguments(
+    output_dir="./Supra-Mini-v2-0.1m",
+    num_train_epochs=3,
+    per_device_train_batch_size=1024,
+    gradient_accumulation_steps=1,
+    save_steps=500,
+    save_total_limit=2,
+    logging_steps=100,
+    weight_decay=0.01,
+    fp16=torch.cuda.is_available(),
+    push_to_hub=False,
+    report_to="none",
+    dataloader_num_workers=2,
+    dataloader_pin_memory=True,
+    learning_rate=6e-4,
+    lr_scheduler_type="cosine",
+    warmup_ratio=0.05,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset,
+    data_collator=collate_fn,
+)
+
+print("[*] Starting training...")
+trainer.train()
+trainer.save_model("./Supra-Mini-v2-0.1m-FINAL")
+tokenizer.save_pretrained("./Supra-Mini-v2-0.1m-FINAL")
+print("[*] Training finished.")

train_tokenizer.py

@@ -0,0 +1,23 @@
+print("[*] Loading libraries...")
+from datasets import load_dataset
+from tokenizers import ByteLevelBPETokenizer
+
+dataset = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT", split="train", streaming=True)
+def get_training_corpus():
+    dataset_iter = iter(dataset)
+    for _ in range(50000):
+        yield next(dataset_iter)["text"]
+
+tokenizer = ByteLevelBPETokenizer()
+
+print("[*] Training tokenizer...")
+
+tokenizer.train_from_iterator(
+    get_training_corpus(),
+    vocab_size=2048,
+    min_frequency=2,
+    special_tokens=["<s>", "<pad>", "</s>", "<unk>", "<mask>"]
+)
+
+tokenizer.save_model(".", "custom_llama_tokenizer")
+print("[*] Tokenizer training complete!")

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:475e4bde8bedd30ee7a419e297d7166e343668fec0ee93cf9fa1436fc10e0a96
+size 5201