diff --git a/.env.example b/.env.example index f3b8bcc..9b05ff0 100644 --- a/.env.example +++ b/.env.example @@ -16,3 +16,6 @@ OPENAI_API_KEY=sk-your-api-key-here # vLLM / Local server # OPENAI_BASE_URL=http://localhost:8000/v1 # OPENAI_API_KEY=dummy-key + +# Optional: Hugging Face token to push generated datasets +# HF_TOKEN=your-huggingface-token-here diff --git a/CHANGELOG.md b/CHANGELOG.md index af247dd..9171026 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,19 @@ The format is based on Keep a Changelog, and this project adheres to Semantic Ve Nothing yet. +## [0.5.0] - 2025-01-11 +### Added +- Hugging Face Hub integration for direct dataset uploads + - `push_to_hub()` function in new `hf_hub` module to upload datasets to HF Hub + - Uploads JSONL files (train.jsonl, val.jsonl), manifest.json, and auto-generated README.md + - CLI flags: `--push-to-hub`, `--repo-id`, `--hf-token`, `--private` + - Support for both public and private repositories + - Auto-generated dataset cards with dataset statistics, model info, usage examples, and citation +- Optional dependency: `huggingface_hub>=0.20.0` (install with `pip install toolsgen[hf]`) +- Example in `examples/hf_hub_upload/` with dotenv configuration +- Test suite for HF Hub functionality in `tests/test_hf_hub.py` +- `push_to_hub` exported from main `toolsgen` package for easier imports + ## [0.4.0] - 2025-01-10 ### Added - Quality tagging system for generated records diff --git a/README.md b/README.md index 57faec9..1e15487 100644 --- a/README.md +++ b/README.md @@ -76,16 +76,26 @@ toolsgen generate \ --num 500 \ --workers 6 \ --worker-batch-size 4 + +# Generate and push directly to Hugging Face Hub +export HF_TOKEN="your-hf-token-here" +toolsgen generate \ + --tools tools.json \ + --out output_dir \ + --num 100 \ + --push-to-hub \ + --repo-id username/dataset-name ``` ### Python API Usage ```python -import os from pathlib import Path +from dotenv import load_dotenv + from toolsgen.core import GenerationConfig, ModelConfig, generate_dataset -os.environ["OPENAI_API_KEY"] = "your-api-key-here" +load_dotenv() # Load from .env file # Configuration tools_path = Path("tools.json") @@ -119,6 +129,50 @@ print(f"Generated {manifest['num_generated']}/{manifest['num_requested']} record print(f"Failed: {manifest['num_failed']} attempts") ``` +### Push to Hugging Face Hub + +```python +from pathlib import Path +from dotenv import load_dotenv + +from toolsgen import GenerationConfig, ModelConfig, generate_dataset, push_to_hub + +load_dotenv() # Load from .env file + +tools_path = Path("tools.json") +output_dir = Path("output") + +gen_config = GenerationConfig( + num_samples=100, + strategy="random", + seed=42, + train_split=0.9, +) + +model_config = ModelConfig( + model="gpt-4o-mini", + temperature=0.7, +) + +# Generate dataset +manifest = generate_dataset( + output_dir=output_dir, + gen_config=gen_config, + model_config=model_config, + tools_path=tools_path, +) + +# Push to Hub +hub_info = push_to_hub( + output_dir=output_dir, + repo_id="username/dataset-name", + private=False, +) + +print(f"Generated: {manifest['num_generated']} records") +print(f"Repository: {hub_info['repo_url']}") +``` + See `examples/` directory for complete working examples. **Note**: The examples in `examples/` use `python-dotenv` for convenience (load API keys from `.env` file). Install it with `pip install python-dotenv` if you want to use this approach. @@ -228,7 +282,7 @@ For detailed information about the system architecture, pipeline, and core compo - [ ] Custom prompt template system - [x] Parallel generation with multiprocessing - [ ] Additional sampling strategies (coverage-based, difficulty-based) -- [ ] Integration with Hugging Face Hub for direct dataset uploads +- [x] Integration with Hugging Face Hub for direct dataset uploads - [ ] Support for more LLM providers (Anthropic, Cohere, etc.) - [ ] Web UI for dataset inspection and curation - [ ] Advanced filtering and deduplication diff --git a/examples/hf_hub_upload/README.md b/examples/hf_hub_upload/README.md new file mode 100644 index 0000000..b6dd730 --- /dev/null +++ b/examples/hf_hub_upload/README.md @@ -0,0 +1,76 @@ +# Hugging Face Hub Upload Example + +This example demonstrates how to generate a dataset and push it directly to Hugging Face Hub. + +## Prerequisites + +1. OpenAI API key +2. Hugging Face account and token with write access + +## Setup + +```bash +# Install dependencies +pip install toolsgen huggingface_hub python-dotenv + +# Create .env file from example +cp .env.example .env + +# Edit .env and add your API keys +# OPENAI_API_KEY=your-openai-api-key +# HF_TOKEN=your-huggingface-token +``` + +## Usage + +### Python API + +```python +python example.py +``` + +Make sure to update the `repo_id` in the script to your own repository name. + +### CLI + +```bash +toolsgen generate \ + --tools ../basic/tools.json \ + --out output \ + --num 50 \ + --push-to-hub \ + --repo-id your-username/your-dataset-name +``` + +## What Gets Uploaded + +The following files are automatically uploaded to your HF Hub repository: + +- `train.jsonl` - Training dataset +- `val.jsonl` - Validation dataset (if train_split < 1.0) +- `manifest.json` - Generation metadata +- `README.md` - Auto-generated dataset card + +## Repository Visibility + +By default, repositories are public. To create a private repository: + +**Python API:** +```python +hub_info = push_to_hub( + output_dir=output_dir, + repo_id="username/dataset-name", + private=True, +) +``` + +**CLI:** +```bash +toolsgen generate ... --push-to-hub --private +``` + +## Notes + +- The HF token can be provided via `--hf-token` flag or `HF_TOKEN` environment variable +- If a repository already exists, it will be updated with new files +- A dataset card (README.md) is automatically generated if not present diff --git a/examples/hf_hub_upload/example.py b/examples/hf_hub_upload/example.py new file mode 100644 index 0000000..2556351 --- /dev/null +++ b/examples/hf_hub_upload/example.py @@ -0,0 +1,45 @@ +"""Example: Generate dataset and push to Hugging Face Hub.""" + +from pathlib import Path + +from dotenv import load_dotenv + +from toolsgen import GenerationConfig, ModelConfig, generate_dataset, push_to_hub + +# Load environment variables from .env file +load_dotenv() + +# Configuration +tools_path = Path(__file__).parent.parent / "basic" / "tools.json" +output_dir = Path(__file__).parent / "output" + +gen_config = GenerationConfig( + num_samples=50, + strategy="random", + seed=42, + train_split=0.9, +) + +model_config = ModelConfig( + model="gpt-4o-mini", + temperature=0.7, +) + +# Generate dataset +manifest = generate_dataset( + output_dir=output_dir, + gen_config=gen_config, + model_config=model_config, + tools_path=tools_path, +) + +# Push to Hub +hub_info = push_to_hub( + output_dir=output_dir, + repo_id="your-username/your-dataset-name", # Change this! + private=False, +) + +print("\n✓ Dataset generated and uploaded!") +print(f" Generated: {manifest['num_generated']} records") +print(f" Repository: {hub_info['repo_url']}") diff --git a/pyproject.toml b/pyproject.toml index cfbf085..c7f0d84 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ toolsgen = ["prompts/*.txt"] [project] name = "toolsgen" -version = "0.4.0" +version = "0.5.0" description = "Generate tool-calling datasets from OpenAI-compatible tool specs" readme = "README.md" requires-python = ">=3.9" @@ -40,6 +40,11 @@ dependencies = [ "tqdm>=4.66.0", ] +[project.optional-dependencies] +hf = [ + "huggingface_hub>=0.20.0", +] + [project.urls] Homepage = "https://github.com/atasoglu/toolsgen" Repository = "https://github.com/atasoglu/toolsgen" diff --git a/requirements-dev.txt b/requirements-dev.txt index 45f4af5..47401cb 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -3,3 +3,4 @@ pytest-cov>=5.0.0 pre-commit>=3.7.0 ruff>=0.6.0 python-dotenv>=1.0.0 +huggingface_hub>=0.20.0 diff --git a/src/toolsgen/__init__.py b/src/toolsgen/__init__.py index 5e7501a..a7a9a4a 100644 --- a/src/toolsgen/__init__.py +++ b/src/toolsgen/__init__.py @@ -14,6 +14,7 @@ load_tool_specs, write_dataset_jsonl, ) +from .hf_hub import push_to_hub from .judge import JudgeResponse, judge_tool_calls from .problem_generator import generate_problem from .tool_caller import generate_tool_calls @@ -55,6 +56,8 @@ "generate_dataset", "load_tool_specs", "write_dataset_jsonl", + # HF Hub + "push_to_hub", # Judge "JudgeResponse", "judge_tool_calls", diff --git a/src/toolsgen/cli.py b/src/toolsgen/cli.py index a52a6e8..8038d6c 100644 --- a/src/toolsgen/cli.py +++ b/src/toolsgen/cli.py @@ -168,6 +168,28 @@ def create_parser() -> argparse.ArgumentParser: help="Temperature for judging (defaults to --temperature)", ) + # Hugging Face Hub options + gen_parser.add_argument( + "--push-to-hub", + action="store_true", + help="Push dataset to Hugging Face Hub after generation", + ) + gen_parser.add_argument( + "--repo-id", + default=None, + help="HF Hub repository ID (e.g., 'username/dataset-name')", + ) + gen_parser.add_argument( + "--hf-token", + default=None, + help="HF API token (defaults to HF_TOKEN env var)", + ) + gen_parser.add_argument( + "--private", + action="store_true", + help="Create private repository on HF Hub", + ) + return parser @@ -275,9 +297,15 @@ def cmd_generate(args: argparse.Namespace) -> None: max_tokens=args.max_tokens, ) + # Validate HF Hub options + if args.push_to_hub and not args.repo_id: + print("Error: --repo-id is required when using --push-to-hub", file=sys.stderr) + sys.exit(1) + # Generate dataset try: print(f"Generating {args.num} samples using {args.model}...") + manifest = generate_dataset( args.out, gen_config, model_config, tools_path=args.tools ) @@ -297,6 +325,20 @@ def cmd_generate(args: argparse.Namespace) -> None: print(f" - Manifest: {args.out / 'manifest.json'}") + if args.push_to_hub: + from .hf_hub import push_to_hub + + print("\nPushing to Hugging Face Hub...") + hub_info = push_to_hub( + output_dir=args.out, + repo_id=args.repo_id, + token=args.hf_token, + private=args.private, + ) + print("✓ Pushed to Hugging Face Hub") + print(f" - Repository: {hub_info['repo_url']}") + print(f" - Files uploaded: {', '.join(hub_info['files_uploaded'])}") + except ValueError as e: print(f"Error: {e}", file=sys.stderr) sys.exit(1) diff --git a/src/toolsgen/hf_hub.py b/src/toolsgen/hf_hub.py new file mode 100644 index 0000000..5508cef --- /dev/null +++ b/src/toolsgen/hf_hub.py @@ -0,0 +1,157 @@ +"""Hugging Face Hub integration for dataset uploads.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Dict, Optional + +try: + from huggingface_hub import HfApi, create_repo + + HF_HUB_AVAILABLE = True +except ImportError: + HF_HUB_AVAILABLE = False + + +def push_to_hub( + output_dir: Path, + repo_id: str, + token: Optional[str] = None, + private: bool = False, + commit_message: Optional[str] = None, +) -> Dict[str, Any]: + """Push generated dataset to Hugging Face Hub. + + Args: + output_dir: Directory containing train.jsonl, val.jsonl, and manifest.json + repo_id: Repository ID on Hugging Face Hub (e.g., "username/dataset-name") + token: HF API token (if None, uses HF_TOKEN env var or cached token) + private: Whether to create a private repository + commit_message: Custom commit message + + Returns: + Dictionary with upload metadata including repo_url + + Raises: + ImportError: If huggingface_hub is not installed + ValueError: If output_dir doesn't contain required files + """ + if not HF_HUB_AVAILABLE: + raise ImportError( + "huggingface_hub is required for Hub uploads. " + "Install with: pip install huggingface_hub" + ) + + if not output_dir.exists(): + raise ValueError(f"Output directory not found: {output_dir}") + + manifest_path = output_dir / "manifest.json" + if not manifest_path.exists(): + raise ValueError(f"manifest.json not found in {output_dir}") + + api = HfApi(token=token) + + # Create repository + create_repo( + repo_id=repo_id, + repo_type="dataset", + private=private, + exist_ok=True, + token=token, + ) + + # Prepare files to upload + files_to_upload = [manifest_path] + + train_path = output_dir / "train.jsonl" + val_path = output_dir / "val.jsonl" + + if train_path.exists(): + files_to_upload.append(train_path) + if val_path.exists(): + files_to_upload.append(val_path) + + # Generate README if not exists + readme_path = output_dir / "README.md" + if not readme_path.exists(): + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + readme_content = _generate_readme(repo_id, manifest) + readme_path.write_text(readme_content, encoding="utf-8") + files_to_upload.append(readme_path) + + # Upload files + message = commit_message or "Upload ToolsGen dataset" + + for file_path in files_to_upload: + api.upload_file( + path_or_fileobj=str(file_path), + path_in_repo=file_path.name, + repo_id=repo_id, + repo_type="dataset", + commit_message=message, + ) + + repo_url = f"https://huggingface.co/datasets/{repo_id}" + + return { + "repo_id": repo_id, + "repo_url": repo_url, + "files_uploaded": [f.name for f in files_to_upload], + "private": private, + } + + +def _generate_readme(repo_id: str, manifest: Dict[str, Any]) -> str: + """Generate a README.md for the dataset.""" + splits_info = manifest.get("splits", {}) + train_count = splits_info.get("train", 0) + val_count = splits_info.get("val", 0) + + splits_section = "" + if val_count > 0: + splits_section = f""" +## Dataset Splits + +| Split | Records | +|-------|---------| +| train | {train_count} | +| val | {val_count} | +""" + else: + splits_section = f"\n## Dataset Size\n\n{train_count} training records\n" + + readme = f"""# {repo_id.split("/")[-1]} + +This dataset was generated using [ToolsGen](https://github.com/atasoglu/toolsgen). + +## Dataset Description + +- **Generated samples**: {manifest.get("num_generated", 0)} +- **Tools count**: {manifest.get("tools_count", 0)} +- **Sampling strategy**: {manifest.get("strategy", "N/A")} +- **Models used**: + - Problem generator: {manifest.get("models", {}).get("problem_generator", "N/A")} + - Tool caller: {manifest.get("models", {}).get("tool_caller", "N/A")} + - Judge: {manifest.get("models", {}).get("judge", "N/A")} +{splits_section} +## Usage + +```python +from datasets import load_dataset + +dataset = load_dataset("{repo_id}") +``` + +## Citation + +```bibtex +@software{{toolsgen2025, + title = {{ToolsGen: Synthetic Tool-Calling Dataset Generator}}, + author = {{Ataşoğlu, Ahmet}}, + year = {{2025}}, + url = {{https://github.com/atasoglu/toolsgen}} +}} +``` +""" + return readme diff --git a/tests/test_hf_hub.py b/tests/test_hf_hub.py new file mode 100644 index 0000000..0baff7a --- /dev/null +++ b/tests/test_hf_hub.py @@ -0,0 +1,171 @@ +"""Tests for Hugging Face Hub integration.""" + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + + +@pytest.fixture +def mock_output_dir(tmp_path): + """Create a mock output directory with dataset files.""" + output_dir = tmp_path / "output" + output_dir.mkdir() + + # Create manifest.json + manifest = { + "version": "0.1.0", + "num_requested": 100, + "num_generated": 95, + "num_failed": 5, + "strategy": "random", + "seed": 42, + "train_split": 0.9, + "tools_count": 10, + "models": { + "problem_generator": "gpt-4o-mini", + "tool_caller": "gpt-4o-mini", + "judge": "gpt-4o-mini", + }, + "splits": {"train": 85, "val": 10}, + } + (output_dir / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8") + + # Create train.jsonl + (output_dir / "train.jsonl").write_text('{"id": "1"}\n', encoding="utf-8") + + # Create val.jsonl + (output_dir / "val.jsonl").write_text('{"id": "2"}\n', encoding="utf-8") + + return output_dir + + +def test_push_to_hub_missing_huggingface_hub(): + """Test error when huggingface_hub is not installed.""" + from toolsgen.hf_hub import push_to_hub + + with patch("toolsgen.hf_hub.HF_HUB_AVAILABLE", False): + with pytest.raises(ImportError, match="huggingface_hub is required"): + push_to_hub(Path("output"), "user/repo") + + +def test_push_to_hub_missing_output_dir(): + """Test error when output directory doesn't exist.""" + from toolsgen.hf_hub import push_to_hub + + with patch("toolsgen.hf_hub.HF_HUB_AVAILABLE", True): + with pytest.raises(ValueError, match="Output directory not found"): + push_to_hub(Path("nonexistent"), "user/repo") + + +def test_push_to_hub_missing_manifest(tmp_path): + """Test error when manifest.json is missing.""" + from toolsgen.hf_hub import push_to_hub + + output_dir = tmp_path / "output" + output_dir.mkdir() + + with patch("toolsgen.hf_hub.HF_HUB_AVAILABLE", True): + with pytest.raises(ValueError, match="manifest.json not found"): + push_to_hub(output_dir, "user/repo") + + +def test_push_to_hub_success(mock_output_dir): + """Test successful push to Hub.""" + import sys + + # Mock huggingface_hub module + mock_hf_hub = MagicMock() + mock_api = MagicMock() + mock_hf_hub.HfApi.return_value = mock_api + sys.modules["huggingface_hub"] = mock_hf_hub + + try: + # Force reimport with mocked module + import importlib + import toolsgen.hf_hub + + importlib.reload(toolsgen.hf_hub) + from toolsgen.hf_hub import push_to_hub + + result = push_to_hub( + output_dir=mock_output_dir, + repo_id="user/test-dataset", + token="test-token", + private=True, + ) + + # Verify create_repo was called + mock_hf_hub.create_repo.assert_called_once_with( + repo_id="user/test-dataset", + repo_type="dataset", + private=True, + exist_ok=True, + token="test-token", + ) + + # Verify upload_file was called for each file + assert mock_api.upload_file.call_count == 4 # manifest, train, val, README + + # Verify result + assert result["repo_id"] == "user/test-dataset" + assert result["repo_url"] == "https://huggingface.co/datasets/user/test-dataset" + assert result["private"] is True + assert len(result["files_uploaded"]) == 4 + finally: + # Cleanup + if "huggingface_hub" in sys.modules: + del sys.modules["huggingface_hub"] + importlib.reload(toolsgen.hf_hub) + + +def test_generate_readme(): + """Test README generation.""" + from toolsgen.hf_hub import _generate_readme + + manifest = { + "num_generated": 100, + "tools_count": 15, + "strategy": "random", + "models": { + "problem_generator": "gpt-4o-mini", + "tool_caller": "gpt-4o", + "judge": "gpt-4o", + }, + "splits": {"train": 90, "val": 10}, + } + + readme = _generate_readme("user/test-dataset", manifest) + + assert "test-dataset" in readme + assert "100" in readme + assert "15" in readme + assert "random" in readme + assert "gpt-4o-mini" in readme + assert "train | 90" in readme + assert "val | 10" in readme + assert "load_dataset" in readme + assert "user/test-dataset" in readme + + +def test_generate_readme_no_val_split(): + """Test README generation without validation split.""" + from toolsgen.hf_hub import _generate_readme + + manifest = { + "num_generated": 100, + "tools_count": 15, + "strategy": "random", + "models": { + "problem_generator": "gpt-4o-mini", + "tool_caller": "gpt-4o", + "judge": "gpt-4o", + }, + "splits": {"train": 100}, + } + + readme = _generate_readme("user/test-dataset", manifest) + + assert "100 training records" in readme + assert "| Split |" not in readme