|
|
""" |
|
|
Upload trained model to Hugging Face Hub |
|
|
""" |
|
|
import argparse |
|
|
import sys |
|
|
import os |
|
|
from huggingface_hub import HfApi, create_repo |
|
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification |
|
|
|
|
|
|
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) |
|
|
|
|
|
|
|
|
def upload_to_hub( |
|
|
model_path: str, |
|
|
repo_name: str, |
|
|
organization: str = None, |
|
|
private: bool = False |
|
|
): |
|
|
""" |
|
|
Upload model to Hugging Face Hub. |
|
|
|
|
|
Args: |
|
|
model_path: Path to the trained model |
|
|
repo_name: Name for the repository on Hugging Face Hub |
|
|
organization: Organization name (optional) |
|
|
private: Whether to make the repository private |
|
|
""" |
|
|
print("=" * 60) |
|
|
print("Uploading Model to Hugging Face Hub") |
|
|
print("=" * 60) |
|
|
|
|
|
|
|
|
if organization: |
|
|
repo_id = f"{organization}/{repo_name}" |
|
|
else: |
|
|
repo_id = repo_name |
|
|
|
|
|
print(f"\nRepository: {repo_id}") |
|
|
print(f"Private: {private}") |
|
|
|
|
|
|
|
|
print("\n[1/3] Loading model...") |
|
|
try: |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path) |
|
|
model = AutoModelForSequenceClassification.from_pretrained(model_path) |
|
|
print("β Model loaded successfully") |
|
|
except Exception as e: |
|
|
print(f"β Error loading model: {e}") |
|
|
return |
|
|
|
|
|
|
|
|
print("\n[2/3] Creating repository...") |
|
|
try: |
|
|
create_repo( |
|
|
repo_id=repo_id, |
|
|
repo_type="model", |
|
|
exist_ok=True, |
|
|
private=private |
|
|
) |
|
|
print(f"β Repository created/verified: {repo_id}") |
|
|
except Exception as e: |
|
|
print(f"β Error creating repository: {e}") |
|
|
print("\nMake sure you're logged in:") |
|
|
print(" huggingface-cli login") |
|
|
return |
|
|
|
|
|
|
|
|
print("\n[3/3] Uploading model and tokenizer...") |
|
|
try: |
|
|
model.push_to_hub(repo_id) |
|
|
tokenizer.push_to_hub(repo_id) |
|
|
print("β Upload complete!") |
|
|
except Exception as e: |
|
|
print(f"β Error uploading: {e}") |
|
|
return |
|
|
|
|
|
print("\n" + "=" * 60) |
|
|
print("Success! π") |
|
|
print("=" * 60) |
|
|
print(f"\nYour model is now available at:") |
|
|
print(f"https://huggingface.co/{repo_id}") |
|
|
|
|
|
print("\nTo use your model:") |
|
|
print(f""" |
|
|
from transformers import AutoTokenizer, AutoModelForSequenceClassification |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("{repo_id}") |
|
|
model = AutoModelForSequenceClassification.from_pretrained("{repo_id}") |
|
|
""") |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
parser = argparse.ArgumentParser(description="Upload model to Hugging Face Hub") |
|
|
parser.add_argument( |
|
|
"--model-path", |
|
|
type=str, |
|
|
default="./results/final_model", |
|
|
help="Path to the trained model" |
|
|
) |
|
|
parser.add_argument( |
|
|
"--repo-name", |
|
|
type=str, |
|
|
required=True, |
|
|
help="Name for the repository on Hugging Face Hub" |
|
|
) |
|
|
parser.add_argument( |
|
|
"--organization", |
|
|
type=str, |
|
|
default=None, |
|
|
help="Organization name (optional)" |
|
|
) |
|
|
parser.add_argument( |
|
|
"--private", |
|
|
action="store_true", |
|
|
help="Make the repository private" |
|
|
) |
|
|
args = parser.parse_args() |
|
|
|
|
|
print("\nBefore uploading, make sure you:") |
|
|
print("1. Have a Hugging Face account") |
|
|
print("2. Are logged in: huggingface-cli login") |
|
|
print("3. Have reviewed the model card (MODEL_CARD.md)") |
|
|
|
|
|
response = input("\nProceed with upload? (yes/no): ") |
|
|
if response.lower() in ['yes', 'y']: |
|
|
upload_to_hub( |
|
|
args.model_path, |
|
|
args.repo_name, |
|
|
args.organization, |
|
|
args.private |
|
|
) |
|
|
else: |
|
|
print("Upload cancelled.") |
|
|
|