| from __future__ import annotations
|
|
|
| from dataclasses import dataclass
|
| from pathlib import Path
|
|
|
|
|
| @dataclass(slots=True)
|
| class TrainerConfig:
|
| model_name: str = "OpenPeerAI/OpenPeerLLM"
|
| fallback_model_name: str = "sshleifer/tiny-gpt2"
|
| gates: int = 512
|
| steps: int = 40
|
| out_path: str = "runs/openpeer_controller.pt"
|
| train_jsonl: str | None = None
|
| device: str = "auto"
|
| demo_mode: bool = False
|
|
|
|
|
| def _build_demo_records() -> list[dict[str, str]]:
|
| return [
|
| {"prompt": "Question: 14 + 27 = ?\nAnswer:", "completion": " 41"},
|
| {"prompt": "Question: 36 + 18 = ?\nAnswer:", "completion": " 54"},
|
| {"prompt": "Question: 47 + 36 = ?\nAnswer:", "completion": " 83"},
|
| {"prompt": "Question: 19 + 8 = ?\nAnswer:", "completion": " 27"},
|
| ]
|
|
|
|
|
| def fit_controller(config: TrainerConfig) -> str:
|
| try:
|
| from transformers import AutoModelForCausalLM, AutoTokenizer
|
| from ntkmirror import ForwardFineTuner, load_jsonl_examples
|
| except ImportError as exc:
|
| raise RuntimeError(
|
| "ntkmirror mode requires transformers, torch, and ntkmirror to be installed"
|
| ) from exc
|
|
|
| train_path = Path(config.train_jsonl) if config.train_jsonl else Path("runs/demo_train.jsonl")
|
| if not train_path.exists():
|
| train_path.parent.mkdir(parents=True, exist_ok=True)
|
| import json
|
|
|
| with train_path.open("w", encoding="utf-8") as handle:
|
| for record in _build_demo_records():
|
| handle.write(json.dumps(record) + "\n")
|
|
|
| model_name = config.model_name
|
| tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
|
| try:
|
| model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
|
| except Exception:
|
| if not config.demo_mode:
|
| raise
|
| model_name = config.fallback_model_name
|
| tokenizer = AutoTokenizer.from_pretrained(model_name)
|
| model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
|
|
|
| tuner = ForwardFineTuner(model, tokenizer, gates=config.gates)
|
| tuner.fit(load_jsonl_examples(str(train_path)), steps=config.steps)
|
|
|
| out_path = Path(config.out_path)
|
| out_path.parent.mkdir(parents=True, exist_ok=True)
|
| tuner.save(str(out_path))
|
| return str(out_path)
|
|
|
