| """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") |
