hf-create-model.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. #!/usr/bin/env python3
  2. from huggingface_hub import HfApi
  3. import argparse
  4. # This script requires that the environment variable HF_TOKEN is set with your
  5. # Hugging Face API token.
  6. api = HfApi()
  7. def load_template_and_substitute(template_path, **kwargs):
  8. try:
  9. with open(template_path, 'r', encoding='utf-8') as f:
  10. template_content = f.read()
  11. return template_content.format(**kwargs)
  12. except FileNotFoundError:
  13. print(f"Template file '{template_path}' not found!")
  14. return None
  15. except KeyError as e:
  16. print(f"Missing template variable: {e}")
  17. return None
  18. parser = argparse.ArgumentParser(description='Create a new Hugging Face model repository')
  19. parser.add_argument('--model-name', '-m', help='Name for the model', required=True)
  20. parser.add_argument('--namespace', '-ns', help='Namespace to add the model to', required=True)
  21. parser.add_argument('--org-base-model', '-b', help='Original Base model name', default="")
  22. parser.add_argument('--no-card', action='store_true', help='Skip creating model card')
  23. parser.add_argument('--private', '-p', action='store_true', help='Create private model')
  24. parser.add_argument('--embedding', '-e', action='store_true', help='Use embedding model card template')
  25. parser.add_argument('--dry-run', '-d', action='store_true', help='Print repository info and template without creating repository')
  26. args = parser.parse_args()
  27. repo_id = f"{args.namespace}/{args.model_name}-GGUF"
  28. print("Repository ID: ", repo_id)
  29. repo_url = None
  30. if not args.dry_run:
  31. repo_url = api.create_repo(
  32. repo_id=repo_id,
  33. repo_type="model",
  34. private=args.private,
  35. exist_ok=False
  36. )
  37. if not args.no_card:
  38. if args.embedding:
  39. template_path = "scripts/embedding/modelcard.template"
  40. else:
  41. template_path = "scripts/causal/modelcard.template"
  42. print("Template path: ", template_path)
  43. model_card_content = load_template_and_substitute(
  44. template_path,
  45. model_name=args.model_name,
  46. namespace=args.namespace,
  47. base_model=args.org_base_model,
  48. )
  49. if args.dry_run:
  50. print("\nTemplate Content:\n")
  51. print(model_card_content)
  52. else:
  53. if model_card_content:
  54. api.upload_file(
  55. path_or_fileobj=model_card_content.encode('utf-8'),
  56. path_in_repo="README.md",
  57. repo_id=repo_id
  58. )
  59. print("Model card created successfully.")
  60. else:
  61. print("Failed to create model card.")
  62. if not args.dry_run and repo_url:
  63. print(f"Repository created: {repo_url}")