Larger corpus: full-text papers, smoltalk2, resumable training

README.md

@@ -18,9 +18,10 @@ tags:
   - webgpu
 datasets:
   - vigneshwar234/vashtra-ml-corpus
+  - common-pile/arxiv_papers_filtered
   - CShorten/ML-ArXiv-Papers
   - donfu/oa-stackexchange
-  - HuggingFaceTB/smol-smoltalk
+  - HuggingFaceTB/smoltalk2
 ---
 
 # Vashtra-0.6B
@@ -43,19 +44,29 @@ does.
 Two stages, both full-parameter. At 0.6B a full fine-tune beats LoRA for domain
 adaptation and still fits in a free Colab T4, using about 10GB of 16GB.
 
-**Stage A, continued pretraining.** 117k arXiv `cs.LG` titles and abstracts, packed
-into dense 1024-token blocks so no compute goes to padding. LR 5e-5, cosine, one
-epoch. The LR is low on purpose. The goal is to absorb the domain without wrecking
-the base model's general ability.
+**Stage A, continued pretraining.** Full-text arXiv ML papers from `common-pile`
+(openly licensed only) plus 117k arXiv `cs.LG` titles and abstracts, packed into
+dense 1024-token blocks so no compute goes to padding. LR 5e-5, cosine, one epoch.
+The LR is low on purpose. The goal is to absorb the domain without wrecking the
+base model's general ability.
+
+Full text matters here. Abstracts teach the model what papers claim. Full text
+teaches it how the argument is built, which is what you want when someone asks it
+to explain a method.
+
+The corpus is about 90M tokens, more than a free Colab session will finish, so
+stage A takes a `--cpt_tokens` budget and pushes checkpoints to the Hub. Passing
+`--resume` picks up where a dropped session stopped.
 
 **Stage B, supervised fine-tuning.** ML, stats and CS StackExchange answers that were
-accepted with score 3 or higher, plus arXiv explanation and titling tasks, plus a
-general chat anchor. LR 1e-5, two epochs, loss masked so only assistant turns count.
+accepted with score 3 or higher, arXiv explanation and titling tasks, and three
+subsets of `smoltalk2`: science reasoning, instruction following, and general chat.
+LR 1e-5, two epochs, loss masked so only assistant turns count.
 
 ### The general chat anchor
 
 Fine-tuning a 0.6B model only on domain data reliably destroys its ability to hold a
-normal conversation. About 20k general examples from `smol-smoltalk` stay in the SFT
+normal conversation. About 22k general examples from `smoltalk2` stay in the SFT
 mixture to stop that. It costs a little domain sharpness and keeps the model usable.
 
 ### A tokenisation detail
@@ -104,16 +115,15 @@ Useful for explaining ML concepts, summarising papers, and on-device or edge cas
 where 0.6B is the size you can afford.
 
 It is still a 0.6B model. It gets specifics wrong, especially exact numbers, named
-papers, and anything recent. The arXiv data is abstracts only, so it knows what
-papers claim, not what they showed. Do not cite it and do not use it for anything
-that matters without checking the answer.
+papers, and anything recent. Do not cite it and do not use it for anything that
+matters without checking the answer.
 
 ## Licence
 
