feat: make including system information optional

src/heretic/main.py

@@ -53,10 +53,10 @@ from .utils import (
     format_duration,
     get_readme_intro,
     get_trial_parameters,
+    is_hf_path,
     load_prompts,
     print,
     print_memory_usage,
-    prompt_confirm,
     prompt_password,
     prompt_path,
     prompt_select,
@@ -816,21 +816,41 @@ def run():
                                 settings.good_evaluation_prompts.dataset,
                                 settings.bad_evaluation_prompts.dataset,
                             ]
-                            can_reproduce = not Path(settings.model).exists() and all(
+                            is_reproducible = is_hf_path(settings.model) and all(
-                                not Path(d).exists() for d in datasets
+                                is_hf_path(dataset) for dataset in datasets
                             )
 
-                            if can_reproduce:
+                            if is_reproducible:
-                                # Pin the number of trials to the number of actual completed trials
+                                print(
-                                # for the reproduction configuration.
+                                    (
-                                settings.n_trials = count_completed_trials()
+                                        "Heretic can add information to the repository that allows others to reproduce the model. "
-
+                                        "This is optional, but valuable to the community as both a learning tool and to preserve computational work already done. "
-                                include_reproduce = prompt_confirm(
+                                        "Guaranteeing reproducibility requires basic system information (Python and OS version, CPU and GPU/accelerator info) "
-                                    """Include 'reproduce' folder?
+                                        "as tensor operations can give different results in different system environments. "
-This saves your exact configuration and system information, along with the study checkpoint, to help others verify your results."""
+                                        "[bold]The information does not include any file system paths or other private data.[/]"
+                                    )
                                 )
+                                reproducibility_information = prompt_select(
+                                    "Which reproducibility information do you want to add?",
+                                    [
+                                        Choice(
+                                            title="Full: Settings, package versions, and system information",
+                                            value="full",
+                                        ),
+                                        Choice(
+                                            title="Basic: Settings and package versions",
+                                            value="basic",
+                                        ),
+                                        Choice(
+                                            title="Don't add any reproducibility information",
+                                            value="none",
+                                        ),
+                                    ],
+                                )
+                                if reproducibility_information is None:
+                                    continue
                             else:
-                                include_reproduce = False
+                                reproducibility_information = "none"
 
                             if strategy == "adapter":
                                 print("Uploading LoRA adapter...")
@@ -880,24 +900,35 @@ This saves your exact configuration and system information, along with the study
                                 card.data.tags.append("uncensored")
                                 card.data.tags.append("decensored")
                                 card.data.tags.append("abliterated")
+                                if reproducibility_information != "none":
+                                    card.data.tags.append("reproducible")
                                 card.text = (
-                                    get_readme_intro(settings, trial) + card.text
+                                    get_readme_intro(
+                                        settings,
+                                        trial,
+                                        reproducibility_information != "none",
+                                    )
+                                    + card.text
                                 )
                                 card.push_to_hub(repo_id, token=token)
 
-                            if include_reproduce:
+                            if reproducibility_information != "none":
+                                # Set the number of trials to the number of actual completed trials
+                                # for the reproduction configuration.
+                                settings.n_trials = count_completed_trials()
+
                                 upload_reproduce_folder(
                                     repo_id,
                                     settings,
                                     token,
                                     checkpoint_path=study_checkpoint_file,
                                     trial=trial,
+                                    include_system_information=(
+                                        reproducibility_information == "full"
+                                    ),
                                 )
-                                print(
+
-                                    f"Model and reproducibility files uploaded to [bold]{repo_id}[/]."
+                            print(f"Model uploaded to [bold]{repo_id}[/].")
-                                )
-                            else:
-                                print(f"Model uploaded to [bold]{repo_id}[/].")
 
                         case "Chat with the model":
                             print()

src/heretic/utils.py

@@ -155,18 +155,6 @@ def prompt_password(message: str) -> str:
         return questionary.password(message).ask()
 
 
-def prompt_confirm(message: str, default: bool = True) -> bool:
-    if is_notebook():
-        print()
-        choices = "[Y/n]" if default else "[y/N]"
-        result = input(f"{message} {choices} ").strip().lower()
-        if not result:
-            return default
-        return result in ("y", "yes")
-    else:
-        return questionary.confirm(message, default=default).ask()
-
-
 def format_duration(seconds: float) -> str:
     seconds = round(seconds)
     hours, seconds = divmod(seconds, 3600)
@@ -180,6 +168,18 @@ def format_duration(seconds: float) -> str:
         return f"{seconds}s"
 
 
+def is_hf_path(path: str) -> bool:
+    """Checks whether a path likely refers to a Hugging Face repository."""
+
+    return (
+        not path.startswith("/")
+        and not path.endswith("/")
+        and path.count("/") == 1
+        and "\\" not in path
+        and not Path(path).exists()
+    )
+
+
 @dataclass
 class Prompt:
     system: str
@@ -270,7 +270,11 @@ def get_trial_parameters(trial: Trial) -> dict[str, str]:
     return params
 
 
-def get_readme_intro(settings: Settings, trial: Trial) -> str:
+def get_readme_intro(
+    settings: Settings,
+    trial: Trial,
+    contains_reproducibility_information: bool,
+) -> str:
     if Path(settings.model).exists():
         # Hide the path, which may contain private information.
         model_link = "a model"
@@ -281,7 +285,11 @@ def get_readme_intro(settings: Settings, trial: Trial) -> str:
     return f"""# This is a decensored version of {
         model_link
     }, made using [Heretic](https://github.com/p-e-w/heretic) v{version_info.version}
-
+{
+        f"{chr(10)}**This model is reproducible!** See the [`reproduce`](reproduce) directory and its [README](reproduce/README.md) for more information.{chr(10)}"
+        if contains_reproducibility_information
+        else ""
+    }
 ## Abliteration parameters
 
 | Parameter | Value |
@@ -356,29 +364,78 @@ def generate_reproduce_readme(
     settings: Settings,
     checkpoint_filename: str,
     trial: Trial,
-    timestamp: str | None = None,
+    include_system_information: bool,
 ) -> str:
     """Generates the contents of a README.md for the reproduce/ folder."""
 
-    torch_version = torch.__version__
-    install_hint = f"pip install torch=={torch_version}"
-    if "+" in torch_version:
-        suffix = torch_version.split("+")[1]
-        if suffix:
-            install_hint += f" --index-url https://download.pytorch.org/whl/{suffix}"
-
     heterogeneous_warning = ""
-    if torch.cuda.is_available():
+
-        count = torch.cuda.device_count()
+    if include_system_information:
-        if count > 1:
+        if torch.cuda.is_available():
-            device_names = {torch.cuda.get_device_name(i) for i in range(count)}
+            count = torch.cuda.device_count()
-            if len(device_names) > 1:
+            if count > 1:
-                heterogeneous_warning = """
+                device_names = {torch.cuda.get_device_name(i) for i in range(count)}
+                if len(device_names) > 1:
+                    heterogeneous_warning = """
 > [!WARNING]
-> **Heterogeneous GPUs Detected!**
+> **Heterogeneous GPUs!**
-> This system uses multiple non-identical GPUs. When operations are distributed across different GPUs (e.g. via `device_map='auto'`), non-deterministic behavior can occur. **Reproducibility ***cannot*** be guaranteed in this environment.**
+> This model was generated using multiple non-identical GPUs. When operations are distributed across different GPUs (e.g. via `device_map='auto'`),
+> non-deterministic behavior can occur. **Reproducibility ***cannot*** be guaranteed in this environment.**
 """
 
+        cpu = get_cpu_info_dict()
+        python_env = get_python_env_info_dict()
+
+        accelerators = get_accelerator_info_dict()
+        if accelerators["type"] is None:
+            accelerator_report = "**No GPU or other accelerator detected.**"
+        else:
+            devices = accelerators["devices"]
+            total_vram = sum(d.get("vram_gb", 0) for d in devices)
+            vram_suffix = (
+                f" (`{total_vram:.2f} GB` total VRAM)" if total_vram > 0 else ""
+            )
+            accelerator_lines = [
+                f"- **{accelerators['type']}:** Detected `{len(devices)}` device(s){vram_suffix}"
+            ]
+
+            if accelerators.get("api_name") and accelerators.get("api_version"):
+                accelerator_lines.append(
+                    f"  - **{accelerators['api_name']}:** `{accelerators['api_version']}`"
+                )
+
+            if accelerators.get("driver_version"):
+                accelerator_lines.append(
+                    f"  - **Driver Version:** `{accelerators['driver_version']}`"
+                )
+
+            accelerator_lines.append("- **Devices:**")
+            for i, dev in enumerate(devices):
+                vram = f" (`{dev['vram_gb']:.2f} GB`)" if dev.get("vram_gb") else ""
+                accelerator_lines.append(
+                    f"  - **{accelerators['type']} {i}:** `{dev['name']}`{vram}"
+                )
+            accelerator_report = "\n".join(accelerator_lines)
+
+        system_report = f"""## System
+
+- **Python:** `{python_env["version"]}` (`{python_env["implementation"]}`, `{python_env["compiler"]}`) [`{python_env["environment"]}`]
+- **Operating system:** `{platform.platform()}` (`{platform.machine()}`)
+- **CPU:** `{cpu["brand"] or "Unknown CPU"}`
+
+### Accelerators
+
+{accelerator_report}
+
+"""
+        system_instructions = (
+            "1. Ensure your system matches the specifications in the **System** section above. "
+            "Exact reproducibility is only guaranteed if all aspects of your system are identical to the one the model was originally generated on.\n"
+        )
+    else:
+        system_report = ""
+        system_instructions = ""
+
     version_info = get_heretic_version_info()
     origin_warning = ""
     if not version_info.is_standard_pypi:
@@ -386,138 +443,91 @@ def generate_reproduce_readme(
             repo_info = version_info.origin.split("Git (")[1].strip(")")
             origin_warning = f"""
 > [!NOTE]
-> **Git Installation Detected**
+> **Git installation!**
-> This system installed `heretic-llm` from source repository: `{repo_info}`.
+> This system installed Heretic from a Git repository: `{repo_info}`.
-> To reproduce these results, you must install Heretic from this exact repository and commit.
+> To reproduce the model, you must install Heretic from this exact repository and commit.
 """
         elif version_info.origin == "Local":
             origin_warning = """
 > [!WARNING]
-> **Local Code Detected!**
+> **Local code!**
-> This system installed `heretic-llm` from a local directory or wheel. Uncommitted or experimental code may have been executed. **Reproducibility ***cannot*** be guaranteed in this environment.**
+> This system installed Heretic from a local directory or wheel. Uncommitted or experimental code may have been executed.
+> **Reproducibility ***cannot*** be guaranteed in this environment.**
 """
         else:
             origin_warning = """
 > [!WARNING]
-> **Non-Standard Installation Detected!**
+> **Non-standard installation!**
-> This system installed `heretic-llm` from an unknown non-standard source. **Reproducibility ***cannot*** be guaranteed in this environment.**
+> This system installed Heretic from an unknown non-standard source.
+> **Reproducibility ***cannot*** be guaranteed in this environment.**
 """
 
-    model_link = format_hf_link(settings.model, settings.model_commit)
+    pytorch_version = torch.__version__
-    dataset_info = f"""## Dataset Information
+    pytorch_install_command = f"pip install torch=={pytorch_version}"
-
+    if "+" in pytorch_version:
-- **Good Prompts:** {format_hf_link(settings.good_prompts.dataset, settings.good_prompts.commit, is_dataset=True)}
+        suffix = pytorch_version.split("+")[1]
-- **Bad Prompts:** {format_hf_link(settings.bad_prompts.dataset, settings.bad_prompts.commit, is_dataset=True)}
+        if suffix:
-- **Good Evaluation Prompts:** {format_hf_link(settings.good_evaluation_prompts.dataset, settings.good_evaluation_prompts.commit, is_dataset=True)}
+            pytorch_install_command += (
-- **Bad Evaluation Prompts:** {format_hf_link(settings.bad_evaluation_prompts.dataset, settings.bad_evaluation_prompts.commit, is_dataset=True)}"""
+                f" --index-url https://download.pytorch.org/whl/{suffix}"
-
-    timestamp_str = f"- **Run started at (UTC):** `{timestamp}`" if timestamp else ""
-
-    # System and Accelerator info using structured dictionaries.
-    cpu = get_cpu_info_dict()
-    python_env = get_python_env_info_dict()
-    accelerator = get_accelerator_info_dict()
-
-    # Build System Environment section.
-    system_env_lines = [
-        f"- **OS:** `{platform.platform()}` (`{platform.machine()}`)",
-        f"- **CPU:** `{cpu['brand'] or 'Unknown CPU'}`",
-        f"  - **Information:** Family `{cpu['family']}`, Model `{cpu['model']}`, Stepping `{cpu['stepping']}`",
-    ]
-
-    system_env_lines.extend(
-        [
-            f"- **Python:** `{python_env['version']}` (`{python_env['implementation']}`, `{python_env['compiler']}`) [`{python_env['environment']}`]",
-            f"- **Heretic:** `v{version_info.version}`"
-            + (f" (Origin: `{version_info.origin}`)" if version_info.origin else ""),
-            f"- **PyTorch:** `{torch.__version__}`",
-        ]
-    )
-    system_environment_report = "\n".join(system_env_lines)
-
-    # Build Accelerators section.
-    if accelerator["type"] is None:
-        accelerator_report = "> [!WARNING]\n> **No GPU or other accelerator detected.**"
-    else:
-        devices = accelerator["devices"]
-        total_vram = sum(d.get("vram_gb", 0) for d in devices)
-        vram_suffix = f" (`{total_vram:.2f} GB` total VRAM)" if total_vram > 0 else ""
-        accelerator_lines = [
-            f"- **{accelerator['type']}:** Detected `{len(devices)}` device(s){vram_suffix}"
-        ]
-
-        if accelerator.get("api_name") and accelerator.get("api_version"):
-            accelerator_lines.append(
-                f"  - **{accelerator['api_name']}:** `{accelerator['api_version']}`"
             )
 
-        if accelerator.get("driver_version"):
+    return f"""# Reproduction guide
-            accelerator_lines.append(
-                f"  - **Driver Version:** `{accelerator['driver_version']}`"
-            )
-
-        accelerator_lines.append("- **Devices:**")
-        for i, dev in enumerate(devices):
-            vram = f" (`{dev['vram_gb']:.2f} GB`)" if dev.get("vram_gb") else ""
-            accelerator_lines.append(
-                f"  - **{accelerator['type']} {i}:** `{dev['name']}`{vram}"
-            )
-        accelerator_report = "\n".join(accelerator_lines)
-
-    return f"""# Reproduction Guide
 
 This directory contains the necessary information and assets to reproduce the results obtained during this Heretic run.{heterogeneous_warning}{origin_warning}
 
-## Model Information
+## Models
 
-- **Base Model:** {model_link}
+- **Base model:** {format_hf_link(settings.model, settings.model_commit)}
-{timestamp_str}
 
-{dataset_info}
+## Datasets
 
-## Selected Trial
+- **Good prompts:** {format_hf_link(settings.good_prompts.dataset, settings.good_prompts.commit, is_dataset=True)}
+- **Bad prompts:** {format_hf_link(settings.bad_prompts.dataset, settings.bad_prompts.commit, is_dataset=True)}
+- **Good evaluation prompts:** {format_hf_link(settings.good_evaluation_prompts.dataset, settings.good_evaluation_prompts.commit, is_dataset=True)}
+- **Bad evaluation prompts:** {format_hf_link(settings.bad_evaluation_prompts.dataset, settings.bad_evaluation_prompts.commit, is_dataset=True)}
 
-- **Trial Number:** `#{trial.user_attrs["index"]}`
+## Selected trial
-- **Refusal Count:** `{trial.user_attrs["refusals"]}/{trial.user_attrs["n_bad_prompts"]}`
-- **KL Divergence:** `{trial.user_attrs["kl_divergence"]:.6f}`
 
-## System Environment
+- **Trial number:** `{trial.user_attrs["index"]}`
+- **KL divergence:** `{trial.user_attrs["kl_divergence"]:.6f}`
+- **Refusals:** `{trial.user_attrs["refusals"]}/{trial.user_attrs["n_bad_prompts"]}`
 
-{system_environment_report}
+{system_report}## Environment
 
-### Accelerators
+- **Heretic:** `v{version_info.version}`{f" (Origin: `{version_info.origin}`)" if version_info.origin else ""}
+- **PyTorch:** `{pytorch_version}`
+- **Other dependencies:** See [`requirements.txt`](requirements.txt).
 
-{accelerator_report}
+## Contents of this directory
 
-## Contents
+- [`requirements.txt`](requirements.txt): The exact versions of all Python packages.
+- [`config.toml`](config.toml): The exact configuration used, including the RNG seed.
+- [`{checkpoint_filename}`]({checkpoint_filename}): The Optuna study journal containing the history of all trials.
+- [`SHA256SUMS`](SHA256SUMS): Cryptographic hashes for all weight files.
+- [`reproduce.json`](reproduce.json): A machine-readable file containing all reproducibility information.
 
-- **config.toml**: The exact configuration used, including the seed `{settings.seed}`.
+## How to reproduce
-- **requirements.txt**: The exact versions of all installed Python packages.
-- **{checkpoint_filename}**: The Optuna study journal containing the history of all trials.
-- **reproduce.json**: A machine-readable version of this report.
-- **SHA256SUMS**: Cryptographic hashes for all uploaded weight files (if applicable).
 
-## How to Reproduce
+{system_instructions}1. Install the exact version of Heretic indicated in the **Environment** section above, from its original source.
-
+1. Install the packages listed in `requirements.txt`: `pip install -r requirements.txt`
-1. Ensure your hardware and environment match the specifications in the **System Environment** section above.
+1. Install the correct version of PyTorch: `{pytorch_install_command}`
-2. Install the exact package versions listed in `requirements.txt`.
+1. Place the provided `config.toml` in your working directory.
-3. Place the provided `config.toml` in your working directory.
+1. Run Heretic without any additional arguments: `heretic`
-4. Run `heretic` without any additional arguments.
+1. Wait for the run to finish, then select trial **{trial.user_attrs["index"]}** and export the model.
-5. Verify the integrity of the reproduced files by comparing their SHA256 hashes against the manifest in `SHA256SUMS`.
+1. Verify that the weight files have been exactly reproduced by comparing their SHA-256 hashes against those in `SHA256SUMS`: `sha256sum -c SHA256SUMS` (or look at the hashes online if you uploaded to Hugging Face)
 
 > [!TIP]
-> To use the included Optuna study journal `{checkpoint_filename}`, place it in the checkpoints directory (usually `checkpoints/`) before running `heretic` on the same model.
+> To use the included Optuna study journal `{checkpoint_filename}`, place it in the checkpoints directory (usually `checkpoints/`) before running Heretic.
-
+> This allows you to export other models from the Pareto front, or to run additional trials without having to re-run the stored trials.
-> [!IMPORTANT]
-> Make sure to install correct PyTorch version from: `{install_hint}`
 """
 
 
 def generate_reproduce_json(
     settings: Settings,
     trial: Trial,
-    timestamp: str | None = None,
+    timestamp: str,
-    uploaded_model_hashes: dict[str, str] | None = None,
+    uploaded_model_hashes: dict[str, str],
+    include_system_information: bool,
 ) -> str:
     """Generates the contents of a reproduce.json file for the reproduce/ folder."""
 
@@ -526,15 +536,7 @@ def generate_reproduce_json(
     data = {
         "version": "1",  # Version number of the reproduce.json file format, to allow for future changes.
         "timestamp": timestamp,
-        "system": {
+        "system": None,  # Defined here to preserve insertion order.
-            "python": get_python_env_info_dict(),
-            "os": {
-                "platform": platform.platform(),
-                "machine": platform.machine(),
-            },
-            "cpu": get_cpu_info_dict(),
-            "accelerator": get_accelerator_info_dict(),
-        },
         "environment": {
             "heretic": {
                 "version": version_info.version,
@@ -555,9 +557,22 @@ def generate_reproduce_json(
             "base_refusals": trial.user_attrs["base_refusals"],
             "n_bad_prompts": trial.user_attrs["n_bad_prompts"],
         },
-        "hashes": uploaded_model_hashes or {},
+        "hashes": uploaded_model_hashes,
     }
 
+    if include_system_information:
+        data["system"] = {
+            "python": get_python_env_info_dict(),
+            "os": {
+                "platform": platform.platform(),
+                "machine": platform.machine(),
+            },
+            "cpu": get_cpu_info_dict(),
+            "accelerators": get_accelerator_info_dict(),
+        }
+    else:
+        del data["system"]
+
     return json.dumps(data, indent=4)
 
 
@@ -578,7 +593,8 @@ def create_reproduce_folder(
     settings: Settings,
     checkpoint_path: str | Path,
     trial: Trial,
-    uploaded_model_hashes: dict[str, str] | None = None,
+    uploaded_model_hashes: dict[str, str],
+    include_system_information: bool,
 ):
     reproduce_dir = path / "reproduce"
     reproduce_dir.mkdir(parents=True, exist_ok=True)
@@ -602,34 +618,39 @@ def create_reproduce_folder(
         datetime.now(timezone.utc).replace(microsecond=0, tzinfo=None).isoformat()
     )
 
-    (reproduce_dir / "config.toml").write_text(
-        generate_config_toml(settings),
-        encoding="utf-8",
-    )
     (reproduce_dir / "requirements.txt").write_text(
         generate_requirements_txt(),
         encoding="utf-8",
     )
-    (reproduce_dir / "README.md").write_text(
+
-        generate_reproduce_readme(
+    (reproduce_dir / "config.toml").write_text(
-            settings,
+        generate_config_toml(settings),
-            checkpoint_filename,
-            trial,
-            timestamp=timestamp,
-        ),
         encoding="utf-8",
     )
+
     if uploaded_model_hashes:
         (reproduce_dir / "SHA256SUMS").write_text(
             generate_sha256sums(uploaded_model_hashes),
             encoding="utf-8",
         )
+
     (reproduce_dir / "reproduce.json").write_text(
         generate_reproduce_json(
             settings,
             trial,
             timestamp=timestamp,
             uploaded_model_hashes=uploaded_model_hashes,
+            include_system_information=include_system_information,
+        ),
+        encoding="utf-8",
+    )
+
+    (reproduce_dir / "README.md").write_text(
+        generate_reproduce_readme(
+            settings,
+            checkpoint_filename,
+            trial,
+            include_system_information=include_system_information,
         ),
         encoding="utf-8",
     )
@@ -646,6 +667,7 @@ def upload_reproduce_folder(
     token: str,
     checkpoint_path: str | Path,
     trial: Trial,
+    include_system_information: bool,
 ):
     api = huggingface_hub.HfApi()
     info = api.model_info(repo_id=repo_id, files_metadata=True, token=token)
@@ -673,6 +695,7 @@ def upload_reproduce_folder(
             checkpoint_path=checkpoint_path,
             trial=trial,
             uploaded_model_hashes=uploaded_model_hashes,
+            include_system_information=include_system_information,
         )
 
         reproduce_dir = tmp_path / "reproduce"