| """PUBLIC INTERFACE of Training_Module. Other modules may import ONLY this file. |
| |
| load_tokenizer(model_dir, recipe) -> (tokenizer, surgery) registers <think> </think> <|im_end|> |
| finalize_model(model, tokenizer, surgery) -> surgery embedding rows + generation ids |
| run_pretrain(model, tokenizer, plan, cfg, surgery, pretrain_path, checkpoint_dir) -> dict Stage 1 (CPT) |
| run_finetune(model, tokenizer, plan, cfg, surgery, sft_path, final_dir, adapter_dir) -> dict Stage 2 (SFT) |
| read_meta(folder) -> dict | None contents of rss_meta.json |
| upload_model(folder, cfg) -> bool push to OUTPUT_HF_REPO (HF_TOKEN) |
| download_trained(repo, dest, token) -> Path pull a trained model from the Hub |
| clear_cache(root) free disk/RAM after a successful upload |
| |
| cfg keys used: base_model_repo, output_hf_repo, hf_token, conv_ratio, tune (dict of optional overrides) |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import Dict, Optional, Tuple |
|
|
| from .internal import elevator_pretrainer as _pretrainer |
| from .internal import hf_auto_uploader as _uploader |
| from .internal import rational_sft_trainer as _sft |
| from .internal import tokenizer_surgeon as _surgeon |
|
|
| MODES = _surgeon.MODES |
|
|
|
|
| def load_tokenizer(model_dir: str, recipe: Optional[Dict] = None) -> Tuple[object, Dict]: |
| from transformers import AutoTokenizer |
| tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=False) |
| surgery = _surgeon.prepare_tokenizer(tokenizer, recipe_override=recipe) |
| if surgery["added"]: |
| print(f" [tokenizer] registered new special tokens: {', '.join(surgery['added'])}") |
| return tokenizer, surgery |
|
|
|
|
| def finalize_model(model, tokenizer, surgery: Dict) -> Dict: |
| return _surgeon.finalize_model_embeddings(model, tokenizer, surgery) |
|
|
|
|
| def run_pretrain(model, tokenizer, plan: Dict, cfg: Dict, surgery: Dict, pretrain_path, checkpoint_dir) -> Dict: |
| return _pretrainer.run_pretrain(model, tokenizer, plan, cfg, surgery, Path(pretrain_path), Path(checkpoint_dir)) |
|
|
|
|
| def run_finetune(model, tokenizer, plan: Dict, cfg: Dict, surgery: Dict, sft_path, final_dir, |
| adapter_dir: Optional[str] = None) -> Dict: |
| return _sft.run_finetune(model, tokenizer, plan, cfg, surgery, Path(sft_path), Path(final_dir), |
| adapter_dir=adapter_dir) |
|
|
|
|
| def read_meta(folder) -> Optional[Dict]: |
| meta = Path(folder) / "rss_meta.json" |
| if not meta.exists(): |
| return None |
| try: |
| return json.loads(meta.read_text(encoding="utf-8")) |
| except ValueError: |
| return None |
|
|
|
|
| def upload_model(folder, cfg: Dict) -> bool: |
| return _uploader.upload_model_folder(Path(folder), cfg["output_hf_repo"], cfg.get("hf_token"), |
| cfg["base_model_repo"]) |
|
|
|
|
| def download_trained(repo: str, dest, token: Optional[str]) -> Path: |
| return _uploader.download_trained(repo, Path(dest), token) |
|
|
|
|
| def clear_cache(root) -> None: |
| _uploader.clear_training_cache(Path(root)) |
