File size: 3,250 Bytes
9d614f1 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """Auto-push trained model to the Hugging Face Hub and clear local training cache (PRIVATE)."""
from __future__ import annotations
import gc
import shutil
from pathlib import Path
from typing import Optional
PLACEHOLDER_MARK = "PUT_YOUR_TOKEN"
_CARD = """---
language: [hi, en]
library_name: transformers
base_model: {base}
tags: [rss-ai, hinglish, thinking-modes, lora]
---
# RSS AI
Trained with the RSS AI Engine (continual pre-training + rational SFT).
Modes: `thinking`, `non_thinking`, `max_thinking` (selected through `[MODE: ...]` in the system prompt).
Base model: `{base}`
"""
def _valid(token: Optional[str]) -> bool:
return bool(token) and PLACEHOLDER_MARK not in token
def upload_model_folder(folder: Path, repo_id: str, token: Optional[str], base_model: str,
private: bool = True) -> bool:
folder = Path(folder)
if not _valid(token):
print(" [upload] HF_TOKEN missing/placeholder -> upload skipped (model stays in "
f"{folder}). Put a WRITE token in .env and run: main.py --mode train_all --upload-only")
return False
if not (folder / "rss_meta.json").exists():
print(f" [upload] {folder} has no rss_meta.json -> refusing to upload an unfinished model")
return False
from huggingface_hub import HfApi
api = HfApi(token=token)
try:
api.create_repo(repo_id, repo_type="model", private=private, exist_ok=True)
(folder / "README.md").write_text(_CARD.format(base=base_model), encoding="utf-8")
print(f" [upload] pushing {folder} -> https://huggingface.co/{repo_id} ...")
api.upload_folder(folder_path=str(folder), repo_id=repo_id, repo_type="model",
commit_message="RSS AI: trained model + tokenizer",
ignore_patterns=[".rss_overlay/*", "*.incomplete"])
print(" [upload] SUCCESS - model is safe on the Hub")
return True
except Exception as exc:
print(f" [upload] FAILED ({type(exc).__name__}: {exc}). Local copy kept in {folder}")
return False
def download_trained(repo_id: str, dest: Path, token: Optional[str]) -> Path:
from huggingface_hub import snapshot_download
dest = Path(dest)
print(f" [download] pulling trained model {repo_id} -> {dest}")
snapshot_download(repo_id=repo_id, local_dir=str(dest), token=token if _valid(token) else None)
if not (dest / "rss_meta.json").exists():
raise RuntimeError(f"{repo_id} does not look like an RSS AI export (rss_meta.json missing)")
return dest
def clear_training_cache(root: Path) -> None:
"""Free disk + RAM after a successful upload. The final trained_model/ folder is kept."""
root = Path(root)
for name in ("checkpoints", "offload_disk"):
target = root / name
if target.exists():
shutil.rmtree(target, ignore_errors=True)
print(f" [cleanup] removed {target}/")
tmp = root / "cleaned_data" / "_repo_text.tmp.jsonl"
if tmp.exists():
tmp.unlink()
gc.collect()
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception:
pass
print(" [cleanup] RAM/VRAM released") |