code-review-agent / upload_model.py
100XZX001's picture
Create upload_model.py
a21f6ae verified
import os
from huggingface_hub import HfApi, create_repo, Repository
from transformers import AutoModelForCausalLM, AutoTokenizer
# Set your Hugging Face username and token (with write permission)
USERNAME = "your-username" # replace
TOKEN = "hf_xxxxx" # replace
REPO_ID = f"{USERNAME}/code-review-agent"
LOCAL_DIR = "./code-review-agent"
# 1. Create repo on Hub (if it doesn't exist)
try:
create_repo(REPO_ID, token=TOKEN, private=False, exist_ok=True)
print(f"Created/verified repo: {REPO_ID}")
except Exception as e:
print(e)
# 2. Clone the repo locally
if not os.path.exists(LOCAL_DIR):
repo = Repository(local_dir=LOCAL_DIR, clone_from=REPO_ID, use_auth_token=TOKEN)
else:
# if already exists, just pull latest
repo = Repository(local_dir=LOCAL_DIR, clone_from=REPO_ID, use_auth_token=TOKEN)
repo.git_pull()
# 3. Download the model and tokenizer
print("Downloading distilgpt2 ...")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
# 4. Save them into the local repo
model.save_pretrained(LOCAL_DIR)
tokenizer.save_pretrained(LOCAL_DIR)
# 5. Create model card (README.md)
readme = """---
language: en
license: mit
tags:
- code-review
- agent
- text-generation
---
# Code Review Agent Model
This is a small language model (distilgpt2) that can be fine‑tuned to generate helpful code review comments.
It is a starting point for building an AI that reviews pull requests.
"""
with open(os.path.join(LOCAL_DIR, "README.md"), "w") as f:
f.write(readme)
# 6. Set up Git LFS for large files
# The Repository class may not handle LFS automatically, so we'll do it manually
import subprocess
subprocess.run(["git", "lfs", "install"], cwd=LOCAL_DIR)
subprocess.run(["git", "lfs", "track", "*.bin", "*.safetensors"], cwd=LOCAL_DIR)
subprocess.run(["git", "add", ".gitattributes"], cwd=LOCAL_DIR)
# 7. Commit and push
repo.git_add()
repo.git_commit("Initial upload of distilgpt2 model")
repo.git_push()
print("Model pushed to:", f"https://huggingface.co/{REPO_ID}")