-Apache-2.0, same as the Qwen3-0.6B base. The corpus mixes licences: arXiv abstracts
-are AFL-3.0 through `CShorten/ML-ArXiv-Papers`, StackExchange content is CC-BY-SA-4.0,
-and `smol-smoltalk` is Apache-2.0. Anything derived from StackExchange stays
-CC-BY-SA-4.0.
+Apache-2.0, same as the Qwen3-0.6B base. The corpus mixes licences: full-text papers
+are CC-BY, CC-BY-SA or CC0 through `common-pile/arxiv_papers_filtered`, abstracts are
+AFL-3.0 through `CShorten/ML-ArXiv-Papers`, StackExchange content is CC-BY-SA-4.0,
+and `smoltalk2` is Apache-2.0. Anything derived from StackExchange stays CC-BY-SA-4.0.
 
 Built by [vigneshwar234](https://huggingface.co/vigneshwar234).
 Code on [GitHub](https://github.com/vignesh2027/Vashtra).

Vashtra_training.ipynb

@@ -4,17 +4,17 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Vashtra-0.6B \u2014 build the corpus, train the model, ship the demo\n",
+    "# Vashtra-0.6B: build the corpus, train the model, ship the demo\n",
     "\n",
     "Two-stage domain adaptation of **Qwen3-0.6B** into a machine-learning specialist.\n",
     "\n",
     "| Stage | Data | What it teaches |\n",
     "|---|---|---|\n",
-    "| A \u2014 continued pretraining | 117k arXiv cs.LG abstracts | ML vocabulary and concepts |\n",
-    "| B \u2014 supervised fine-tuning | ML Q&A + general chat | answering, instruction following |\n",
+    "| A, continued pretraining | full-text arXiv ML papers + 117k cs.LG abstracts | ML vocabulary, concepts, how arguments are made |\n",
+    "| B, supervised fine-tuning | ML Q&A + general chat | answering, instruction following |\n",
     "\n",
     "**Runtime \u2192 Change runtime type \u2192 T4 GPU**, then **Runtime \u2192 Run all**.\n",
-    "Everything pushes to the Hub automatically. Budget ~2\u20133 hours on the free tier.\n"
+    "Everything pushes to the Hub automatically. Budget ~2 to 3 hours on the free tier.\n"
    ]
   },
   {
@@ -22,7 +22,7 @@
    "metadata": {},
    "source": [
     "### 1. Confirm you have a GPU\n",
-    "If this cell fails you are on CPU \u2014 switch the runtime first, or nothing below will finish."
+    "If this cell fails you are on CPU. Switch the runtime first, or nothing below will finish."
    ]
   },
   {
@@ -102,12 +102,12 @@
    "source": [
     "### 5. Build the corpus\n",
     "\n",
-    "Streams arXiv cs.LG, ML/stats/CS StackExchange, and the general-chat anchor into\n",
-    "`cpt.jsonl` + `sft.jsonl`, then pushes them to\n",
-    "`vigneshwar234/vashtra-ml-corpus`. Takes ~15\u201325 minutes, mostly the\n",
-    "StackExchange scan.\n",
+    "Streams full-text arXiv papers, cs.LG abstracts, ML/stats/CS StackExchange, and\n",
+    "the smoltalk2 mixture into `cpt.jsonl` + `sft.jsonl`, then pushes them to\n",
+    "`vigneshwar234/vashtra-ml-corpus`. Takes ~25 to 40 minutes, mostly the\n",
+    "StackExchange and full-text scans.\n",
     "\n",
-    "Skip this cell if the corpus is already built \u2014 training pulls it from the Hub."
+    "Skip this cell if the corpus is already built. Training pulls it from the Hub."
    ]
   },
   {
@@ -128,8 +128,30 @@
     "Stage A (continued pretraining) then Stage B (SFT), then the model is pushed to\n",
     "`vigneshwar234/Vashtra-0.6B`.\n",
     "\n",
-    "Running short on session time? Use `--cpt_docs 40000 --sft_epochs 1` to cut\n",
-    "stage A and B down substantially."
+    "The corpus is about 90M tokens, which is more than a free Colab session will get\n",
+    "through. Stage A therefore takes a token budget, defaulting to 30M, which is\n",
+    "roughly two hours on a T4. Checkpoints are pushed to the Hub as it goes.\n",
+    "\n",
+    "If the session drops, open a new one and run the same cell with `--resume`. It\n",
+    "pulls the last checkpoint back down and carries on instead of starting over.\n",
+    "To train on more of the corpus, raise `--cpt_tokens` (or pass `0` for all of it)\n",
+    "and expect to resume across two or three sessions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python train_vashtra.py --stage all --cpt_tokens 30000000"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If the session above was cut off, run this instead of starting over."
    ]
   },
   {
@@ -138,7 +160,7 @@
    "execution_count": null,
    "outputs": [],
    "source": [
-    "!python train_vashtra.py --stage all"
+    "# !python train_vashtra.py --stage all --cpt_tokens 30000000 --resume"
    ]
   },
   {
@@ -147,7 +169,7 @@
    "source": [
     "### 7. Smoke-test what you trained\n",
     "\n",
-    "Sanity check, not an evaluation. You are looking for coherent, on-topic ML answers \u2014\n",
+    "Sanity check, not an evaluation. You are looking for coherent, on-topic ML answers,\n",
     "and for the general question to still work, which is what the anchor data buys."
    ]
   },
