File size: 2,385 Bytes
3f2dde4 | 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 | 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: # pragma: no cover - dependency gate
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)
|