Blame for RSS_AI_Engine/main.py

#!/usr/bin/env python3
"""RSS AI Engine - application entry point.

    python RSS_AI_Engine/main.py --mode clean       stream + clean data only
    python RSS_AI_Engine/main.py --mode pretrain    Stage-1 continual pre-training only
    python RSS_AI_Engine/main.py --mode finetune    Stage-2 rational SFT only
    python RSS_AI_Engine/main.py --mode train_all   purge -> fresh HF model -> clean -> pretrain -> finetune -> upload
    python RSS_AI_Engine/main.py --mode serve       load the trained model -> Telegram bot or CLI chat

Extras:  --fresh-model  --upload-only  --skip-upload  --limit 150|full  --tier 1|2|3  --threads N
"""
from __future__ import annotations

import argparse
import importlib.util
import os
import signal
import subprocess
import sys
from datetime import datetime
from pathlib import Path

# Set native tokenizer/OpenMP limits BEFORE importing torch/transformers/tokenizers.
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["MKL_NUM_THREADS"] = "2"

ROOT = Path(__file__).resolve().parent

ENV_TEMPLATE = """HF_TOKEN="hf_PUT_YOUR_TOKEN_HERE"
TELEGRAM_BOT_TOKEN=""
TELEGRAM_CHAT_ID=""
BASE_MODEL_REPO="Qwen/Qwen2.5-0.5B-Instruct"
OUTPUT_HF_REPO="Biraboro/rss-ai-trained"
SOURCE_REPOS='[{"repo": "Biraboro/Youtop-trenscprit", "limit": 150}]'
LOCAL_CHAT_FILE="./chat.jsonl"
JSON_PRETRAIN_QUOTA="0.15"
PRETRAIN_CONV_RATIO="0.25"
PURGE_OLD_MODEL_ON_START="true"
EXTRA_BLOCKED_EXTENSIONS=""
"""

# (import name, pip name)
REQUIRED = [("torch", "torch"), ("transformers", "transformers"), ("accelerate", "accelerate"),
            ("peft", "peft"), ("datasets", "datasets"), ("huggingface_hub", "huggingface_hub"),
            ("safetensors", "safetensors"), ("psutil", "psutil"), ("dotenv", "python-dotenv"),
            ("requests", "requests")]


def ensure_env_file() -> None:
    env = ROOT / ".env"
    if not env.exists():
        env.write_text(ENV_TEMPLATE, encoding="utf-8")
        print(f"[Setup] created {env} - put your HF_TOKEN (write access) in it before uploading")


def _pip(package: str) -> None:
    base = [sys.executable, "-m", "pip", "install", "-q", package]
    if subprocess.call(base) != 0:
        # PEP 668 "externally managed" interpreters (Debian/Ubuntu system python) need this flag
        subprocess.check_call(base + ["--break-system-packages"])


def _cuda_present() -> bool:
    try:
        import torch
        return bool(torch.cuda.is_available())
    except Exception:
        return False


def ensure_dependencies() -> None:
    missing = [pip for mod, pip in REQUIRED if importlib.util.find_spec(mod) is None]
    for package in missing:
        print(f"[Setup] installing {package} ...")
        _pip(package)
    if _cuda_present() and importlib.util.find_spec("bitsandbytes") is None:
        print("[Setup] CUDA detected -> installing bitsandbytes (4-bit NF4)")
        try:
            _pip("bitsandbytes")
        except Exception as exc:
            print(f"[Setup] bitsandbytes install failed ({exc}); GPU tier will run without 4-bit")
    if not missing:
        print("[Setup] all dependencies present")


def parse_args(argv=None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="RSS AI Engine", formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("--mode", choices=["clean", "pretrain", "finetune", "train_all", "serve"], default="serve",
                        help="clean | pretrain | finetune | train_all | serve   (default: serve)")
    parser.add_argument("--fresh-model", action="store_true",
                        help="purge old checkpoints/offload files and download the base model again")
    parser.add_argument("--upload-only", action="store_true", help="only upload ./trained_model to OUTPUT_HF_REPO")
    parser.add_argument("--skip-upload", action="store_true", help="train_all without the final upload")
    parser.add_argument("--limit", default=None, help="override every repo limit: a number or 'full'")
    parser.add_argument("--tier", type=int, choices=[1, 2, 3], default=None, help="force execution tier")
    parser.add_argument("--threads", type=int, default=None, help="CPU threads for torch")
    parser.add_argument("--no-install", action="store_true", help="do not auto-install missing dependencies")
    return parser.parse_args(argv)


def _handle_sigterm(signum, _frame) -> None:
    stamp = datetime.now().astimezone().isoformat(timespec="seconds")
    print(f"\n[Runtime] received SIGTERM ({signum}) at {stamp}; "
          "external supervisor/timeout or user termination is possible; sender is not identifiable here.",
          flush=True)
    raise SystemExit(128 + int(signum))


def main(argv=None) -> int:
    try:
        signal.signal(signal.SIGTERM, _handle_sigterm)
    except (AttributeError, ValueError):
        pass
    args = parse_args(argv)
    ensure_env_file()
    if not args.no_install:
        ensure_dependencies()
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
    try:
        import torch
        torch.set_num_threads(2)
        try:
            torch.set_num_interop_threads(2)
        except RuntimeError:
            # Inter-op pool may already be initialised by an embedding host.
            pass
        print(f"[Runtime] torch threads={torch.get_num_threads()} interop={torch.get_num_interop_threads()}",
              flush=True)
    except ImportError:
        pass
    from Integration_Layer import engine_orchestrator as orchestrator
    cfg = orchestrator.build_config(args)
    try:
        orchestrator.run(args.mode, cfg, args)
    except FileNotFoundError as exc:
        print(f"\n[RSS] {exc}")
        return 2
    except KeyboardInterrupt:
        print("\n[RSS] interrupted")
        return 130
    return 0


if __name__ == "__main__":
    sys.exit(main())