@@ -183,8 +205,8 @@
     "### 8. Measure it against the base model\n",
     "\n",
     "Compare perplexity on held-out arXiv ML text. If stage A worked, Vashtra should be\n",
-    "meaningfully lower than Qwen3-0.6B here. Put the real number in the model card \u2014\n",
-    "this is the claim that makes the model credible."
+    "meaningfully lower than Qwen3-0.6B here. Put the real number in the model card.\n",
+    "This is the claim that makes the model credible."
    ]
   },
   {
@@ -228,7 +250,7 @@
     "### 9. Export to ONNX for the in-browser demo\n",
     "\n",
     "Produces `vigneshwar234/Vashtra-0.6B-ONNX`, which the static WebGPU Space loads\n",
-    "directly \u2014 so the demo costs nothing to host and runs on the visitor's own GPU."
+    "directly, so the demo costs nothing to host and runs on the visitor's own GPU."
    ]
   },
   {
@@ -257,7 +279,7 @@
     "\n",
     "**Last step:** put the real perplexity number from cell 8 into the model card, and\n",
     "replace the \"weights are training\" banner. Measured results are what make people\n",
-    "trust a small model \u2014 ship them.\n"
+    "trust a small model, so ship them.\n"
    ]
   }
  ],

build_corpus.py

@@ -1,16 +1,21 @@
 """
 Build the Vashtra training corpus and push it to the Hub.
 
-Runs on Colab (fast network, plenty of disk). Produces two files in
-vigneshwar234/vashtra-ml-corpus:
+Runs on Colab, where network and disk are not a constraint. Produces two files
+in vigneshwar234/vashtra-ml-corpus:
 
-  cpt.jsonl  {"text": ...}      raw ML literature for continued pretraining
-  sft.jsonl  {"messages": [..]} chat data for instruction tuning
+  cpt.jsonl  {"text": ...}       ML text for continued pretraining
+  sft.jsonl  {"messages": [..]}  chat data for instruction tuning
 
 Sources
-  CShorten/ML-ArXiv-Papers      117k arXiv cs.LG titles + abstracts
-  donfu/oa-stackexchange        ML / stats / CS Q&A with accepted answers
-  HuggingFaceTB/smol-smoltalk   general chat, anchors against forgetting
+  common-pile/arxiv_papers_filtered   full-text openly-licensed arXiv papers
+  CShorten/ML-ArXiv-Papers            117k arXiv cs.LG titles + abstracts
+  donfu/oa-stackexchange              ML / stats / CS Q&A, accepted answers
+  HuggingFaceTB/smoltalk2             science reasoning, instruction following,
+                                      general chat anchor
+
+Full text is the point. Abstracts teach the model what papers claim; full text
+teaches it how the argument is actually made.
 """
 import json, random, re, os, collections
 from datasets import load_dataset
@@ -20,36 +25,78 @@ OUT = "corpus"
 os.makedirs(OUT, exist_ok=True)
 REPO = "vigneshwar234/vashtra-ml-corpus"
 
-WS = re.compile(r"\s+")
-def clean(s: str) -> str:
-    return WS.sub(" ", (s or "")).strip()
+# How much full-text to keep. Each paper is truncated to CHARS_PER_PAPER, so
+# the budget is predictable rather than dominated by a few 80-page papers.
+FULLTEXT_PAPERS   = 22000
+CHARS_PER_PAPER   = 12000
+ANCHOR_TARGET     = 22000
+SCIENCE_TARGET    = 30000
+IFEVAL_TARGET     = 12000
 
+WS = re.compile(r"\s+")
+def clean(s): return WS.sub(" ", (s or "")).strip()
+
+# Terms that mark a paper as machine-learning rather than, say, cosmology.
+ML_TERMS = [
+    "neural network", "deep learning", "machine learning", "gradient descent",
+    "training data", "transformer", "convolutional", "reinforcement learning",
+    "classifier", "supervised", "unsupervised", "embedding", "backpropagation",
+    "loss function", "overfitting", "attention mechanism", "language model",
+    "generalization", "stochastic gradient", "fine-tuning", "benchmark dataset",
+]
 
