psi-bench 0.1.0__tar.gz

psi_bench-0.1.0/PKG-INFO

@@ -0,0 +1,150 @@
+Metadata-Version: 2.4
+Name: psi-bench
+Version: 0.1.0
+Summary: Ψ-Bench: Evaluating Persona-Sensitive Influencing in Persuasive Dialogues
+Author-email: Peixuan Han <ph16@illinois.edu>
+License: MIT
+Project-URL: Repository, https://github.com/yourusername/psi_bench
+Project-URL: Documentation, https://github.com/yourusername/psi_bench#readme
+Project-URL: Bug Tracker, https://github.com/yourusername/psi_bench/issues
+Keywords: benchmark,persuasion,personalization,LLM,evaluation
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: openai>=1.0.0
+Requires-Dist: transformers>=4.30.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: pandas>=1.5.0
+Requires-Dist: sentence-transformers>=2.2.0
+Requires-Dist: scikit-learn>=1.2.0
+Requires-Dist: scipy>=1.10.0
+Requires-Dist: tqdm>=4.65.0
+Requires-Dist: nltk>=3.8.0
+Requires-Dist: json-repair>=0.17.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: black>=22.0; extra == "dev"
+Requires-Dist: flake8>=4.0; extra == "dev"
+Requires-Dist: isort>=5.0; extra == "dev"
+
+<div align="center">
+  <img src="figs/title_fig.jpg" width="100%">
+</div>
+
+<!-- <h1 align="center">Ψ-Bench: Evaluating Persona-Sensitive Influencing in Persuasive Dialogues</h1> -->
+<h2 align="center">Peixuan Han, Hongyi Du, Jiayu Liu, Yihang Sun, Yutong Liu, Jiaxuan You</h2>
+
+<p align="center">
+  <a href="https://pypi.org/project/psi-bench">
+    <img src="https://img.shields.io/pypi/v/your-package?logo=pypi&logoColor=white" />
+  </a>
+  ·
+  <a href="https://arxiv.org/pdf/2606.02754">
+    <img src="https://img.shields.io/badge/arXiv-2401.12345-b31b1b?logo=arxiv" />
+  </a>
+</p>
+
+## Introduction
+**Ψ-Bench (Psi-Bench)** is a benchmark for assessing LLMs' ability to influence realistic users through conversation. We design three real-world interaction scenarios involving persuasion in Psi-Bench and endow simulated clients with personal characteristics via explicit user profiles derived from dialogue histories.
+
+![](figs/main_fig.png)
+
+## Benchmarking Results
+
+Try Psi-Bench to see if your LLM is a personalized expert!
+
+![](figs/main_table.png)
+
+
+## Get Started
+
+Before using Psi-bench, you need to configure your API through environment variables. Psi-bench uses DeepSeek-v3.2 to serve as the client and judge. Since the DeepSeek official API no longer supports this model, it's recommended to access it via  [vocanic platform](https://console.volcengine.com/ark) (the model identifier is deepseek-v3-2-251201). Other LLMs like GPT-4o and DeepSeek-v4 can also serve as clients and judges; however, be cautious when comparing scores from different judges.
+
+For example:
+```bash
+export CLIENT_BASE_URL=https://ark.cn-beijing.volces.com/api/v3
+export CLIENT_API_KEY=sk-...
+export JUDGE_BASE_URL=https://ark.cn-beijing.volces.com/api/v3
+export JUDGE_API_KEY=sk-...
+export PERSUADER_BASE_URL=... # If you want to evaluate API-based persuader models
+export PERSUADER_API_KEY=sk-...
+```
+
+### Download the Package
+This is the easiest way of using Psi-bench.
+```bash
+pip install psi-bench
+psi-bench download all # data will be saved in ./data
+
+# Run evaluation with local persuader model, results will be saved in ./eval
+CUDA_VISIBLE_DEVICES=0 psi-bench eval all \
+  --tested_model Qwen/Qwen3-8B \
+  --persuader_local \
+  --client_model deepseek-v3-2-251201 \
+  --judge_model deepseek-v3-2-251201
+```
+
+### Clone the Repository
+If you wish to develop using Psi-bench or evaluate in more advanced settings (Oracle, profile analyzer, ...), you can clone the Git repo. Below are some examples:
+
+```bash
+git clone https://github.com/Hanpx20/Psi-Bench
+cd Psi-Bench
+
+# Basic evaluation with local persuader
+CUDA_VISIBLE_DEVICES=0 bash eval.sh all \
+  --tested_model Qwen/Qwen3-8B \
+  --client_model deepseek-v3-2-251201 \
+  --judge_model deepseek-v3-2-251201 \
+  --persuader_local
+
+# Inference with oracle setting (client profile provided)
+CUDA_VISIBLE_DEVICES=0 bash eval.sh all \
+  --tested_model Qwen/Qwen3-8B \
+  --client_model deepseek-v3-2-251201 \
+  --judge_model deepseek-v3-2-251201 \
+  --persuader_local \
+  --test_oracle
+
+# Inference with profile analyzer (client profile predicted by an LLM)
+CUDA_VISIBLE_DEVICES=0 python psi_bench/inference.py \
+  --client_model deepseek-v3-2-251201 \
+  --task request \
+  --conv_file data/request/queries.json \
+  --persona_file data/request/persona_profile.json \
+  --persuader_model Qwen/Qwen3-8B \
+  --persuader_local \
+  --profile_mode infer \
+  --persona_infer_model deepseek-v3.2 \
+  --output eval/test.json
+```
+
+### Notations in the Repo
+- CMV, counsel, and request correspond to "Viewpoint Debate", "Psychological Consultation," and "Everyday Request" scenarios, respectively.
+
+- `size` is set to 500 by default, as the first 500 queries in CMV are the test set; the other two scenarios only have 90 and 100 queries in total.
+
+- The LLM judge returns 4 metrics, whereas "general_conversation_quality", "personalized_response" and "persuasion_effect" are what's shown in the paper; "personality_perception" is mainly for investigation purposes.
+
+## Cite this paper
+If you find this repo or the paper useful, please cite:
+```
+@article{han2026psibench,
+  title={$\Psi$-Bench: Evaluating Persona-Sensitive Influencing in Persuasive Dialogues}, 
+  author={Peixuan Han and Hongyi Du and Jiayu Liu and Yihang Sun and Yutong Liu and Jiaxuan You},
+  journal={arXiv preprint arXiv:2606.02754},
+  year={2026},
+  url={https://arxiv.org/abs/2606.02754}, 
+}
+```
+
+Reach out to [Peixuan Han](mailto:ph16@illinois.edu) for any questions.

