File size: 3,675 Bytes
7845694 | 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 | """Create-only JSONL command line scorer."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from .core import load_causal_model, validate_row
from .direct import score as direct_score
from .reranker import score as reranker_score
from .serial import SerialPrefixScorer
from .shared import score_shared
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--mode", choices=("direct", "serial", "shared", "reranker"), required=True)
parser.add_argument("--backend", choices=("torch", "mlx"), default="torch")
parser.add_argument("--mlx-bits", type=int, choices=(4, 8), help="Quantize MLX weights in memory; default preserves source precision")
parser.add_argument("--mlx-cache-limit-mib", type=int,
help="MLX inactive allocation cache in MiB (default: 256; 0 disables caching)")
parser.add_argument("--model", required=True)
parser.add_argument("--revision", required=True)
parser.add_argument("--input", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--max-tokens", type=int, default=4096)
args = parser.parse_args()
if args.output.exists() or args.max_tokens < 1:
parser.error("Output must be new and max-tokens must be positive")
if args.mlx_bits and args.backend != "mlx":
parser.error("--mlx-bits requires --backend mlx")
if args.mlx_cache_limit_mib is not None:
if args.backend != "mlx":
parser.error("--mlx-cache-limit-mib requires --backend mlx")
if args.mlx_cache_limit_mib < 0:
parser.error("--mlx-cache-limit-mib must be nonnegative")
if args.backend == "mlx" and args.mode == "reranker":
parser.error("MLX supports direct, serial, and shared modes; reranker requires torch")
rows = [json.loads(line) for line in args.input.read_text().splitlines() if line.strip()]
if not rows:
parser.error("Input is empty")
for row in rows:
validate_row(row)
direct, serial, shared = direct_score, SerialPrefixScorer, score_shared
if args.backend == "mlx":
from . import mlx_backend
cache_limit_mib = (mlx_backend.DEFAULT_CACHE_LIMIT_MIB if args.mlx_cache_limit_mib is None
else args.mlx_cache_limit_mib)
model, tokenizer, metadata = mlx_backend.load_model(
args.model, args.revision, args.mlx_bits, cache_limit_mib=cache_limit_mib)
direct, serial, shared = mlx_backend.score, mlx_backend.SerialPrefixScorer, mlx_backend.score_shared
else:
model, tokenizer, metadata = load_causal_model(args.model, args.revision)
args.output.parent.mkdir(parents=True, exist_ok=True)
with args.output.open("x") as destination:
if args.mode == "shared":
results, timing = shared(model, tokenizer, rows, metadata, args.max_tokens)
for result in results:
destination.write(json.dumps({**result, "shared_timing": timing}, allow_nan=False) + "\n")
elif args.mode == "serial":
scorer = serial(model, tokenizer, metadata, args.max_tokens)
for row in rows:
destination.write(json.dumps(scorer.score(row), allow_nan=False) + "\n")
destination.flush()
else:
scorer = direct if args.mode == "direct" else reranker_score
for row in rows:
destination.write(json.dumps(scorer(model, tokenizer, row, metadata, args.max_tokens), allow_nan=False) + "\n")
destination.flush()
if __name__ == "__main__":
main()
|