| """Integration layer: the ONLY place that wires the modules together. |
| |
| It imports exactly four public interfaces - data_api, hardware_api, training_api, serving_api - and never |
| touches anything inside an `internal/` folder. |
| |
| Purge policy (PURGE_OLD_MODEL_ON_START=true or --fresh-model): |
| train_all / pretrain start a NEW training lineage -> old artifacts are purged and the base model re-downloaded |
| clean / finetune / serve never purge (finetune needs the stage-1 checkpoint, |
| serve needs the trained model) unless --fresh-model is given |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import sys |
| import time |
| from pathlib import Path |
| from typing import Dict, List, Optional |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| if str(ROOT) not in sys.path: |
| sys.path.insert(0, str(ROOT)) |
|
|
| from Data_Module import data_api |
| from Hardware_Module import hardware_api |
| from Serving_Module import serving_api |
| from Training_Module import training_api |
|
|
| DEFAULT_SOURCES = [{"repo": "Biraboro/Youtop-trenscprit", "limit": 150}] |
|
|
|
|
| |
| def _path(root: Path, value: str) -> Path: |
| p = Path(value).expanduser() |
| return p if p.is_absolute() else (root / p).resolve() |
|
|
|
|
| def _parse_sources(raw: str) -> List[Dict]: |
| raw = (raw or "").strip() |
| if not raw: |
| return list(DEFAULT_SOURCES) |
| try: |
| data = json.loads(raw) |
| except ValueError: |
| data = [{"repo": part.strip(), "limit": 150} for part in raw.split(",") if part.strip()] |
| if isinstance(data, dict): |
| data = [data] |
| out = [] |
| for item in data if isinstance(data, list) else []: |
| if isinstance(item, str): |
| out.append({"repo": item, "limit": 150}) |
| elif isinstance(item, dict) and item.get("repo"): |
| out.append({"repo": str(item["repo"]), "limit": item.get("limit", 150)}) |
| if not out: |
| print("[Config] SOURCE_REPOS could not be parsed -> using the default repo") |
| return list(DEFAULT_SOURCES) |
| return out |
|
|
|
|
| def _flag(value: str, default: bool = False) -> bool: |
| return str(value if value is not None else default).strip().lower() in ("1", "true", "yes", "y", "on") |
|
|
|
|
| def _float(value: str, default: float) -> float: |
| try: |
| return float(value) |
| except (TypeError, ValueError): |
| return default |
|
|
|
|
| def build_config(args=None) -> Dict: |
| env_file = ROOT / ".env" |
| try: |
| from dotenv import load_dotenv |
| load_dotenv(env_file) |
| except ImportError: |
| if env_file.exists(): |
| for line in env_file.read_text(encoding="utf-8").splitlines(): |
| if "=" in line and not line.strip().startswith("#"): |
| key, _, val = line.partition("=") |
| os.environ.setdefault(key.strip(), val.strip().strip("'\"")) |
| env = os.environ.get |
| tune: Dict = {} |
| raw_tune = env("TUNE_JSON", "").strip() |
| if raw_tune: |
| try: |
| tune = json.loads(raw_tune) |
| except ValueError: |
| print("[Config] TUNE_JSON is not valid JSON - ignored") |
| extra = [e.strip() for e in (env("EXTRA_BLOCKED_EXTENSIONS", "") or "").replace(";", ",").split(",") if e.strip()] |
| cfg: Dict = { |
| "root": str(ROOT), |
| "hf_token": (env("HF_TOKEN", "") or "").strip(), |
| "telegram_token": (env("TELEGRAM_BOT_TOKEN", "") or "").strip(), |
| "telegram_chat_id": (env("TELEGRAM_CHAT_ID", "") or "").strip(), |
| "base_model_repo": (env("BASE_MODEL_REPO", "Qwen/Qwen2.5-0.5B-Instruct") or "").strip(), |
| "output_hf_repo": (env("OUTPUT_HF_REPO", "Biraboro/rss-ai-trained") or "").strip(), |
| "source_repos": _parse_sources(env("SOURCE_REPOS", "")), |
| "chat_file": str(_path(ROOT, env("LOCAL_CHAT_FILE", "./chat.jsonl") or "./chat.jsonl")), |
| "json_quota": _float(env("JSON_PRETRAIN_QUOTA", "0.15"), 0.15), |
| "conv_ratio": _float(env("PRETRAIN_CONV_RATIO", "0.25"), 0.25), |
| "purge_old_model": _flag(env("PURGE_OLD_MODEL_ON_START", "true"), True), |
| "extra_blocked_extensions": extra, |
| "tune": tune, |
| "force_tier": None, |
| "cpu_threads": None, |
| "limit_override": None, |
| } |
| if args is not None: |
| if getattr(args, "limit", None): |
| cfg["limit_override"] = args.limit |
| if getattr(args, "tier", None): |
| cfg["force_tier"] = args.tier |
| if getattr(args, "threads", None): |
| cfg["cpu_threads"] = args.threads |
| return cfg |
|
|
|
|
| def _dirs(cfg: Dict) -> Dict[str, Path]: |
| root = Path(cfg["root"]) |
| return {"stage1": root / "checkpoints" / "stage1", "final": root / "trained_model", "root": root} |
|
|
|
|
| def _valid_stage1_checkpoint(folder: Path) -> bool: |
| """Validate Stage-1 metadata AND the actual weight payload before any resume decision.""" |
| folder = Path(folder) |
| meta = training_api.read_meta(folder) |
| if not meta or int(meta.get("stage", 0)) != 1: |
| return False |
| kind = meta.get("kind") |
| if kind == "full": |
| if not (folder / "config.json").is_file(): |
| return False |
| single = folder / "model.safetensors" |
| index = folder / "model.safetensors.index.json" |
| if single.is_file() and single.stat().st_size > 0: |
| return True |
| if index.is_file(): |
| try: |
| weight_map = json.loads(index.read_text(encoding="utf-8")).get("weight_map", {}) |
| shards = {str(name) for name in weight_map.values()} |
| return bool(shards) and all((folder / name).is_file() and (folder / name).stat().st_size > 0 |
| for name in shards) |
| except Exception: |
| return False |
| return False |
| if kind == "adapter": |
| adapter = folder / (meta.get("adapter_dir") or "adapter") |
| if not (adapter / "adapter_config.json").is_file(): |
| return False |
| return any(path.is_file() and path.stat().st_size > 0 for path in |
| (adapter / "adapter_model.safetensors", adapter / "adapter_model.bin")) |
| return False |
|
|
|
|
| def _banner(text: str) -> None: |
| print("\n" + "=" * 72 + f"\n {text}\n" + "=" * 72) |
|
|
|
|
| |
| def _should_purge(cfg: Dict, args, mode: str) -> bool: |
| |
| |
| return False |
|
|
|
|
| def _load(model_dir: str, cfg: Dict, recipe: Optional[Dict] = None): |
| """profile -> plan -> tokenizer surgery -> load -> finalize. Returns (model, tokenizer, plan, surgery).""" |
| prof = hardware_api.profile_system(model_dir, cfg) |
| plan = hardware_api.plan_execution(prof, cfg) |
| tokenizer, surgery = training_api.load_tokenizer(model_dir, recipe=recipe) |
| model = hardware_api.load_model(model_dir, plan, cfg, row_recipe=surgery["recipe"] or None) |
| training_api.finalize_model(model, tokenizer, surgery) |
| return model, tokenizer, plan, surgery |
|
|
|
|
| def _ensure_clean(cfg: Dict) -> None: |
| if not data_api.cleaned_data_ready(cfg): |
| print("[Orchestrator] cleaned data not found -> running the clean step first") |
| data_api.run_clean(cfg) |
|
|
|
|
| |
| def mode_clean(cfg: Dict, args=None) -> Dict: |
| _banner("MODE clean - stream + clean data") |
| data_api.ensure_chat_file(cfg) |
| return data_api.run_clean(cfg) |
|
|
|
|
| def mode_pretrain(cfg: Dict, args=None, model_dir: Optional[str] = None) -> Dict: |
| _banner("MODE pretrain - Stage 1 continual pre-training") |
| dirs = _dirs(cfg) |
| if _valid_stage1_checkpoint(dirs["stage1"]): |
| meta = training_api.read_meta(dirs["stage1"]) or {} |
| print(f"[Orchestrator] valid Stage-1 checkpoint found -> SKIP Stage 1; resume from {dirs['stage1']}", |
| flush=True) |
| return {"kind": meta.get("kind"), "checkpoint": str(dirs["stage1"]), "skipped": True} |
| if dirs["stage1"].exists(): |
| raise RuntimeError(f"Stage-1 directory exists but failed validation; refusing to overwrite/delete it: {dirs['stage1']}") |
| data_api.ensure_chat_file(cfg) |
| _ensure_clean(cfg) |
| dirs = _dirs(cfg) |
| if model_dir is None: |
| model_dir = hardware_api.prepare_base_model(cfg, fresh=False) |
| paths = data_api.cleaned_paths(cfg) |
| model, tokenizer, plan, surgery = _load(model_dir, cfg) |
| try: |
| result = training_api.run_pretrain(model, tokenizer, plan, cfg, surgery, paths["pretrain"], dirs["stage1"]) |
| finally: |
| del model |
| hardware_api.release_memory() |
| print(f"[Orchestrator] Stage 1 done (loss {result['loss']:.4f}) -> {result['checkpoint']}") |
| return result |
|
|
|
|
| def mode_finetune(cfg: Dict, args=None, model_dir: Optional[str] = None) -> Dict: |
| _banner("MODE finetune - Stage 2 rational SFT") |
| data_api.ensure_chat_file(cfg) |
| _ensure_clean(cfg) |
| dirs = _dirs(cfg) |
| paths = data_api.cleaned_paths(cfg) |
| meta = training_api.read_meta(dirs["stage1"]) |
| if dirs["stage1"].exists() and not _valid_stage1_checkpoint(dirs["stage1"]): |
| raise RuntimeError(f"Stage-1 directory exists but checkpoint metadata/weights are missing/invalid; refusing fallback or overwrite: {dirs['stage1']}") |
| adapter_dir: Optional[str] = None |
| recipe = None |
| if meta is None: |
| print("[Orchestrator] no Stage-1 checkpoint found -> fine-tuning directly from the base model") |
| source = model_dir or hardware_api.prepare_base_model(cfg, fresh=False) |
| elif meta["kind"] == "adapter": |
| source = model_dir or hardware_api.prepare_base_model(cfg, fresh=False) |
| adapter_dir = str(dirs["stage1"] / meta["adapter_dir"]) |
| recipe = meta.get("row_recipe") or None |
| print(f"[Orchestrator] continuing the Stage-1 LoRA adapter on top of {cfg['base_model_repo']}") |
| else: |
| source = str(dirs["stage1"]) |
| recipe = meta.get("row_recipe") or None |
| print(f"[Orchestrator] loading the Stage-1 checkpoint {source}") |
| model, tokenizer, plan, surgery = _load(source, cfg, recipe=recipe) |
| try: |
| result = training_api.run_finetune(model, tokenizer, plan, cfg, surgery, paths["sft"], dirs["final"], |
| adapter_dir=adapter_dir) |
| finally: |
| del model |
| hardware_api.release_memory() |
| print(f"[Orchestrator] Stage 2 done (loss {result['loss']:.4f}) -> {result['final_dir']}") |
| return result |
|
|
|
|
| def mode_upload(cfg: Dict, args=None) -> bool: |
| final = _dirs(cfg)["final"] |
| ok = training_api.upload_model(final, cfg) |
| if ok: |
| training_api.clear_cache(cfg["root"]) |
| return ok |
|
|
|
|
| def mode_train_all(cfg: Dict, args=None) -> Dict: |
| _banner("MODE train_all - checkpoint-aware clean -> pretrain/skip -> finetune -> upload") |
| started = time.time() |
| data_api.ensure_chat_file(cfg) |
| dirs = _dirs(cfg) |
| stage1_valid = _valid_stage1_checkpoint(dirs["stage1"]) |
| if dirs["stage1"].exists() and not stage1_valid: |
| raise RuntimeError(f"Stage-1 directory exists but is not a valid checkpoint; preserving it and refusing overwrite: {dirs['stage1']}") |
|
|
| if stage1_valid: |
| stage1_meta = training_api.read_meta(dirs["stage1"]) or {} |
| print(f"[1/5] Stage-1 checkpoint verified ({stage1_meta.get('kind')}) -> preserving; no purge/download", flush=True) |
| model_dir = None |
| if stage1_meta.get("kind") == "adapter": |
| model_dir = hardware_api.prepare_base_model(cfg, fresh=False) |
| else: |
| print("[1/5] base model (reuse verified local copy; destructive purge disabled)", flush=True) |
| model_dir = hardware_api.prepare_base_model(cfg, fresh=False) |
|
|
| print("[2/5] cleaning data", flush=True) |
| if data_api.cleaned_data_ready(cfg): |
| print("[Data] cleaned SFT/pretrain files already exist -> reuse", flush=True) |
| else: |
| data_api.run_clean(cfg) |
|
|
| if stage1_valid: |
| print(f"[3/5] stage 1 -> SKIPPED (verified checkpoint {dirs['stage1']})", flush=True) |
| else: |
| print("[3/5] stage 1 -> train once (checkpoint path is write-once)", flush=True) |
| mode_pretrain(cfg, args, model_dir=model_dir) |
|
|
| print("[4/5] stage 2 -> rational SFT", flush=True) |
| result = mode_finetune(cfg, args, model_dir=model_dir) |
| uploaded = False |
| if getattr(args, "skip_upload", False): |
| print("[5/5] upload skipped (--skip-upload)", flush=True) |
| else: |
| print("[5/5] uploading to Hugging Face", flush=True) |
| uploaded = mode_upload(cfg, args) |
| result["uploaded"] = uploaded |
| print(f"\n[Orchestrator] train_all finished in {(time.time() - started) / 60:.1f} min " |
| f"(uploaded={uploaded}). Chat with it: python RSS_AI_Engine/main.py --mode serve", flush=True) |
| return result |
|
|
|
|
| def mode_serve(cfg: Dict, args=None) -> str: |
| _banner("MODE serve - load trained model and start the chat front-end") |
| data_api.ensure_chat_file(cfg) |
| trained = serving_api.resolve_trained_dir(cfg, fetch_remote=training_api.download_trained) |
| meta = training_api.read_meta(trained) or {} |
| if meta.get("has_adapter") and not meta.get("base_is_self"): |
| weights_dir = hardware_api.prepare_base_model(cfg, fresh=False) |
| else: |
| weights_dir = str(trained) |
| recipe = meta.get("row_recipe") or None |
| cfg["active_vocab_file"] = str(trained / "active_vocab_ids.json") |
| model, tokenizer, plan, surgery = _load(weights_dir, cfg, recipe=recipe) |
| if weights_dir != str(trained): |
| tokenizer, surgery = training_api.load_tokenizer(str(trained), recipe=recipe) |
| model = serving_api.prepare_model(model, trained, meta) |
| engine = serving_api.build_engine(model, tokenizer, cfg) |
| return serving_api.start(engine, cfg) |
|
|
|
|
| def run(mode: str, cfg: Dict, args=None): |
| table = {"clean": mode_clean, "pretrain": mode_pretrain, "finetune": mode_finetune, |
| "train_all": mode_train_all, "serve": mode_serve} |
| if getattr(args, "upload_only", False): |
| return mode_upload(cfg, args) |
| return table[mode](cfg, args) |
|
|
