Publish HALT-CoT

.gitignore

@@ -0,0 +1,17 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+.venv/
+venv/
+env/
+build/
+dist/
+*.egg-info/
+.env
+.env.*
+!.env.example
+.DS_Store
+logs/
+outputs/

CITATION.cff

@@ -0,0 +1,8 @@
+cff-version: 1.2.0
+message: "If you use HALT-CoT, please cite the paper."
+title: "HALT-CoT: Model-Agnostic Early Stopping for Chain-of-Thought Reasoning via Answer Entropy"
+authors:
+  - family-names: "Laaouach"
+    given-names: "Yassir"
+year: 2026
+license: "MIT"

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Yassir Laaouach
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,136 @@
+---
+title: HALT-CoT
+sdk: gradio
+app_file: app.py
+license: mit
+tags:
+  - chain-of-thought
+  - reasoning
+  - early-stopping
+  - entropy
+  - transformers
+  - inference-optimization
+---
+
+# HALT-CoT
+
+HALT-CoT is an inference-time early stopping method for chain-of-thought reasoning. After each generated reasoning step, it computes the Shannon entropy of the model's answer distribution over candidate answers. When entropy stays below a threshold, generation stops and the current most likely answer is returned.
+
+This repository packages the method from:
+
+**HALT-CoT: Model-Agnostic Early Stopping for Chain-of-Thought Reasoning via Answer Entropy**  
+Yassir Laaouach
+
+The release is a method implementation and Hugging Face Space, not trained model weights.
+
+## What Is Included
+
+- Dependency-free core entropy and halting controller in `halt_cot/core.py`.
+- Hugging Face Transformers backend in `halt_cot/transformers_backend.py`.
+- CLI entrypoint: `halt-cot`.
+- Gradio Space app in `app.py`.
+- Publish helper for `huggingface_hub`.
+- Focused tests for entropy, candidate sets, and halting behavior.
+
+## Install
+
+```bash
+pip install -e ".[transformers]"
+```
+
+For the Gradio Space locally:
+
+```bash
+pip install -e ".[demo]"
+python app.py
+```
+
+## CLI Example
+
+```bash
+python -m halt_cot.cli \
+  --model Qwen/Qwen2.5-0.5B-Instruct \
+  --question "If a shop has 12 apples and sells 5, how many apples are left?" \
+  --candidates 5 6 7 8 9 \
+  --theta 0.6 \
+  --consecutive 2
+```
+
+## Python Example
+
+```python
+from halt_cot import HaltCoTConfig
+from halt_cot.transformers_backend import HaltCoTForCausalLM
+
+runner = HaltCoTForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
+config = HaltCoTConfig(theta=0.6, consecutive_low_entropy=2, max_steps=8)
+
+result = runner.run(
+    "If a shop has 12 apples and sells 5, how many apples are left?",
+    candidates=["5", "6", "7", "8", "9"],
+    config=config,
+)
+
+print(result.answer)
+print(result.reasoning)
+```
+
+## Method
+
+At reasoning step `i`, HALT-CoT scores candidate answers `a in A` from the next-token logits after the current partial chain. It normalizes those candidate scores into `p_i(a)` and computes:
+
+```text
+H_i = - sum_a p_i(a) log p_i(a)
+```
+
+The controller halts when `H_i < theta`, optionally requiring the condition for multiple consecutive steps. This implementation uses bits by default, which aligns with the entropy traces in the paper. If you tune thresholds in nats, set `entropy_unit="nats"`.
+
+Suggested starting thresholds from the paper:
+
+- Math-style numeric tasks: `theta` in `[0.5, 0.7]`.
+- StrategyQA-style yes/no tasks: `theta = 0.8`.
+- CommonsenseQA-style multiple choice: `theta = 0.7`.
+
+For numeric tasks, use task-specific candidate answers where possible. The helper `numeric_candidates_from_texts` can build candidates from training answers plus an integer range.
+
+## Publish To Hugging Face
+
+Log in once:
+
+```bash
+python -m huggingface_hub.commands.huggingface_cli login
+```
+
+### Free Code/Method Repository
+
+Hugging Face may require a PRO subscription for hosted Gradio Spaces on `cpu-basic`. To publish the method, code, docs, and tests for free, upload this folder as a normal Hugging Face model repository:
+
+```bash
+python scripts/publish_to_hf.py --repo-id <your-username>/halt-cot --repo-type model
+```
+
+For this account, the repo id is:
+
+```bash
+python scripts/publish_to_hf.py --repo-id yass4/halt-cot --repo-type model
+```
+
+### Interactive Gradio Space
+
+If your account can create Gradio Spaces, upload the same folder as a Space:
+
+```bash
+python scripts/publish_to_hf.py --repo-id <your-username>/halt-cot --repo-type space
+```
+
+The default Space model is controlled by `HALT_COT_MODEL_ID`. Set it in the Space variables to use a different causal LM. Use `HALT_COT_DEVICE_MAP=auto` on hardware that supports Accelerate device placement.
+
+## Citation
+
+```bibtex
+@misc{laaouach2026haltcot,
+  title = {HALT-CoT: Model-Agnostic Early Stopping for Chain-of-Thought Reasoning via Answer Entropy},
+  author = {Laaouach, Yassir},
+  year = {2026}
+}
+```

