| |
| """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 |
|
|
| |
| 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="" |
| """ |
|
|
| |
| 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: |
| |
| 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: |
| |
| 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()) |
|
|