psi_bench-0.1.0/README.md

@@ -0,0 +1,112 @@
+<div align="center">
+  <img src="figs/title_fig.jpg" width="100%">
+</div>
+
+<!-- <h1 align="center">Ψ-Bench: Evaluating Persona-Sensitive Influencing in Persuasive Dialogues</h1> -->
+<h2 align="center">Peixuan Han, Hongyi Du, Jiayu Liu, Yihang Sun, Yutong Liu, Jiaxuan You</h2>
+
+<p align="center">
+  <a href="https://pypi.org/project/psi-bench">
+    <img src="https://img.shields.io/pypi/v/your-package?logo=pypi&logoColor=white" />
+  </a>
+  ·
+  <a href="https://arxiv.org/pdf/2606.02754">
+    <img src="https://img.shields.io/badge/arXiv-2401.12345-b31b1b?logo=arxiv" />
+  </a>
+</p>
+
+## Introduction
+**Ψ-Bench (Psi-Bench)** is a benchmark for assessing LLMs' ability to influence realistic users through conversation. We design three real-world interaction scenarios involving persuasion in Psi-Bench and endow simulated clients with personal characteristics via explicit user profiles derived from dialogue histories.
+
+![](figs/main_fig.png)
+
+## Benchmarking Results
+
+Try Psi-Bench to see if your LLM is a personalized expert!
+
+![](figs/main_table.png)
+
+
+## Get Started
+
+Before using Psi-bench, you need to configure your API through environment variables. Psi-bench uses DeepSeek-v3.2 to serve as the client and judge. Since the DeepSeek official API no longer supports this model, it's recommended to access it via  [vocanic platform](https://console.volcengine.com/ark) (the model identifier is deepseek-v3-2-251201). Other LLMs like GPT-4o and DeepSeek-v4 can also serve as clients and judges; however, be cautious when comparing scores from different judges.
+
+For example:
+```bash
+export CLIENT_BASE_URL=https://ark.cn-beijing.volces.com/api/v3
+export CLIENT_API_KEY=sk-...
+export JUDGE_BASE_URL=https://ark.cn-beijing.volces.com/api/v3
+export JUDGE_API_KEY=sk-...
+export PERSUADER_BASE_URL=... # If you want to evaluate API-based persuader models
+export PERSUADER_API_KEY=sk-...
+```
+
+### Download the Package
+This is the easiest way of using Psi-bench.
+```bash
+pip install psi-bench
+psi-bench download all # data will be saved in ./data
+
+# Run evaluation with local persuader model, results will be saved in ./eval
+CUDA_VISIBLE_DEVICES=0 psi-bench eval all \
+  --tested_model Qwen/Qwen3-8B \
+  --persuader_local \
+  --client_model deepseek-v3-2-251201 \
+  --judge_model deepseek-v3-2-251201
+```
+
+### Clone the Repository
+If you wish to develop using Psi-bench or evaluate in more advanced settings (Oracle, profile analyzer, ...), you can clone the Git repo. Below are some examples:
+
+```bash
+git clone https://github.com/Hanpx20/Psi-Bench
+cd Psi-Bench
+
+# Basic evaluation with local persuader
+CUDA_VISIBLE_DEVICES=0 bash eval.sh all \
+  --tested_model Qwen/Qwen3-8B \
+  --client_model deepseek-v3-2-251201 \
+  --judge_model deepseek-v3-2-251201 \
+  --persuader_local
+
+# Inference with oracle setting (client profile provided)
+CUDA_VISIBLE_DEVICES=0 bash eval.sh all \
+  --tested_model Qwen/Qwen3-8B \
+  --client_model deepseek-v3-2-251201 \
+  --judge_model deepseek-v3-2-251201 \
+  --persuader_local \
+  --test_oracle
+
+# Inference with profile analyzer (client profile predicted by an LLM)
+CUDA_VISIBLE_DEVICES=0 python psi_bench/inference.py \
+  --client_model deepseek-v3-2-251201 \
+  --task request \
+  --conv_file data/request/queries.json \
+  --persona_file data/request/persona_profile.json \
+  --persuader_model Qwen/Qwen3-8B \
+  --persuader_local \
+  --profile_mode infer \
+  --persona_infer_model deepseek-v3.2 \
+  --output eval/test.json
+```
+
+### Notations in the Repo
+- CMV, counsel, and request correspond to "Viewpoint Debate", "Psychological Consultation," and "Everyday Request" scenarios, respectively.
+
+- `size` is set to 500 by default, as the first 500 queries in CMV are the test set; the other two scenarios only have 90 and 100 queries in total.
+
+- The LLM judge returns 4 metrics, whereas "general_conversation_quality", "personalized_response" and "persuasion_effect" are what's shown in the paper; "personality_perception" is mainly for investigation purposes.
+
+## Cite this paper
+If you find this repo or the paper useful, please cite:
+```
+@article{han2026psibench,
+  title={$\Psi$-Bench: Evaluating Persona-Sensitive Influencing in Persuasive Dialogues}, 
+  author={Peixuan Han and Hongyi Du and Jiayu Liu and Yihang Sun and Yutong Liu and Jiaxuan You},
+  journal={arXiv preprint arXiv:2606.02754},
+  year={2026},
+  url={https://arxiv.org/abs/2606.02754}, 
+}
+```
+
+Reach out to [Peixuan Han](mailto:ph16@illinois.edu) for any questions.