app.py

@@ -0,0 +1,125 @@
+"""Hugging Face Space app for HALT-CoT."""
+
+from __future__ import annotations
+
+import os
+
+import gradio as gr
+
+from halt_cot import HaltCoTConfig
+from halt_cot.transformers_backend import HaltCoTForCausalLM
+
+
+DEFAULT_MODEL_ID = os.getenv("HALT_COT_MODEL_ID", "Qwen/Qwen2.5-0.5B-Instruct")
+DEFAULT_DEVICE_MAP = os.getenv("HALT_COT_DEVICE_MAP") or None
+DEFAULT_THETA = float(os.getenv("HALT_COT_THETA", "0.6"))
+
+_RUNNER: HaltCoTForCausalLM | None = None
+
+
+def get_runner() -> HaltCoTForCausalLM:
+    global _RUNNER
+    if _RUNNER is None:
+        _RUNNER = HaltCoTForCausalLM.from_pretrained(
+            DEFAULT_MODEL_ID,
+            device_map=DEFAULT_DEVICE_MAP,
+        )
+    return _RUNNER
+
+
+def parse_candidates(candidate_text: str) -> list[str]:
+    lines = []
+    for raw_line in candidate_text.replace(",", "\n").splitlines():
+        line = raw_line.strip()
+        if line:
+            lines.append(line)
+    if len(lines) < 2:
+        raise gr.Error("Provide at least two candidate answers.")
+    return lines
+
+
+def run_halt_cot(
+    question: str,
+    candidate_text: str,
+    theta: float,
+    consecutive: int,
+    max_steps: int,
+    step_max_new_tokens: int,
+):
+    if not question.strip():
+        raise gr.Error("Question is required.")
+
+    config = HaltCoTConfig(
+        theta=float(theta),
+        consecutive_low_entropy=int(consecutive),
+        max_steps=int(max_steps),
+        step_max_new_tokens=int(step_max_new_tokens),
+    )
+    result = get_runner().run(question.strip(), parse_candidates(candidate_text), config=config)
+    trace_rows = [
+        [
+            step.index,
+            round(step.entropy, 4),
+            step.prediction,
+            step.generated_tokens,
+            step.halted,
+            step.text,
+        ]
+        for step in result.steps
+    ]
+    final_probs = []
+    if result.steps:
+        final_probs = sorted(
+            result.steps[-1].probabilities.items(),
+            key=lambda item: item[1],
+            reverse=True,
+        )
+        final_probs = [[label, round(probability, 6)] for label, probability in final_probs]
+
+    status = "halted" if result.halted else "reached max steps"
+    summary = (
+        f"Answer: {result.answer}\n"
+        f"Status: {status}\n"
+        f"Generated reasoning tokens: {result.generated_tokens}"
+    )
+    return summary, result.reasoning, trace_rows, final_probs
+
+
+with gr.Blocks(title="HALT-CoT") as demo:
+    gr.Markdown("# HALT-CoT")
+    with gr.Row():
+        with gr.Column(scale=2):
+            question = gr.Textbox(
+                label="Question",
+                lines=5,
+                value="If a shop has 12 apples and sells 5, how many apples are left?",
+            )
+            candidates = gr.Textbox(
+                label="Candidate answers",
+                lines=6,
+                value="5\n6\n7\n8\n9",
+            )
+            run_button = gr.Button("Run HALT-CoT", variant="primary")
+        with gr.Column(scale=1):
+            theta = gr.Slider(0.0, 2.0, value=DEFAULT_THETA, step=0.05, label="Entropy threshold")
+            consecutive = gr.Slider(1, 4, value=2, step=1, label="Consecutive low-entropy steps")
+            max_steps = gr.Slider(1, 12, value=8, step=1, label="Maximum steps")
+            step_max_new_tokens = gr.Slider(8, 160, value=64, step=8, label="Step token cap")
+
+    summary = gr.Textbox(label="Result", lines=3)
+    reasoning = gr.Textbox(label="Generated reasoning", lines=8)
+    trace = gr.Dataframe(
+        headers=["Step", "Entropy", "Prediction", "Tokens", "Halted", "Text"],
+        label="HALT trace",
+    )
+    probabilities = gr.Dataframe(headers=["Candidate", "Probability"], label="Final answer distribution")
+
+    run_button.click(
+        run_halt_cot,
+        inputs=[question, candidates, theta, consecutive, max_steps, step_max_new_tokens],
+        outputs=[summary, reasoning, trace, probabilities],
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

examples/run_halt_cot.py

@@ -0,0 +1,19 @@
+"""Minimal HALT-CoT example using a Hugging Face causal LM."""
+
+from halt_cot import HaltCoTConfig
+from halt_cot.transformers_backend import HaltCoTForCausalLM
+
+
+def main() -> None:
+    runner = HaltCoTForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
+    result = runner.run(
+        "If a shop has 12 apples and sells 5, how many apples are left?",
+        candidates=["5", "6", "7", "8", "9"],
+        config=HaltCoTConfig(theta=0.6, consecutive_low_entropy=2, max_steps=8),
+    )
+    print(result.answer)
+    print(result.reasoning)
+
+
+if __name__ == "__main__":
+    main()

halt_cot/__init__.py

@@ -0,0 +1,39 @@
+"""HALT-CoT: entropy-based early stopping for chain-of-thought reasoning."""
+
+from .core import (
+    AnswerCandidate,
+    AnswerDistribution,
+    EntropyHaltingController,
+    HaltCoTConfig,
+    HaltCoTResult,
+    HaltCoTStep,
+    HaltDecision,
+    answer_distribution_from_scores,
+    entropy_from_probabilities,
+    format_reasoning,
+    integer_candidates,
+    multiple_choice_candidates,
+    normalize_candidates,
+    numeric_candidates_from_texts,
+    softmax,
+    yes_no_candidates,
+)
+
+__all__ = [
+    "AnswerCandidate",
+    "AnswerDistribution",
+    "EntropyHaltingController",
+    "HaltCoTConfig",
+    "HaltCoTResult",
+    "HaltCoTStep",
+    "HaltDecision",
+    "answer_distribution_from_scores",
+    "entropy_from_probabilities",
+    "format_reasoning",
+    "integer_candidates",
+    "multiple_choice_candidates",
+    "normalize_candidates",
+    "numeric_candidates_from_texts",
+    "softmax",
+    "yes_no_candidates",
+]

halt_cot/cli.py

@@ -0,0 +1,118 @@
+"""Command line entrypoint for HALT-CoT."""
+
+from __future__ import annotations
+
+import argparse
+import json
+
+from .core import (
+    HaltCoTConfig,
+    integer_candidates,
+    multiple_choice_candidates,
+    yes_no_candidates,
+)
+from .transformers_backend import HaltCoTForCausalLM
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run HALT-CoT with a Hugging Face causal LM.")
+    parser.add_argument("--model", required=True, help="Hugging Face model id or local model path.")
+    parser.add_argument("--question", required=True, help="Question to answer.")
+    parser.add_argument(
+        "--candidate-set",
+        choices=("custom", "yes-no", "multiple-choice", "integers"),
+        default="custom",
+        help="Candidate answer set.",
+    )
+    parser.add_argument(
+        "--candidates",
+        nargs="*",
+        default=None,
+        help="Custom candidate labels, e.g. --candidates Yes No or --candidates A B C D E.",
+    )
+    parser.add_argument("--integer-start", type=int, default=0)
+    parser.add_argument("--integer-end", type=int, default=100)
+    parser.add_argument("--theta", type=float, default=0.6)
+    parser.add_argument("--max-steps", type=int, default=12)
+    parser.add_argument("--min-steps", type=int, default=1)
+    parser.add_argument("--consecutive", type=int, default=2)
+    parser.add_argument("--step-max-new-tokens", type=int, default=96)
+    parser.add_argument("--entropy-unit", choices=("bits", "nats"), default="bits")
+    parser.add_argument("--device-map", default=None, help="Optional Transformers device_map, e.g. auto.")
+    parser.add_argument("--trust-remote-code", action="store_true")
+    parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON.")
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    candidates = _resolve_candidates(args)
+    config = HaltCoTConfig(
+        theta=args.theta,
+        max_steps=args.max_steps,
+        min_steps=args.min_steps,
+        consecutive_low_entropy=args.consecutive,
+        step_max_new_tokens=args.step_max_new_tokens,
+        entropy_unit=args.entropy_unit,
+    )
+
+    runner = HaltCoTForCausalLM.from_pretrained(
+        args.model,
+        config=config,
+        device_map=args.device_map,
+        trust_remote_code=args.trust_remote_code,
+    )
+    result = runner.run(args.question, candidates)
+
+    if args.json:
+        print(
+            json.dumps(
+                {
+                    "answer": result.answer,
+                    "halted": result.halted,
+                    "generated_tokens": result.generated_tokens,
+                    "steps": [
+                        {
+                            "index": step.index,
+                            "text": step.text,
+                            "entropy": step.entropy,
+                            "prediction": step.prediction,
+                            "probabilities": step.probabilities,
+                            "halted": step.halted,
+                            "generated_tokens": step.generated_tokens,
+                        }
+                        for step in result.steps
+                    ],
+                },
+                indent=2,
+            )
+        )
+    else:
+        print(f"Answer: {result.answer}")
+        print(f"Halted: {result.halted}")
+        print(f"Generated tokens: {result.generated_tokens}")
+        print("\nTrace:")
+        for step in result.steps:
+            marker = " HALT" if step.halted else ""
+            print(
+                f"- Step {step.index}: H={step.entropy:.3f} "
+                f"pred={step.prediction!r}{marker} | {step.text}"
+            )
+    return 0
+
+
+def _resolve_candidates(args: argparse.Namespace):
+    if args.candidate_set == "yes-no":
+        return yes_no_candidates()
+    if args.candidate_set == "multiple-choice":
+        labels = args.candidates or ["A", "B", "C", "D", "E"]
+        return multiple_choice_candidates(labels)
+    if args.candidate_set == "integers":
+        return integer_candidates(args.integer_start, args.integer_end)
+    if not args.candidates:
+        raise SystemExit("--candidates is required when --candidate-set custom")
+    return args.candidates
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

halt_cot/core.py

@@ -0,0 +1,294 @@
+"""Core HALT-CoT primitives.
+
+The functions in this module are dependency-free so they can be used in API
+integrations, tests, and backends that are not based on Hugging Face
+Transformers.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import math
+import re
+from typing import Iterable, Mapping, Sequence
+
+
+EntropyUnit = str
+
+
+@dataclass(frozen=True)
+class AnswerCandidate:
+    """A possible final answer and optional textual aliases for token scoring."""
+
+    label: str
+    aliases: tuple[str, ...] = ()
+
+    def __post_init__(self) -> None:
+        label = self.label.strip()
+        if not label:
+            raise ValueError("AnswerCandidate.label must not be empty")
+        object.__setattr__(self, "label", label)
+        object.__setattr__(
+            self,
+            "aliases",
+            tuple(alias.strip() for alias in self.aliases if alias.strip()),
+        )
+
+    @property
+    def texts(self) -> tuple[str, ...]:
+        return (self.label, *self.aliases)
+
+
+@dataclass(frozen=True)
+class AnswerDistribution:
+    """Normalized answer probabilities and their entropy."""
+
+    probabilities: dict[str, float]
+    entropy: float
+    prediction: str
+
+
+@dataclass(frozen=True)
+class HaltCoTConfig:
+    """Configuration for entropy-based chain-of-thought halting."""
+
+    theta: float = 0.6
+    max_steps: int = 12
+    min_steps: int = 1
+    consecutive_low_entropy: int = 2
+    entropy_unit: EntropyUnit = "bits"
+    step_max_new_tokens: int = 96
+    step_min_new_tokens: int = 4
+    prompt_template: str = "{question}\n\nLet's think step by step.\n"
+    step_prefix: str = "Step {step}: "
+    step_stop_strings: tuple[str, ...] = ("\nStep", "\nAnswer:", "\nFinal answer:", "\n")
+    answer_probe: str = "\nTherefore, the final answer is "
+    do_sample: bool = False
+    temperature: float = 0.0
+    top_p: float = 1.0
+
+    def __post_init__(self) -> None:
+        if self.theta < 0:
+            raise ValueError("theta must be non-negative")
+        if self.max_steps < 1:
+            raise ValueError("max_steps must be at least 1")
+        if self.min_steps < 1:
+            raise ValueError("min_steps must be at least 1")
+        if self.min_steps > self.max_steps:
+            raise ValueError("min_steps cannot exceed max_steps")
+        if self.consecutive_low_entropy < 1:
+            raise ValueError("consecutive_low_entropy must be at least 1")
+        if self.entropy_unit not in {"bits", "nats"}:
+            raise ValueError("entropy_unit must be 'bits' or 'nats'")
+        if self.step_max_new_tokens < 1:
+            raise ValueError("step_max_new_tokens must be at least 1")
+        if self.step_min_new_tokens < 0:
+            raise ValueError("step_min_new_tokens must be non-negative")
+        if self.temperature < 0:
+            raise ValueError("temperature must be non-negative")
+        if not 0 < self.top_p <= 1:
+            raise ValueError("top_p must be in (0, 1]")
+
+
+@dataclass(frozen=True)
+class HaltDecision:
+    """Controller state after observing one entropy value."""
+
+    should_halt: bool
+    low_entropy_streak: int
+
+
+@dataclass(frozen=True)
+class HaltCoTStep:
+    """One generated reasoning step plus HALT-CoT diagnostics."""
+
+    index: int
+    text: str
+    entropy: float
+    prediction: str
+    probabilities: dict[str, float]
+    halted: bool
+    generated_tokens: int = 0
+
+
+@dataclass(frozen=True)
+class HaltCoTResult:
+    """Final output of a HALT-CoT run."""
+
+    answer: str
+    halted: bool
+    steps: tuple[HaltCoTStep, ...]
+    reasoning: str
+    generated_tokens: int
+
+
+class EntropyHaltingController:
+    """Implements the threshold and consecutive-step guard from the paper."""
+
+    def __init__(self, config: HaltCoTConfig):
+        self.config = config
+        self.low_entropy_streak = 0
+
+    def observe(self, entropy: float, step_index: int) -> HaltDecision:
+        if step_index < self.config.min_steps:
+            self.low_entropy_streak = 0
+            return HaltDecision(False, self.low_entropy_streak)
+
+        if entropy < self.config.theta:
+            self.low_entropy_streak += 1
+        else:
+            self.low_entropy_streak = 0
+
+        return HaltDecision(
+            should_halt=self.low_entropy_streak >= self.config.consecutive_low_entropy,
+            low_entropy_streak=self.low_entropy_streak,
+        )
+
+
+def normalize_candidates(
+    candidates: Sequence[str | AnswerCandidate],
+) -> tuple[AnswerCandidate, ...]:
+    """Convert user-provided strings or candidates into validated candidates."""
+
+    normalized: list[AnswerCandidate] = []
+    seen: set[str] = set()
+    for candidate in candidates:
+        item = candidate if isinstance(candidate, AnswerCandidate) else AnswerCandidate(str(candidate))
+        key = item.label.casefold()
+        if key in seen:
+            raise ValueError(f"Duplicate answer candidate: {item.label!r}")
+        seen.add(key)
+        normalized.append(item)
+
+    if len(normalized) < 2:
+        raise ValueError("At least two answer candidates are required")
+    return tuple(normalized)
+
+
+def softmax(scores: Mapping[str, float]) -> dict[str, float]:
+    """Stable softmax over a label-to-score mapping."""
+
+    if not scores:
+        raise ValueError("scores must not be empty")
+
+    max_score = max(scores.values())
+    weights = {label: math.exp(score - max_score) for label, score in scores.items()}
+    total = sum(weights.values())
+    if total == 0 or not math.isfinite(total):
+        raise ValueError("scores produced an invalid softmax normalization")
+    return {label: weight / total for label, weight in weights.items()}
+
+
+def entropy_from_probabilities(
+    probabilities: Mapping[str, float],
+    entropy_unit: EntropyUnit = "bits",
+) -> float:
+    """Compute Shannon entropy for an answer distribution."""
+
+    if entropy_unit not in {"bits", "nats"}:
+        raise ValueError("entropy_unit must be 'bits' or 'nats'")
+    if not probabilities:
+        raise ValueError("probabilities must not be empty")
+
+    total = sum(probabilities.values())
+    if not math.isclose(total, 1.0, rel_tol=1e-6, abs_tol=1e-6):
+        raise ValueError(f"probabilities must sum to 1, got {total}")
+
+    log = math.log2 if entropy_unit == "bits" else math.log
+    return -sum(p * log(p) for p in probabilities.values() if p > 0)
+
+
+def answer_distribution_from_scores(
+    scores: Mapping[str, float],
+    entropy_unit: EntropyUnit = "bits",
+) -> AnswerDistribution:
+    """Turn answer logits/scores into probabilities, entropy, and argmax answer."""
+
+    probabilities = softmax(scores)
+    entropy = entropy_from_probabilities(probabilities, entropy_unit)
+    prediction = max(probabilities, key=probabilities.__getitem__)
+    return AnswerDistribution(
+        probabilities=dict(probabilities),
+        entropy=entropy,
+        prediction=prediction,
+    )
+
+
+def yes_no_candidates() -> tuple[AnswerCandidate, AnswerCandidate]:
+    """Standard StrategyQA-style answer set."""
+
+    return (
+        AnswerCandidate("Yes", aliases=("yes", "YES")),
+        AnswerCandidate("No", aliases=("no", "NO")),
+    )
+
+
+def multiple_choice_candidates(
+    labels: Iterable[str] = ("A", "B", "C", "D", "E"),
+) -> tuple[AnswerCandidate, ...]:
+    """Create option candidates with common first-token answer spellings."""
+
+    candidates = []
+    for raw_label in labels:
+        label = raw_label.strip()
+        if not label:
+            continue
+        candidates.append(
+            AnswerCandidate(
+                label,
+                aliases=(
+                    label.lower(),
+                    f"({label})",
+                    f"{label}.",
+                    f"{label})",
+                ),
+            )
+        )
+    return normalize_candidates(candidates)
+
+
+def integer_candidates(start: int = 0, end: int = 100) -> tuple[AnswerCandidate, ...]:
+    """Numeric answer candidates for math-style tasks."""
+
+    if start > end:
+        raise ValueError("start cannot exceed end")
+    return tuple(AnswerCandidate(str(value)) for value in range(start, end + 1))
+
+
+_NUMBER_RE = re.compile(r"[-+]?(?:\d{1,3}(?:,\d{3})+|\d+)(?:\.\d+)?")
+
+
+def numeric_candidates_from_texts(
+    texts: Iterable[str],
+    extra_integers: tuple[int, int] | None = (0, 100),
+    max_candidates: int | None = None,
+) -> tuple[AnswerCandidate, ...]:
+    """Build a numeric candidate set from answer strings plus an integer range."""
+
+    labels: dict[str, None] = {}
+    for text in texts:
+        for match in _NUMBER_RE.findall(text):
+            labels[match.replace(",", "")] = None
+
+    if extra_integers is not None:
+        start, end = extra_integers
+        for value in range(start, end + 1):
+            labels[str(value)] = None
+
+    sorted_labels = sorted(labels, key=_numeric_sort_key)
+    if max_candidates is not None:
+        sorted_labels = sorted_labels[:max_candidates]
+    return normalize_candidates([AnswerCandidate(label) for label in sorted_labels])
+
+
+def _numeric_sort_key(label: str) -> tuple[int, float | str]:
+    try:
+        return (0, float(label))
+    except ValueError:
+        return (1, label)
+
+
+def format_reasoning(steps: Sequence[HaltCoTStep]) -> str:
+    """Render generated steps for logs, demos, and model cards."""
+
+    return "\n".join(f"Step {step.index}: {step.text}".rstrip() for step in steps)

halt_cot/transformers_backend.py

@@ -0,0 +1,291 @@
+"""Hugging Face Transformers backend for HALT-CoT."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+import math
+from typing import Sequence
+
+from .core import (
+    AnswerCandidate,
+    AnswerDistribution,
+    EntropyHaltingController,
+    HaltCoTConfig,
+    HaltCoTResult,
+    HaltCoTStep,
+    answer_distribution_from_scores,
+    format_reasoning,
+    normalize_candidates,
+)
+
+
+class TextStopCriteria:
+    """Transformers stopping criterion that halts when generated text hits a marker."""
+
+    def __init__(
+        self,
+        tokenizer,
+        prompt_length: int,
+        stop_strings: Sequence[str],
+        min_new_tokens: int,
+    ):
+        from transformers import StoppingCriteria
+
+        class _Criterion(StoppingCriteria):
+            def __call__(inner_self, input_ids, scores, **kwargs) -> bool:
+                new_len = input_ids.shape[-1] - prompt_length
+                if new_len < min_new_tokens:
+                    return False
+                tail_start = max(prompt_length, input_ids.shape[-1] - 64)
+                tail = tokenizer.decode(
+                    input_ids[0, tail_start:],
+                    skip_special_tokens=False,
+                )
+                return any(stop in tail for stop in stop_strings)
+
+        self.criterion = _Criterion()
+
+
+class HaltCoTForCausalLM:
+    """Run HALT-CoT with any causal language model exposing next-token logits."""
+
+    def __init__(self, model, tokenizer, config: HaltCoTConfig | None = None):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.config = config or HaltCoTConfig()
+
+        if getattr(self.tokenizer, "pad_token_id", None) is None:
+            eos_token = getattr(self.tokenizer, "eos_token", None)
+            if eos_token is not None:
+                self.tokenizer.pad_token = eos_token
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id: str,
+        *,
+        config: HaltCoTConfig | None = None,
+        torch_dtype: str | None = "auto",
+        device_map: str | None = None,
+        trust_remote_code: bool = False,
+        **kwargs,
+    ) -> "HaltCoTForCausalLM":
+        """Load a Transformers causal LM and tokenizer."""
+
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_id,
+            trust_remote_code=trust_remote_code,
+        )
+
+        model_kwargs = dict(kwargs)
+        if torch_dtype is not None:
+            model_kwargs["torch_dtype"] = torch_dtype
+        if device_map is not None:
+            model_kwargs["device_map"] = device_map
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            trust_remote_code=trust_remote_code,
+            **model_kwargs,
+        )
+        if device_map is None:
+            model.to(_default_device())
+        model.eval()
+        return cls(model=model, tokenizer=tokenizer, config=config)
+
+    def run(
+        self,
+        question: str,
+        candidates: Sequence[str | AnswerCandidate],
+        config: HaltCoTConfig | None = None,
+    ) -> HaltCoTResult:
+        """Generate reasoning steps until answer entropy crosses the threshold."""
+
+        run_config = config or self.config
+        answer_candidates = normalize_candidates(candidates)
+        controller = EntropyHaltingController(run_config)
+        context = run_config.prompt_template.format(question=question)
+        steps: list[HaltCoTStep] = []
+        generated_tokens = 0
+        last_distribution: AnswerDistribution | None = None
+
+        for index in range(1, run_config.max_steps + 1):
+            step_prompt = _ensure_trailing_newline(context) + run_config.step_prefix.format(step=index)
+            step_text, token_count = self.generate_step(step_prompt, run_config)
+            generated_tokens += token_count
+            step_text = step_text.strip()
+            context = f"{step_prompt}{step_text}\n"
+
+            distribution = self.answer_distribution(
+                context + run_config.answer_probe,
+                answer_candidates,
+                entropy_unit=run_config.entropy_unit,
+            )
+            last_distribution = distribution
+            decision = controller.observe(distribution.entropy, index)
+
+            steps.append(
+                HaltCoTStep(
+                    index=index,
+                    text=step_text,
+                    entropy=distribution.entropy,
+                    prediction=distribution.prediction,
+                    probabilities=distribution.probabilities,
+                    halted=decision.should_halt,
+                    generated_tokens=token_count,
+                )
+            )
+
+            if decision.should_halt:
+                return HaltCoTResult(
+                    answer=distribution.prediction,
+                    halted=True,
+                    steps=tuple(steps),
+                    reasoning=format_reasoning(steps),
+                    generated_tokens=generated_tokens,
+                )
+
+        if last_distribution is None:
+            raise RuntimeError("HALT-CoT run produced no steps")
+
+        return HaltCoTResult(
+            answer=last_distribution.prediction,
+            halted=False,
+            steps=tuple(steps),
+            reasoning=format_reasoning(steps),
+            generated_tokens=generated_tokens,
+        )
+
+    def generate_step(self, prompt: str, config: HaltCoTConfig) -> tuple[str, int]:
+        """Generate one bounded reasoning step."""
+
+        import torch
+        from transformers import StoppingCriteriaList
+
+        inputs = self._encode(prompt)
+        prompt_length = inputs["input_ids"].shape[-1]
+        stopping_criteria = None
+        if config.step_stop_strings:
+            stopping_criteria = StoppingCriteriaList(
+                [
+                    TextStopCriteria(
+                        self.tokenizer,
+                        prompt_length,
+                        config.step_stop_strings,
+                        config.step_min_new_tokens,
+                    ).criterion
+                ]
+            )
+
+        generation_kwargs = {
+            "max_new_tokens": config.step_max_new_tokens,
+            "do_sample": config.do_sample,
+            "pad_token_id": self.tokenizer.pad_token_id,
+            "eos_token_id": self.tokenizer.eos_token_id,
+            "stopping_criteria": stopping_criteria,
+        }
+        if config.do_sample:
+            generation_kwargs["temperature"] = max(config.temperature, 1e-6)
+            generation_kwargs["top_p"] = config.top_p
+
+        with torch.inference_mode():
+            output_ids = self.model.generate(**inputs, **generation_kwargs)[0]
+
+        new_ids = output_ids[prompt_length:]
+        text = self.tokenizer.decode(new_ids, skip_special_tokens=True).lstrip()
+        text = _truncate_at_first_stop(text, config.step_stop_strings)
+        return text, int(new_ids.shape[-1])
+
+    def answer_distribution(
+        self,
+        context: str,
+        candidates: Sequence[AnswerCandidate],
+        *,
+        entropy_unit: str = "bits",
+    ) -> AnswerDistribution:
+        """Compute answer entropy from next-token logits over candidates."""
+
+        import torch
+
+        candidate_token_ids = first_token_ids_by_candidate(self.tokenizer, candidates)
+        inputs = self._encode(context)
+        with torch.inference_mode():
+            logits = self.model(**inputs).logits[0, -1, :].float()
+
+        scores: dict[str, float] = {}
+        for candidate in candidates:
+            token_ids = candidate_token_ids[candidate.label]
+            token_logits = logits[token_ids]
+            scores[candidate.label] = float(torch.logsumexp(token_logits, dim=0).item())
+
+        return answer_distribution_from_scores(scores, entropy_unit=entropy_unit)
+
+    def _encode(self, text: str):
+        inputs = self.tokenizer(text, return_tensors="pt")
+        return inputs.to(_model_device(self.model))
+
+
+def first_token_ids_by_candidate(
+    tokenizer,
+    candidates: Sequence[AnswerCandidate],
+) -> dict[str, list[int]]:
+    """Map each candidate to first-token ids for its label and aliases."""
+
+    mapping: dict[str, list[int]] = {}
+    for candidate in candidates:
+        token_ids: set[int] = set()
+        for text in candidate.texts:
+            for variant in _candidate_variants(text):
+                encoded = tokenizer.encode(variant, add_special_tokens=False)
+                if encoded:
+                    token_ids.add(int(encoded[0]))
+        if not token_ids:
+            raise ValueError(f"Candidate {candidate.label!r} did not tokenize to any ids")
+        mapping[candidate.label] = sorted(token_ids)
+    return mapping
+
+
+def with_config(base: HaltCoTConfig, **updates) -> HaltCoTConfig:
+    """Return a modified config while preserving dataclass validation."""
+
+    return replace(base, **updates)
+
+
+def _candidate_variants(text: str) -> tuple[str, ...]:
+    stripped = text.strip()
+    variants = {stripped, f" {stripped}"}
+    if stripped[:1].isalpha():
+        variants.add(stripped.lower())
+        variants.add(stripped.upper())
+        variants.add(f" {stripped.lower()}")
+        variants.add(f" {stripped.upper()}")
+    return tuple(variant for variant in variants if variant)
+
+
+def _truncate_at_first_stop(text: str, stop_strings: Sequence[str]) -> str:
+    if not stop_strings:
+        return text
+    indexes = [text.find(stop) for stop in stop_strings if text.find(stop) >= 0]
+    if not indexes:
+        return text
+    return text[: min(indexes)]
+
+
+def _ensure_trailing_newline(text: str) -> str:
+    return text if text.endswith("\n") else f"{text}\n"
+
+
+def _model_device(model):
+    try:
+        return next(model.parameters()).device
+    except StopIteration:
+        return _default_device()
+
+
+def _default_device():
+    import torch
+
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")

