File size: 2,610 Bytes
3ce19a2 | 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 | import os
import json
from huggingface_hub import HfApi
def get_param_count():
# The default log path from train_cifar10.sh and train.py
log_file = "./runs/cifar10_rtm/cifar10/log.jsonl"
if not os.path.exists(log_file):
return "unknown"
with open(log_file, "r") as f:
for line in f:
try:
data = json.loads(line)
if "message" in data and "Number of parameters in IMLE:" in data["message"]:
parts = data["message"].split("Number of parameters in IMLE:")
if len(parts) > 1:
params_str = parts[1].strip()
return params_str
except:
pass
return "unknown"
def main():
api = HfApi()
# 1. Get parameter count and format repo name
param_count = get_param_count()
repo_id = f"JerMa88/rtm-latent-refinement-cifar10-{param_count}params"
print(f"Creating repository: {repo_id}")
# Create the repo (it will use the implicitly logged-in user)
try:
api.create_repo(repo_id, exist_ok=True)
except Exception as e:
print(f"Error creating repo: {e}")
# Upload training checkpoints and logs
run_dir = "./runs/cifar10_rtm/cifar10"
if os.path.exists(run_dir):
print(f"Uploading run directory {run_dir}...")
api.upload_folder(
folder_path=run_dir,
repo_id=repo_id,
path_in_repo="runs/cifar10",
)
else:
print(f"Run directory {run_dir} not found! Check if training finished successfully.")
# Upload current repository code
print("Uploading repository code...")
api.upload_folder(
folder_path=".",
repo_id=repo_id,
path_in_repo="code",
ignore_patterns=["runs/*", "datasets/*", "venv/*", ".git/*"]
)
# Upload README to Hugging Face
readme_content = f"""
# RTM Latent Refinement (Fast Variant)
This model was trained with a reduced architecture ({param_count} parameters) compared to the original paper's baseline to speed up training.
- **Dataset**: CIFAR-10
- **W&B Logs**: Check the W&B project `rtm-latent-refinement` under the account `JerMa88`.
- **Source Code**: Included in the `code/` directory.
"""
with open("HF_README.md", "w") as f:
f.write(readme_content)
api.upload_file(
path_or_fileobj="HF_README.md",
path_in_repo="README.md",
repo_id=repo_id,
)
print("Upload complete! You can view the model on Hugging Face.")
if __name__ == "__main__":
main()
|