File size: 5,448 Bytes
5acccb6 | 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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | #!/usr/bin/env python3
"""Run Laya with the original PyTorch checkpoint using the AX650 request schema."""
import argparse
import json
import os
import sys
import time
from pathlib import Path
from typing import Any, Dict, Optional, Tuple
# Avoid importing TensorFlow through Transformers. Some TensorFlow installations can
# delay or deadlock Laya model construction, and TensorFlow is not used here.
os.environ.setdefault("USE_TF", "0")
MODEL_SPECS: Dict[str, Tuple[str, Optional[str]]] = {
"english": ("convaiinnovations/laya", None),
"multilingual": ("convaiinnovations/laya", "multilingual"),
"typed-decisions": ("convaiinnovations/laya", "typed-decisions"),
}
MODEL_NAMES = {
"english": "laya-english",
"multilingual": "laya-multilingual",
"typed-decisions": "laya-typed-decisions",
}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description=(
"Run an original Laya PyTorch checkpoint with the same state/questions "
"JSON schema used by the packaged AX650 runtime."
)
)
parser.add_argument(
"--variant",
required=True,
choices=tuple(MODEL_SPECS),
help="Checkpoint to load.",
)
parser.add_argument(
"--input",
type=Path,
help="Request JSON file. Omit for resident JSON Lines mode on stdin.",
)
parser.add_argument(
"--device",
default="auto",
help="PyTorch device: auto, cpu, cuda, cuda:0, or mps (default: auto).",
)
return parser.parse_args()
def load_agent(variant: str, device: str):
try:
import laya
except ImportError as exc:
raise SystemExit(
"The Python dependencies are missing. Run: "
"python -m pip install -r python/requirements.txt"
) from exc
repo, subfolder = MODEL_SPECS[variant]
selected_device = None if device == "auto" else device
agent = laya.load(
repo,
subfolder=subfolder,
device=selected_device,
)
# Match the fixed sequence and option budgets used by the packaged AXModels.
# The original upstream checkpoints support larger contexts.
agent.cfg["max_len"] = 256
agent.cfg["head_max_len"] = 128
return agent
def validate_request(request: Any) -> Dict[str, Any]:
if not isinstance(request, dict):
raise ValueError("request must be a JSON object")
if "state" not in request:
raise ValueError("request is missing required field: state")
questions = request.get("questions")
if not isinstance(questions, dict) or not questions:
raise ValueError("request.questions must be a non-empty object")
for question_id, question in questions.items():
if not isinstance(question, dict):
raise ValueError(f"question {question_id!r} must be an object")
question_type = question.get("type")
if question_type not in {"choice", "score", "noul"}:
raise ValueError(
f"question {question_id!r} has unsupported type {question_type!r}"
)
if question_type in {"choice", "score"}:
criteria = question.get("criteria")
if not isinstance(criteria, (dict, list)):
raise ValueError(
f"question {question_id!r}.criteria must be an object or list"
)
if not 2 <= len(criteria) <= 4:
raise ValueError(
f"question {question_id!r} must contain 2 to 4 criteria"
)
return request
def predict(agent, variant: str, request: Dict[str, Any]) -> Dict[str, Any]:
request = validate_request(request)
started = time.perf_counter()
result = agent.predict(request["state"], request["questions"])
latency_ms = (time.perf_counter() - started) * 1000.0
# Keep the primary result fields aligned with `axllm run`. Python adds its
# backend and wall-clock timing under `python_runtime`.
result["model"] = MODEL_NAMES[variant]
result["python_runtime"] = {
"backend": "pytorch",
"device": str(agent.device),
"latency_ms": round(latency_ms, 3),
"sequence_length": 256,
"max_options": 4,
}
return result
def run_file(agent, variant: str, input_path: Path) -> None:
request = json.loads(input_path.read_text(encoding="utf-8"))
result = predict(agent, variant, request)
print(json.dumps(result, indent=2, ensure_ascii=False))
def run_json_lines(agent, variant: str) -> None:
for line_number, line in enumerate(sys.stdin, start=1):
line = line.strip()
if not line:
continue
if line == "/exit":
return
try:
request = json.loads(line)
result = predict(agent, variant, request)
print(json.dumps(result, ensure_ascii=False), flush=True)
except Exception as exc: # Keep the resident process available after a bad request.
error = {
"error": str(exc),
"line": line_number,
}
print(json.dumps(error, ensure_ascii=False), flush=True)
def main() -> None:
args = parse_args()
agent = load_agent(args.variant, args.device)
if args.input is not None:
run_file(agent, args.variant, args.input)
else:
run_json_lines(agent, args.variant)
if __name__ == "__main__":
main()
|