# Training: STaR round 3 LoRA adapter

This folder holds the code, settings and data that made the adapter in [SinhalaMMLU-SharedTask/team-sinhalanuwana-gemma4-e4b-star3-lora](https://huggingface.co/SinhalaMMLU-SharedTask/team-sinhalanuwana-gemma4-e4b-star3-lora). The adapter trains from the base model `google/gemma-4-E4B-it` (revision `ee0ef6023621cff504d758262d4e04895a5af4a2`). It uses no earlier adapter.

## Files

| File | Purpose |
| --- | --- |
| `review_overlap.py` | Compares each training question with every dev and test question, by wording (character 6-gram Jaccard and containment, same choices) and by meaning (LaBSE cosine). Lists the matches for removal. |
| `prepare_star_data.py` | Removes the listed questions and splits the rest into 2 shards with seed 0 |
| `star.py` | Lets the model think about each training question and keeps the thoughts that finish and give the correct answer |
| `train_lora.py` | Trains the LoRA adapter on the kept thoughts |
| `kaggle/star3-gen-0`, `kaggle/star3-gen-1` | The Kaggle notebooks that ran `star.py` on the two shards (2 x T4 each, about 4.2 hours each) |
| `kaggle/star3-train` | The Kaggle notebook that joined their kept thoughts and ran `train_lora.py` (2 x T4, 4.3 hours) |
| `train_log.json` | Settings, parameter counts and the loss of every step of the training run |
| `data/shard-0.jsonl`, `data/shard-1.jsonl` | The 8,894 training questions (4,447 each), the input of `star.py` |
| `data/excluded_ids.txt` | The ids of the 1,106 questions that `review_overlap.py` removed |
| `data/overlap-summary.json` | The counts of the overlap review, by rule |
| `data/star3-thoughts.jsonl` | The 4,073 training examples: prompt, thought and answer. This is the exact file the adapter trained on. |

## Training data

Anthropic claude-haiku-5-5 generated Sinhala multiple-choice questions from Sri Lankan school textbooks, NIE teacher guides and resource books, A/L learning outcome notes and Wikibooks. Anthropic claude-sonnet-5-5 (9,228 questions) and claude-opus-5-5 (772 questions) checked the answer keys of 10,000 of them. The shared task allows closed models for training data.

Each question in `data/shard-*.jsonl` has its `id`, `subject`, `level`, `grade`, `source`, `question`, the 4 `choices`, the checked `answer` (1 to 4) and an `explanation`.

We removed the questions that resemble an organiser question before any training. `review_overlap.py` compared the 10,000 questions with the 1,851 dev and 2,238 test questions, and with an earlier copy of the test set (2,146 questions), and removed 1,106 by these rules. A question can match more than one rule, so the counts add up to more than 1,106:

| Rule | Removed |
| --- | --- |
| Same question text | 0 |
| Question Jaccard of 0.5 or more | 17 |
| Question containment of 0.8 or more | 48 |
| 3 or more same choices and containment of 0.3 or more | 11 |
| All 4 choices the same | 43 |
| LaBSE cosine of 0.70 or more | 1,033 |

The removed questions are not in this repository. The dev and test questions were not used for training in any form.

## Steps

Run the commands from the root of this repository. The training scripts import `predict.py` from the root, so set `PYTHONPATH=.`.

1. Remove the questions that resemble a dev or test question, and make the shards. These steps need the full 10,000-question file and the organiser dev and test sets, which are not in this repository. Their output is in `data/`.

   ```bash
   PYTHONPATH=. python training/review_overlap.py --questions verified-mcq.jsonl \
     --eval dev=DEV_DIR test=TEST_DIR --out_dir overlap
   PYTHONPATH=. python training/prepare_star_data.py --questions verified-mcq.jsonl \
     --exclude overlap/excluded_ids.txt --out_dir training/data --shards 2 --seed 0
   ```

2. Let the base model think about each question, with no adapter and no hinted retry. One run for each shard, a thought of up to 512 tokens, greedy, batches of 32:

   ```bash
   PYTHONPATH=. python training/star.py --train training/data/shard-0.jsonl --no_hint_retry \
     --revision ee0ef6023621cff504d758262d4e04895a5af4a2 --out_dir star-0
   PYTHONPATH=. python training/star.py --train training/data/shard-1.jsonl --no_hint_retry \
     --revision ee0ef6023621cff504d758262d4e04895a5af4a2 --out_dir star-1
   ```

   `star.py` keeps a thought only if it ends within 512 tokens and the answer after it is the checked answer. It kept 4,073 of the 8,894.

3. Join the kept thoughts of the two runs:

   ```bash
   cat star-0/star.jsonl star-1/star.jsonl > star3.jsonl
   ```

   The result is byte for byte `data/star3-thoughts.jsonl`. `train_lora.py` shuffles the examples with seed 0.

4. Train the adapter:

   ```bash
   PYTHONPATH=. python training/train_lora.py --data star3.jsonl \
     --revision ee0ef6023621cff504d758262d4e04895a5af4a2 --out_dir adapter
   ```

   To train on the published examples directly, pass `--data training/data/star3-thoughts.jsonl`.

## Settings

| Setting | Value |
| --- | --- |
| LoRA | Rank 16, alpha 32, dropout 0.05, on all 342 linear layers of the 42 decoder layers (`q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`, `per_layer_input_gate`, `per_layer_projection`) |
| Trainable parameters | 38,666,240 |
| Loss | Cross-entropy on the thought, the answer cue and the answer number, not on the prompt |
| Optimiser | AdamW, learning rate 1e-4, cosine schedule with 3% warmup |
| Batch | 4 examples, 4 accumulation steps, so 16 examples per step |
| Epochs | 2 (510 steps) |
| Maximum length | 1,024 tokens (no example was dropped) |
| Precision | float16 with gradient scaling (T4 GPUs have no bfloat16) |
| Seed | 0 |
| Hardware and time | 2 x Tesla T4 on Kaggle, 15,552 seconds (4.3 hours) |

The Kaggle notebooks used the same scripts from a snapshot of our code with all scripts in one folder, and the shards from a private Kaggle dataset. Their `run.py` installs `requirements.txt`, removes `torchao`, downloads the model and runs the command above.