pyproject.toml

@@ -0,0 +1,54 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "halt-cot"
+version = "0.1.0"
+description = "Entropy-based early stopping for chain-of-thought reasoning."
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT License" }
+authors = [{ name = "Yassir Laaouach" }]
+keywords = ["chain-of-thought", "entropy", "early-stopping", "reasoning", "transformers"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = []
+
+[project.optional-dependencies]
+transformers = [
+  "torch>=2.2",
+  "transformers>=4.44",
+  "accelerate>=0.33",
+]
+demo = [
+  "torch>=2.2",
+  "transformers>=4.44",
+  "accelerate>=0.33",
+  "gradio>=4.44",
+  "huggingface_hub>=0.24",
+]
+publish = [
+  "huggingface_hub>=0.24",
+]
+dev = [
+  "pytest>=8",
+]
+
+[project.scripts]
+halt-cot = "halt_cot.cli:main"
+
+[tool.setuptools.packages.find]
+include = ["halt_cot*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q"

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.2
+transformers>=4.44
+accelerate>=0.33
+gradio>=4.44
+huggingface_hub>=0.24

scripts/publish_to_hf.py

@@ -0,0 +1,63 @@
+"""Upload the current HALT-CoT folder to Hugging Face."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from huggingface_hub import create_repo, upload_folder
+
+
+DEFAULT_IGNORE = [
+    ".git",
+    ".git/*",
+    ".pytest_cache",
+    ".pytest_cache/*",
+    "__pycache__",
+    "**/__pycache__/*",
+    "*.pyc",
+    "*.egg-info",
+    "*.egg-info/*",
+    "build",
+    "build/*",
+    "dist",
+    "dist/*",
+    ".env",
+    ".env.*",
+]
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Publish HALT-CoT to Hugging Face.")
+    parser.add_argument("--repo-id", required=True, help="Example: username/halt-cot")
+    parser.add_argument("--repo-type", choices=("space", "model", "dataset"), default="space")
+    parser.add_argument("--private", action="store_true")
+    parser.add_argument("--commit-message", default="Publish HALT-CoT")
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    project_root = Path(__file__).resolve().parents[1]
+
+    create_repo(
+        repo_id=args.repo_id,
+        repo_type=args.repo_type,
+        private=args.private,
+        exist_ok=True,
+        space_sdk="gradio" if args.repo_type == "space" else None,
+    )
+    upload_folder(
+        repo_id=args.repo_id,
+        repo_type=args.repo_type,
+        folder_path=str(project_root),
+        path_in_repo=".",
+        ignore_patterns=DEFAULT_IGNORE,
+        commit_message=args.commit_message,
+    )
+    print(f"Uploaded {project_root} to {args.repo_type} repo {args.repo_id}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/test_core.py

@@ -0,0 +1,92 @@
+import math
+
+import pytest
+
+from halt_cot import (
+    AnswerCandidate,
+    EntropyHaltingController,
+    HaltCoTConfig,
+    answer_distribution_from_scores,
+    entropy_from_probabilities,
+    integer_candidates,
+    multiple_choice_candidates,
+    normalize_candidates,
+    numeric_candidates_from_texts,
+    yes_no_candidates,
+)
+
+
+def test_uniform_binary_entropy_is_one_bit():
+    entropy = entropy_from_probabilities({"Yes": 0.5, "No": 0.5}, entropy_unit="bits")
+    assert entropy == pytest.approx(1.0)
+
+
+def test_answer_distribution_uses_softmax_and_argmax():
+    distribution = answer_distribution_from_scores({"A": 4.0, "B": 0.0}, entropy_unit="bits")
+    assert distribution.prediction == "A"
+    assert distribution.probabilities["A"] > 0.98
+    assert 0.0 < distribution.entropy < 0.2
+
+
+def test_entropy_validates_probability_mass():
+    with pytest.raises(ValueError, match="sum to 1"):
+        entropy_from_probabilities({"A": 0.2, "B": 0.2})
+
+
+def test_halting_controller_requires_consecutive_low_entropy_steps():
+    controller = EntropyHaltingController(
+        HaltCoTConfig(theta=0.5, consecutive_low_entropy=2, max_steps=4)
+    )
+    first = controller.observe(entropy=0.4, step_index=1)
+    second = controller.observe(entropy=0.3, step_index=2)
+    assert not first.should_halt
+    assert second.should_halt
+    assert second.low_entropy_streak == 2
+
+
+def test_halting_controller_respects_min_steps():
+    controller = EntropyHaltingController(
+        HaltCoTConfig(theta=0.5, min_steps=3, consecutive_low_entropy=1, max_steps=4)
+    )
+    assert not controller.observe(entropy=0.1, step_index=1).should_halt
+    assert not controller.observe(entropy=0.1, step_index=2).should_halt
+    assert controller.observe(entropy=0.1, step_index=3).should_halt
+
+
+def test_candidate_normalization_rejects_duplicates():
+    with pytest.raises(ValueError, match="Duplicate"):
+        normalize_candidates(["A", "a"])
+
+
+def test_yes_no_candidates_include_aliases():
+    yes, no = yes_no_candidates()
+    assert yes.label == "Yes"
+    assert "YES" in yes.aliases
+    assert no.label == "No"
+
+
+def test_multiple_choice_candidates_include_common_forms():
+    candidates = multiple_choice_candidates(["A", "B"])
+    assert candidates[0] == AnswerCandidate("A", aliases=("a", "(A)", "A.", "A)"))
+    assert candidates[1].label == "B"
+
+
+def test_integer_candidates_are_inclusive():
+    candidates = integer_candidates(2, 4)
+    assert [candidate.label for candidate in candidates] == ["2", "3", "4"]
+
+
+def test_numeric_candidates_from_texts_extracts_numbers_and_range():
+    candidates = numeric_candidates_from_texts(
+        ["The answer is 1,234.5", "Then subtract -7."],
+        extra_integers=(0, 2),
+    )
+    labels = [candidate.label for candidate in candidates]
+    assert "-7" in labels
+    assert "1234.5" in labels
+    assert {"0", "1", "2"}.issubset(labels)
+
+
+def test_config_validation():
+    with pytest.raises(ValueError, match="theta"):
+        HaltCoTConfig(theta=-math.inf)