psi_bench-0.1.0/psi_bench/cli.py

@@ -0,0 +1,242 @@
+#!/usr/bin/env python
+"""
+psi_bench CLI - Command line interface for running evaluations
+"""
+
+import argparse
+import sys
+import subprocess
+import os
+from pathlib import Path
+
+def print_help():
+    """Print help message"""
+    help_text = """
+psi_bench Evaluation CLI
+
+USAGE:
+    psi_bench [COMMAND] [OPTIONS]
+
+COMMANDS:
+    eval          - Run evaluations on benchmark datasets
+    download      - Download benchmark datasets
+    help          - Show this help message
+
+EVAL USAGE:
+    psi_bench eval [TASK] [OPTIONS]
+
+    TASKS:
+        cmv       - Evaluate on CMV (Change My View) dataset
+        counsel   - Evaluate on Counsel dataset
+        request   - Evaluate on Request dataset
+        all       - Run all three evaluations
+
+    OPTIONS:
+        --tested_model MODEL          (required) Model to test
+        --client_model MODEL          Client model (default: deepseek-v3-2-251201)
+        --judge_model MODEL           Judge model (default: deepseek-v3-2-251201)
+        --n_turns N                   Number of turns (default: 3)
+        --size M                      Number of conversations to evaluate (default: 500 for cmv, full for other scenarios)
+        --test_oracle                 Test oracle mode
+        --persuader_local             Use local persuader model
+        --inference_parallel N        Parallel workers (default: 16)
+
+DOWNLOAD USAGE:
+    psi_bench download [DATASET]
+
+    DATASETS:
+        cmv       - Change My View dataset
+        counsel   - Counsel dataset
+        request   - Request dataset
+        all       - All datasets (default)
+
+    OPTIONS:
+        --output PATH                 Output directory (default: ./data)
+
+EXAMPLES:
+    # Evaluate model on all datasets
+    psi_bench eval all --tested_model gpt-4
+
+    # Evaluate on specific dataset with local model
+    psi_bench eval cmv --tested_model Qwen/Qwen3-8B --persuader_local
+
+    # Download all datasets
+    psi_bench download all
+
+For more details, see:
+- API Key Setup: https://github.com/yourusername/psi_bench/blob/main/API_KEY_CONFIG.md
+- User Guide: https://github.com/yourusername/psi_bench/blob/main/USER_GUIDE.md
+"""
+    print(help_text)
+
+
+def run_eval_inference(task, tested_model, client_model, judge_model, n_turns, size, test_oracle,
+                       persuader_local, inference_parallel):
+    """Run inference for a specific task"""
+    project_root = Path(__file__).resolve().parent.parent
+    DATA_ROOT = Path.cwd()
+
+    # Data and output file configurations
+    data_config = {
+        "cmv": {
+            "persona_file": "data/cmv/persona_profile.json",
+            "conv_file": "data/cmv/queries.json",
+        },
+        "counsel": {
+            "persona_file": "data/counsel/persona_profile.json",
+            "conv_file": "data/counsel/queries.json",
+        },
+        "request": {
+            "persona_file": "data/request/persona_profile.json",
+            "conv_file": "data/request/queries.json",
+        },
+    }
+
+    if task not in data_config:
+        print(f"Error: Unknown task '{task}'")
+        return 1
+
+    config = data_config[task]
+    model_name = tested_model.split("/")[-1]
+
+    # Build inference.py command
+    inference_cmd = [
+        sys.executable,
+        os.path.join(project_root, "psi_bench", "inference.py"),
+        "--client_model", client_model,
+        "--persona_file", os.path.join(DATA_ROOT, config["persona_file"]),
+        "--conv_file", os.path.join(DATA_ROOT, config["conv_file"]),
+        "--persuader_model", tested_model,
+        "--output", os.path.join(DATA_ROOT, f"eval/{task}/convs/{model_name}{'_oracle' if test_oracle else ''}.json"),
+        "--size", str(size),
+        "--n_turns", str(n_turns),
+        "--inference_parallel", str(inference_parallel),
+    ]
+
+    if test_oracle:
+        inference_cmd.extend(["--profile_mode", "oracle"])
+
+    if persuader_local:
+        inference_cmd.append("--persuader_local")
+
+    print(f"\n{'='*60}")
+    print(f"Running {task.upper()} Inference")
+    print(f"{'='*60}")
+    print(f"Command: {' '.join(inference_cmd)}\n")
+
+    result = subprocess.run(inference_cmd, cwd=project_root)
+    if result.returncode != 0:
+        print(f"Error: Inference failed for {task}")
+        return result.returncode
+
+    # Build llm_judge_eval.py command
+    judge_cmd = [
+        sys.executable,
+        os.path.join(project_root, "psi_bench", "llm_judge_eval.py"),
+        "--judge_model", judge_model,
+        "--conv_file", os.path.join(DATA_ROOT, f"eval/{task}/convs/{model_name}{'_oracle' if test_oracle else ''}.json"),
+        "--persona_file", os.path.join(DATA_ROOT, config["persona_file"]),
+        "--output", os.path.join(DATA_ROOT, f"eval/{task}/{model_name}{'_oracle' if test_oracle else ''}_judge.json"),
+        "--inference_parallel", str(inference_parallel),
+    ]
+
+    print(f"\n{'='*60}")
+    print(f"Running {task.upper()} Judge Evaluation")
+    print(f"{'='*60}")
+    print(f"Command: {' '.join(judge_cmd)}\n")
+
+    result = subprocess.run(judge_cmd, cwd=project_root)
+    if result.returncode != 0:
+        print(f"Error: Judge evaluation failed for {task}")
+        return result.returncode
+
+    return 0
+
+
+def main():
+    """Main CLI entry point"""
+    parser = argparse.ArgumentParser(
+        description="psi_bench: Evaluating Persona-Sensitive Influencing in Persuasive Dialogues",
+        add_help=False
+    )
+
+    parser.add_argument('command', nargs='?', help='Command to run (eval, download, help)')
+    parser.add_argument('task', nargs='?', help='Task for eval command (cmv, counsel, request, all)')
+
+    # Eval options
+    parser.add_argument('--tested_model', type=str, help='Model to test')
+    parser.add_argument('--client_model', type=str, default='deepseek-v3-2-251201', help='Client model')
+    parser.add_argument('--judge_model', type=str, default='deepseek-v3-2-251201', help='Judge model')
+    parser.add_argument('--n_turns', type=int, default=3, help='Number of turns')
+    parser.add_argument('--size', type=int, default=500, help='Number of conversations to evaluate')
+    parser.add_argument('--test_oracle', action='store_true', help='Test oracle mode')
+    parser.add_argument('--persuader_local', action='store_true', help='Use local persuader model')
+    parser.add_argument('--inference_parallel', type=int, default=16, help='Number of parallel workers')
+
+    # Download options
+    parser.add_argument('--output', type=str, default='./data', help='Output directory')
+
+    parser.add_argument('--help', '-h', action='store_true', help='Show help message')
+
+    args, remaining = parser.parse_known_args()
+
+    # Handle help
+    if args.help or not args.command or args.command in ['help', '-h', '--help']:
+        print_help()
+        return 0
+
+    if args.command == 'eval':
+        # Validate eval command
+        if not args.task:
+            print("Error: eval command requires a task (cmv, counsel, request, or all)")
+            print("\nUsage: psi_bench eval [cmv|counsel|request|all] [OPTIONS]")
+            return 1
+
+        if not args.tested_model:
+            print("Error: --tested_model is required")
+            return 1
+
+        if args.task not in ['cmv', 'counsel', 'request', 'all']:
+            print(f"Error: Unknown task '{args.task}'")
+            return 1
+
+        # Run evaluation(s)
+        tasks = ['cmv', 'counsel', 'request'] if args.task == 'all' else [args.task]
+
+        # Ensure output directories exist
+        DATA_ROOT = Path.cwd()
+
+        for task in tasks:
+            (DATA_ROOT / f"eval/{task}/convs").mkdir(parents=True, exist_ok=True)
+            
+            result = run_eval_inference(
+                task=task,
+                tested_model=args.tested_model,
+                client_model=args.client_model,
+                judge_model=args.judge_model,
+                n_turns=args.n_turns,
+                size=args.size,
+                test_oracle=args.test_oracle,
+                persuader_local=args.persuader_local,
+                inference_parallel=args.inference_parallel
+            )
+            if result != 0:
+                return result
+
+        print(f"\n{'='*60}")
+        print("✓ All evaluations completed successfully!")
+        print(f"{'='*60}\n")
+        return 0
+
+    elif args.command == 'download':
+        from . import download_data
+        sys.argv = ['psi_bench-download'] + ([args.task] if args.task else []) + remaining
+        return download_data.main()
+
+    else:
+        print(f"Error: Unknown command '{args.command}'")
+        print_help()
+        return 1
+
+if __name__ == '__main__':
+    sys.exit(main())