-# ------------------------------------------------------------------ 1. arXiv
-print("[1/4] arXiv cs.LG …")
+cpt_lines = 0
+cpt_chars = 0
+fh_cpt = open(f"{OUT}/cpt.jsonl", "w")
+
+
+# ------------------------------------------- 1. full-text arXiv (the upgrade)
+print("[1/5] full-text arXiv (common-pile) ...")
+kept = 0
+try:
+    ft = load_dataset("common-pile/arxiv_papers_filtered", split="train",
+                      streaming=True)
+    for r in ft:
+        txt = r.get("text") or ""
+        if len(txt) < 4000:
+            continue
+        head = txt[:4000].lower()
+        score = sum(1 for t in ML_TERMS if t in head)
+        if score < 2:                      # not an ML paper
+            continue
+        body = clean(txt)[:CHARS_PER_PAPER]
+        fh_cpt.write(json.dumps({"text": body}) + "\n")
+        cpt_lines += 1; cpt_chars += len(body); kept += 1
+        if kept >= FULLTEXT_PAPERS:
+            break
+    print(f"      full-text papers kept: {kept}")
+except Exception as e:
+    print(f"      full-text SKIPPED ({type(e).__name__}: {e})")
+
+
+# --------------------------------------------------- 2. arXiv cs.LG abstracts
+print("[2/5] arXiv cs.LG abstracts ...")
 ax = load_dataset("CShorten/ML-ArXiv-Papers", split="train")
 papers, seen_title = [], set()
 for r in ax:
     t, a = clean(r.get("title")), clean(r.get("abstract"))
-    if not t or not a:
+    if not t or not a or not (300 < len(a) < 2200) or len(t) < 15:
         continue
-    if not (300 < len(a) < 2200) or len(t) < 15:
+    k = t.lower()[:90]
+    if k in seen_title:
         continue
-    key = t.lower()[:90]
-    if key in seen_title:
-        continue
-    seen_title.add(key)
+    seen_title.add(k)
     papers.append((t, a))
-print(f"      usable papers: {len(papers)}")
+for t, a in papers:
+    doc = f"# {t}\n\n{a}"
+    fh_cpt.write(json.dumps({"text": doc}) + "\n")
+    cpt_lines += 1; cpt_chars += len(doc)
+fh_cpt.close()
+print(f"      abstracts: {len(papers)}")
+print(f"      cpt.jsonl: {cpt_lines} docs, ~{cpt_chars/4/1e6:.0f}M tokens")
 
-with open(f"{OUT}/cpt.jsonl", "w") as fh:
-    for t, a in papers:
-        fh.write(json.dumps({"text": f"# {t}\n\n{a}"}) + "\n")
-print(f"      cpt.jsonl: {len(papers)} docs "
-      f"(~{sum(len(t)+len(a) for t, a in papers)/4/1e6:.1f}M tokens)")
 
