File size: 4,330 Bytes
9a3ed32 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """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())
|