1
0

hf-create-collection.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106
  1. #!/usr/bin/env python3
  2. from huggingface_hub import HfApi
  3. import argparse
  4. import os
  5. import sys
  6. def create_collection(title, description, private=False, namespace=None, return_slug=False):
  7. """
  8. Create a new collection on Hugging Face
  9. Args:
  10. title: Collection title
  11. description: Collection description
  12. private: Whether the collection should be private (default: False)
  13. namespace: Optional namespace (defaults to your username)
  14. Returns:
  15. Collection object if successful, None if failed
  16. """
  17. # Check if HF_TOKEN is available
  18. token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
  19. if not token:
  20. print("❌ No HF_TOKEN or HUGGINGFACE_HUB_TOKEN found in environment variables")
  21. print("Please set your Hugging Face token as an environment variable")
  22. return None
  23. # Initialize API
  24. api = HfApi()
  25. try:
  26. # Test authentication first
  27. user_info = api.whoami()
  28. if not return_slug:
  29. print(f"✅ Authenticated as: {user_info['name']}")
  30. # Create the collection
  31. if not return_slug:
  32. print(f"📚 Creating collection: '{title}'...")
  33. collection = api.create_collection(
  34. title=title,
  35. description=description,
  36. private=private,
  37. namespace=namespace
  38. )
  39. if not return_slug:
  40. print(f"✅ Collection created successfully!")
  41. print(f"📋 Collection slug: {collection.slug}")
  42. print(f"🔗 Collection URL: https://huggingface.co/collections/{collection.slug}")
  43. return collection
  44. except Exception as e:
  45. print(f"❌ Error creating collection: {e}")
  46. return None
  47. def main():
  48. # This script requires that the environment variable HF_TOKEN is set with your
  49. # Hugging Face API token.
  50. api = HfApi()
  51. parser = argparse.ArgumentParser(description='Create a Huggingface Collection')
  52. parser.add_argument('--name', '-n', help='The name/title of the Collection', required=True)
  53. parser.add_argument('--description', '-d', help='The description for the Collection', required=True)
  54. parser.add_argument('--namespace', '-ns', help='The namespace to add the Collection to', required=True)
  55. parser.add_argument('--private', '-p', help='Create a private Collection', action='store_true') # Fixed
  56. parser.add_argument('--return-slug', '-s', help='Only output the collection slug', action='store_true') # Fixed
  57. args = parser.parse_args()
  58. name = args.name
  59. description = args.description
  60. private = args.private
  61. namespace = args.namespace
  62. return_slug = args.return_slug
  63. if not return_slug:
  64. print("🚀 Creating Hugging Face Collection")
  65. print(f"Title: {name}")
  66. print(f"Description: {description}")
  67. print(f"Namespace: {namespace}")
  68. print(f"Private: {private}")
  69. collection = create_collection(
  70. title=name,
  71. description=description,
  72. private=private,
  73. namespace=namespace,
  74. return_slug=return_slug
  75. )
  76. if collection:
  77. if return_slug:
  78. print(collection.slug)
  79. else:
  80. print("\n🎉 Collection created successfully!")
  81. print(f"Use this slug to add models: {collection.slug}")
  82. else:
  83. print("\n❌ Failed to create collection")
  84. sys.exit(1)
  85. if __name__ == "__main__":
  86. main()