-# arXiv-derived instruction tasks. Varied phrasings so the model learns the
-# task rather than the template.
+# ------------------------------------------------ 3. arXiv instruction tasks
 EXPLAIN = [
     'Explain the main contribution of the paper "{t}".',
     'What problem does the paper "{t}" address, and how?',
@@ -63,7 +110,6 @@ TITLE = [
     "What would this machine-learning paper be titled?\n\n{a}",
     "Read this abstract and give the paper's title:\n\n{a}",
 ]
-
 sft, stats = [], collections.Counter()
 for i, (t, a) in enumerate(random.sample(papers, min(26000, len(papers)))):
     if i % 4 == 0:
@@ -72,12 +118,12 @@ for i, (t, a) in enumerate(random.sample(papers, min(26000, len(papers)))):
         u, r = random.choice(EXPLAIN).format(t=t), a
     sft.append({"messages": [{"role": "user", "content": u},
                              {"role": "assistant", "content": r}]})
-    stats["arxiv"] += 1
-print(f"      arXiv SFT tasks: {stats['arxiv']}")
+    stats["arxiv_tasks"] += 1
+print(f"[3/5] arXiv SFT tasks: {stats['arxiv_tasks']}")
 
 
-# ---------------------------------------------------- 2. ML/CS StackExchange
-print("[2/4] StackExchange (ML / stats / CS) …")
+# ---------------------------------------------------- 4. ML/CS StackExchange
+print("[4/5] StackExchange (ML / stats / CS) ...")
 SITES = {
     "stackexchange-datascience", "stackexchange-ai", "stackexchange-stats",
     "stackexchange-cs", "stackexchange-cstheory",
@@ -85,8 +131,7 @@ SITES = {
     "stackexchange-crossvalidated",
 }
 se = load_dataset("donfu/oa-stackexchange", split="train", streaming=True)
-seen_q, n_se = set(), 0
-by_site = collections.Counter()
+seen_q, by_site = set(), collections.Counter()
 for r in se:
     if r["SOURCE"] not in SITES:
         continue
@@ -96,56 +141,70 @@ for r in se:
     q, a = clean(r["INSTRUCTION"]), clean(r["RESPONSE"])
     if not (40 < len(q) < 1200) or not (80 < len(a) < 1600):
         continue
-    key = q.lower()[:110]
-    if key in seen_q:
+    k = q.lower()[:110]
+    if k in seen_q:
         continue
-    seen_q.add(key)
+    seen_q.add(k)
     sft.append({"messages": [{"role": "user", "content": q},
                              {"role": "assistant", "content": a}]})
     by_site[r["SOURCE"]] += 1
-    n_se += 1
-stats["stackexchange"] = n_se
-print(f"      StackExchange SFT: {n_se}")
+    stats["stackexchange"] += 1
+print(f"      StackExchange SFT: {stats['stackexchange']}")
 for s, c in by_site.most_common():
     print(f"        {c:>6}  {s}")
 
 
-# ------------------------------------------------------- 3. general anchor
-print("[3/4] general-chat anchor …")
-anchor = load_dataset("HuggingFaceTB/smol-smoltalk", split="train", streaming=True)
-n_anchor = 0
-TARGET_ANCHOR = 20000
-for row in anchor:
-    m = row["messages"]
-    if not (2 <= len(m) <= 6):
-        continue
-    if sum(len(x["content"]) for x in m) > 3000:
-        continue
-    if m[0]["role"] not in ("user", "system"):
-        continue
-    sft.append({"messages": m})
-    n_anchor += 1
-    if n_anchor >= TARGET_ANCHOR:
-        break
-stats["anchor"] = n_anchor
-print(f"      anchor SFT: {n_anchor}")
-
-
-# ------------------------------------------------------------------ 4. write
-print("[4/4] shuffle + write …")
+# ---------------------------------------------------------- 5. smoltalk2 mix
+def take(subset, target, label):
+    """Pull `target` conversations from a smoltalk2 subset."""
+    n = 0
+    try:
+        d = load_dataset("HuggingFaceTB/smoltalk2", subset, split="train",
+                         streaming=True)
+        for row in d:
+            m = row.get("messages")
+            if not m or not (2 <= len(m) <= 6):
+                continue
+            if sum(len(x["content"]) for x in m) > 4000:
+                continue
+            if m[0]["role"] not in ("user", "system"):
+                continue
+            sft.append({"messages": m})
+            n += 1
+            if n >= target:
+                break
+    except Exception as e:
+        print(f"      {label} SKIPPED ({type(e).__name__}: {e})")
+    stats[label] = n
+    print(f"      {label}: {n}")
+    return n
+
+print("[5/5] smoltalk2 ...")
+take("SFT/Mixture_of_Thoughts_science_no_think", SCIENCE_TARGET, "science_reasoning")
+take("SFT/tulu_3_sft_personas_instruction_following_no_think", IFEVAL_TARGET,
+     "instruction_following")
+take("SFT/smoltalk_smollm3_smol_magpie_ultra_no_think", ANCHOR_TARGET, "general_anchor")
+
+
+# ------------------------------------------------------------------- write
+print("shuffle + write ...")
 random.shuffle(sft)
 with open(f"{OUT}/sft.jsonl", "w") as fh:
     for r in sft:
         fh.write(json.dumps(r) + "\n")
 
-print(f"      sft.jsonl: {len(sft)} examples")
+print(f"sft.jsonl: {len(sft)} examples")
 for k, v in stats.most_common():
-    print(f"        {v:>6}  {k}  ({v/len(sft)*100:.1f}%)")
-
-with open(f"{OUT}/stats.json", "w") as fh:
-    json.dump({"cpt_docs": len(papers), "sft_total": len(sft),
-               "sft_breakdown": dict(stats),
-               "stackexchange_by_site": dict(by_site)}, fh, indent=2)
+    print(f"  {v:>7}  {k}  ({v/len(sft)*100:.1f}%)")
+
+json.dump({"cpt_docs": cpt_lines,
+           "cpt_approx_tokens": int(cpt_chars / 4),
+           "fulltext_papers": kept,
+           "abstracts": len(papers),
+           "sft_total": len(sft),
+           "sft_breakdown": dict(stats),
+           "stackexchange_by_site": dict(by_site)},
+          open(f"{OUT}/stats.json", "w"), indent=2)
 
 from huggingface_hub import HfApi
 api = HfApi()

train_vashtra.py

@@ -7,10 +7,15 @@ Vashtra-0.6B - two-stage domain adaptation of Qwen3-0.6B for machine learning.
 Full-parameter fine-tune with 8-bit Adam. At 0.6B this beats LoRA for domain
 adaptation and still fits a free Colab T4 (~10GB of 16GB).
 
+The corpus is larger than one free Colab session can train on, so stage A takes
+a token budget. Checkpoints are pushed to the Hub, so --resume continues a run
+that was cut short in an earlier session.
+
 Run:
-    python train_vashtra.py                      # both stages, push to Hub
-    python train_vashtra.py --stage sft          # SFT only, from base
-    python train_vashtra.py --cpt_docs 40000     # shorter stage A
+    python train_vashtra.py                        # both stages, push to Hub
+    python train_vashtra.py --stage sft            # SFT only, from base
+    python train_vashtra.py --cpt_tokens 60000000  # longer stage A
+    python train_vashtra.py --resume               # continue a dropped run
 """
 import argparse, os, json
 import torch
@@ -28,11 +33,16 @@ SEED = 17
 def parse():
     p = argparse.ArgumentParser()
     p.add_argument("--stage", choices=["all", "cpt", "sft"], default="all")
-    p.add_argument("--cpt_docs", type=int, default=0, help="0 = all")
+    p.add_argument("--cpt_tokens", type=int, default=30_000_000,
+                   help="token budget for stage A; 0 = use the whole corpus. "
+                        "The corpus is ~90M tokens, which is more than one free "
+                        "Colab session can finish. 30M is about 2 hours on a T4.")
     p.add_argument("--sft_examples", type=int, default=0, help="0 = all")
     p.add_argument("--bs", type=int, default=8)
     p.add_argument("--accum", type=int, default=4)
     p.add_argument("--sft_epochs", type=float, default=2.0)
+    p.add_argument("--resume", action="store_true",
+                   help="resume from the last checkpoint pushed to the Hub")
     p.add_argument("--no_push", action="store_true")
     return p.parse_args()
 
@@ -92,9 +102,15 @@ def load_model(path):
     return m
 
 
-def targs(out, epochs, lr, args):
+def targs(out, epochs, lr, args, hub_id=None):
+    """Checkpoints go to the Hub so a dropped Colab session can be resumed."""
+    hub = {}
+    if hub_id and not args.no_push:
+        hub = dict(push_to_hub=True, hub_model_id=hub_id,
+                   hub_strategy="checkpoint", hub_private_repo=False)
     return TrainingArguments(
         output_dir=out,
+        **hub,
         per_device_train_batch_size=args.bs,
         gradient_accumulation_steps=args.accum,
         num_train_epochs=epochs,
@@ -117,12 +133,38 @@ def targs(out, epochs, lr, args):
     )
 
 
+def fetch_checkpoint(hub_id, out):
+    """Pull `last-checkpoint` from the Hub into out/ so training can resume.
+
+    Trainer's resume_from_checkpoint only looks on local disk. A fresh Colab
+    session has none, so without this a dropped run restarts from zero.
+    Returns the local checkpoint path, or None if there is nothing to resume.
+    """
+    import glob
+    from huggingface_hub import snapshot_download
+
+    local = glob.glob(os.path.join(out, "checkpoint-*")) or \
+            glob.glob(os.path.join(out, "last-checkpoint"))
+    if local:
+        print(f"[resume] using local checkpoint {local[0]}")
+        return sorted(local)[-1]
+    try:
+        p = snapshot_download(repo_id=hub_id, allow_patterns="last-checkpoint/*",
+                              local_dir=out)
+        p = os.path.join(p, "last-checkpoint")
+        if os.path.isdir(p):
+            print(f"[resume] downloaded checkpoint from {hub_id}")
+            return p
+    except Exception as e:
+        print(f"[resume] nothing to resume from {hub_id} ({type(e).__name__})")
+    return None
+
+
 # --------------------------------------------------------------- Stage A: CPT
 def stage_cpt(args, out="vashtra-cpt"):
     ds = load_dataset(CORPUS, data_files="cpt.jsonl", split="train")
-    if args.cpt_docs:
-        ds = ds.select(range(min(args.cpt_docs, len(ds))))
-    print(f"[cpt] documents: {len(ds)}")
+    ds = ds.shuffle(seed=SEED)   # mix full-text papers and abstracts together
+    print(f"[cpt] documents available: {len(ds)}")
 
     eos = tok.eos_token_id
 
@@ -142,11 +184,23 @@ def stage_cpt(args, out="vashtra-cpt"):
 
     ds = ds.map(pack, batched=True, batch_size=2000,
                 remove_columns=ds.column_names, num_proc=2)
-    print(f"[cpt] packed blocks: {len(ds)}  (~{len(ds)*SEQ/1e6:.1f}M tokens)")
-
+    print(f"[cpt] corpus packs to {len(ds)} blocks (~{len(ds)*SEQ/1e6:.0f}M tokens)")
+
+    if args.cpt_tokens:
+        keep = min(len(ds), args.cpt_tokens // SEQ)
+        if keep < len(ds):
+            print(f"[cpt] budget {args.cpt_tokens/1e6:.0f}M tokens -> training on "
+                  f"{keep} of {len(ds)} blocks. Raise --cpt_tokens (or set 0) to "
+                  f"use the whole corpus over multiple sessions.")
+        ds = ds.select(range(keep))
+    print(f"[cpt] training on {len(ds)} blocks (~{len(ds)*SEQ/1e6:.1f}M tokens)")
+
+    ckpt_repo = f"{HUB_MODEL}-cpt-ckpt"
+    resume = fetch_checkpoint(ckpt_repo, out) if args.resume else None
     model = load_model(BASE_MODEL)
-    Trainer(model=model, args=targs(out, 1, 5e-5, args),
-            train_dataset=ds, data_collator=pad_collate).train()
+    Trainer(model=model, args=targs(out, 1, 5e-5, args, hub_id=ckpt_repo),
+            train_dataset=ds, data_collator=pad_collate
+            ).train(resume_from_checkpoint=resume)
     model.save_pretrained(out); tok.save_pretrained(out)
     print(f"[cpt] saved -> {out}")
     return out
@@ -164,9 +218,13 @@ def stage_sft(args, init_path, out="vashtra-sft"):
     ds = ds.filter(lambda r: any(l != -100 for l in r["labels"]), num_proc=2)
     print(f"[sft] with supervised tokens: {len(ds)}")
 
+    ckpt_repo = f"{HUB_MODEL}-sft-ckpt"
+    resume = fetch_checkpoint(ckpt_repo, out) if args.resume else None
     model = load_model(init_path)
-    Trainer(model=model, args=targs(out, args.sft_epochs, 1e-5, args),
-            train_dataset=ds, data_collator=pad_collate).train()
+    Trainer(model=model, args=targs(out, args.sft_epochs, 1e-5, args,
+                                    hub_id=ckpt_repo),
+            train_dataset=ds, data_collator=pad_collate
+            ).train(resume_from_checkpoint=resume)
     model.save_pretrained(out); tok.save_pretrained(out)
     print(f"[sft] saved -> {out}")
     return out