File size: 3,110 Bytes
9d614f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | """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)) |