psi_bench-0.1.0/psi_bench/download_data.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python
+"""
+Data download script for psi_bench
+Downloads datasets from GitHub or Hugging Face
+"""
+
+import os
+import sys
+import json
+import argparse
+from pathlib import Path
+import urllib.request
+import tarfile
+import zipfile
+
+def download_file(url, output_path, chunk_size=8192):
+    """Download file with progress bar"""
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        print(f"Downloading from: {url}")
+        with urllib.request.urlopen(url) as response:
+            total_size = int(response.headers.get('content-length', 0))
+            downloaded = 0
+
+            with open(output_path, 'wb') as f:
+                while True:
+                    chunk = response.read(chunk_size)
+                    if not chunk:
+                        break
+                    f.write(chunk)
+                    downloaded += len(chunk)
+
+                    if total_size > 0:
+                        percent = (downloaded / total_size) * 100
+                        print(f"  Progress: {percent:.1f}% ({downloaded}/{total_size} bytes)", end='\r')
+
+        print(f"\n✅ Downloaded to: {output_path}")
+        return True
+    except Exception as e:
+        print(f"❌ Download failed: {e}")
+        return False
+
+def extract_archive(archive_path, extract_to):
+    """Extract tar.gz or zip file"""
+    extract_to = Path(extract_to)
+    extract_to.mkdir(parents=True, exist_ok=True)
+
+    try:
+        if str(archive_path).endswith('.tar.gz'):
+            print(f"Extracting tar.gz file...")
+            with tarfile.open(archive_path, 'r:gz') as tar:
+                tar.extractall(extract_to)
+        elif str(archive_path).endswith('.zip'):
+            print(f"Extracting zip file...")
+            with zipfile.ZipFile(archive_path, 'r') as zf:
+                zf.extractall(extract_to)
+        else:
+            print("❌ Unsupported archive format")
+            return False
+
+        print(f"✅ Extracted to: {extract_to}")
+        return True
+    except Exception as e:
+        print(f"❌ Extraction failed: {e}")
+        return False
+
+def download_from_github(dataset_name, output_dir="./data"):
+    """
+    Download dataset from GitHub releases
+    Requires: psi_bench releases with data attached
+    """
+    base_url = "https://github.com/Hanpx20/psi_bench/releases/download"
+
+    datasets = {
+        "cmv": f"{base_url}/v0.1.0/psi_bench-data-cmv.tar.gz",
+        "counsel": f"{base_url}/v0.1.0/psi_bench-data-counsel.tar.gz",
+        "request": f"{base_url}/v0.1.0/psi_bench-data-request.tar.gz",
+        "all": f"{base_url}/v0.1.0/psi_bench-data-all.tar.gz",
+    }
+
+    if dataset_name not in datasets:
+        print(f"Available datasets: {', '.join(datasets.keys())}")
+        return False
+
+    url = datasets[dataset_name]
+    archive_path = Path(output_dir) / f"psi_bench-data-{dataset_name}.tar.gz"
+
+    # Download
+    if not download_file(url, archive_path):
+        return False
+
+    # Extract
+    if not extract_archive(archive_path, output_dir):
+        return False
+
+    # Cleanup
+    archive_path.unlink()
+    print(f"✅ {dataset_name} dataset ready in {output_dir}/")
+    return True
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Download psi_bench datasets from GitHub"
+    )
+    parser.add_argument(
+        "dataset",
+        nargs="?",
+        choices=["cmv", "counsel", "request", "all"],
+        help="Dataset to download (or 'all' for complete dataset)"
+    )
+    parser.add_argument(
+        "--output", "-o",
+        default="./data",
+        help="Output directory (default: ./data)"
+    )
+
+    args = parser.parse_args()
+
+    if not args.dataset:
+        parser.print_help()
+        print("\n💡 Hint: Run 'psi_bench download all' to get all datasets")
+        return 1
+
+    print(f"📥 psi_bench Data Download (from GitHub Hanpx20/psi_bench)\n")
+
+    success = download_from_github(args.dataset, args.output)
+
+    return 0 if success else 1
+
+if __name__ == "__main__":
+    sys.exit(main())