| """Fixed-shape input preparation that reuses the upstream laya sequence builder.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import dataclass |
|
|
| import numpy as np |
| from laya.common import QTYPES, build_sequence, render_options |
|
|
| MAX_OPTIONS = 32 |
|
|
|
|
| def package_name(variant: str, length: int, max_options: int, precision: str = "fp16") -> str: |
| return f"laya_{variant}_{precision}_L{length}_options{max_options}" |
|
|
|
|
| @dataclass(frozen=True) |
| class Shape: |
| length: int |
| max_options: int = MAX_OPTIONS |
|
|
|
|
| def to_internal(question: dict) -> dict: |
| """Mirror laya.agent.Agent._to_internal.""" |
| kind = question["type"] |
| criteria = question.get("criteria") |
| if kind == "choice" and isinstance(criteria, list): |
| criteria = {c: None for c in criteria} |
| instructions = question["instructions"] |
| if not isinstance(instructions, str): |
| instructions = json.dumps(instructions) |
| return {"t": kind, "ins": instructions, "crit": criteria} |
|
|
|
|
| def encode(tok, state, question: dict, shape: Shape, head_max_len: int) -> tuple[list[int], list[int], int]: |
| """Token ids, marker positions and question type for one question at the bucket length.""" |
| q = to_internal(question) |
| ids, markers = build_sequence(tok, state, q, shape.length, head_max_len) |
| if len(markers) != len(render_options(q)): |
| raise ValueError(f"options exceed head_max_len={head_max_len}") |
| if len(markers) > shape.max_options: |
| raise ValueError(f"{len(markers)} options exceed the exported capacity {shape.max_options}") |
| return ids, markers, QTYPES[q["t"]] |
|
|
|
|
| def prepare_arrays(ids: list[int], markers: list[int], qtype: int, shape: Shape, pad_id: int) -> dict[str, np.ndarray]: |
| length, k = shape.length, shape.max_options |
| if len(ids) > length: |
| raise ValueError(f"sequence of {len(ids)} tokens exceeds bucket length {length}") |
| input_ids = np.full((1, length), pad_id, dtype=np.int32) |
| input_ids[0, : len(ids)] = ids |
| attention_mask = np.zeros((1, length), dtype=np.int32) |
| attention_mask[0, : len(ids)] = 1 |
| marker_map = np.zeros((1, k, length), dtype=np.float32) |
| for row, position in enumerate(markers): |
| marker_map[0, row, position] = 1.0 |
| question_type = np.zeros((1, 3), dtype=np.float32) |
| question_type[0, qtype] = 1.0 |
| return { |
| "input_ids": input_ids, |
| "attention_mask": attention_mask, |
| "marker_map": marker_map, |
| "question_type": question_type, |
| } |
|
|
|
|
| def prepare_inputs(tok, state, question: dict, shape: Shape, head_max_len: int) -> tuple[dict[str, np.ndarray], int]: |
| ids, markers, qtype = encode(tok, state, question, shape, head_max_len) |
| return prepare_arrays(ids, markers, qtype, shape, tok.pad_token_id), len(markers) |
|
|
