Spaces:
Paused
Paused
File size: 5,331 Bytes
751ad26 d0350dd 751ad26 31c489c 751ad26 7cbeabe 751ad26 d0350dd 751ad26 31c489c 751ad26 d0350dd 751ad26 d0350dd 751ad26 d0350dd 751ad26 d0350dd 751ad26 d0350dd 751ad26 | 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 | #!/usr/bin/env python3
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from dotcache.integrations.qwen35 import Qwen35TextHarness # noqa: E402
from engines.live_request import DEFAULT_LIVE_DECODE_STEPS, resolve_live_runtime_settings # noqa: E402
from scripts.space_runner_common import ( # noqa: E402
configure_model_cache_env,
decode_generated_text,
load_request_from_stdin,
print_json,
tok_per_sec_from_latency,
)
from scripts.space_task_prompts import _task_specs # noqa: E402
def _build_exact_length_inputs(harness: Qwen35TextHarness, *, prompt_unit: str, prompt_length: int):
import torch
if harness.tokenizer is None:
raise ValueError("tokenizer is unavailable for exact-length prompt construction")
if prompt_length <= 0:
raise ValueError("prompt_length must be positive")
tokenizer = harness.tokenizer
unit_ids = tokenizer(prompt_unit, add_special_tokens=False)["input_ids"]
if not unit_ids:
raise ValueError("prompt text tokenized to an empty sequence")
token_ids: list[int] = []
if tokenizer.bos_token_id is not None:
token_ids.append(int(tokenizer.bos_token_id))
while len(token_ids) < prompt_length:
token_ids.extend(int(token_id) for token_id in unit_ids)
token_ids = token_ids[:prompt_length]
device = harness.adapter.device
input_ids = torch.tensor([token_ids], dtype=torch.long, device=device)
attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=device)
return input_ids, attention_mask
def _task_prompt_inputs(harness: Qwen35TextHarness, settings):
task_args = argparse.Namespace(
max_new_tokens_retrieval=64,
max_new_tokens_reasoning=64,
max_new_tokens_instruction=32,
)
task_specs = _task_specs(
harness,
prompt_length=settings.context_length,
args=task_args,
)
for task_spec in task_specs:
if task_spec["task_name"] == settings.compact_task_name:
return task_spec
raise ValueError(f"Unsupported compact-task replay task: {settings.compact_task_name}")
def main() -> int:
configure_model_cache_env()
request = load_request_from_stdin()
settings = resolve_live_runtime_settings(
request,
decode_steps=int(os.getenv("DOTCACHE_SPACE_LIVE_DECODE_STEPS", str(DEFAULT_LIVE_DECODE_STEPS))),
max_live_context=int(os.getenv("DOTCACHE_SPACE_MAX_LIVE_CONTEXT", "4096")),
)
harness = Qwen35TextHarness.from_pretrained(
settings.model_id,
device=os.getenv("DOTCACHE_SPACE_DEVICE"),
torch_dtype=os.getenv("DOTCACHE_SPACE_TORCH_DTYPE", "float16"),
weight_quantization=os.getenv("DOTCACHE_SPACE_WEIGHT_QUANTIZATION", "none"),
)
if settings.is_custom_prompt:
input_ids, attention_mask = harness.tokenize_prompt(settings.prompt_text)
prompt_length = int(input_ids.shape[1])
if prompt_length > settings.context_length:
raise ValueError(
f"Custom prompt tokenized to {prompt_length} tokens, which exceeds the selected context limit "
f"of {settings.context_length}. Increase Context length or shorten the prompt."
)
decode_steps = int(settings.decode_steps)
elif settings.benchmark_suite == "compact_task":
task_spec = _task_prompt_inputs(harness, settings)
input_ids = task_spec["input_ids"]
attention_mask = task_spec["attention_mask"]
decode_steps = int(task_spec["decode_steps"])
elif settings.use_exact_length_prompt:
input_ids, attention_mask = _build_exact_length_inputs(
harness,
prompt_unit=settings.prompt_text,
prompt_length=settings.context_length,
)
decode_steps = int(settings.decode_steps)
else:
raise ValueError(
"Live replay for this benchmark section is not wired in the Space yet. "
"Use the preset-backed compare for the valid benchmark row."
)
record = harness.generate_greedy(
input_ids=input_ids,
attention_mask=attention_mask,
max_new_tokens=decode_steps + 1,
)
generated_ids = list(record.get("dense_generated_ids") or [])
text = decode_generated_text(harness.tokenizer, generated_ids, limit=decode_steps)
latency = float(record.get("dense_decode_ms_per_step") or 0.0)
prefill_ms = float(record.get("prefill_ms") or 0.0)
payload = {
"text": text,
"tok_per_sec": tok_per_sec_from_latency(latency),
"latency_ms_per_token": latency,
"kv_bytes": int(record.get("dense_final_cache_bytes") or 0),
"trace": [
{"name": "prefill_ms", "value": prefill_ms, "unit": "ms"},
{"name": "prompt_length", "value": int(record.get("prompt_length") or input_ids.shape[1]), "unit": "tokens"},
{
"name": "decode_steps",
"value": int(record.get("decode_steps") or decode_steps),
"unit": "tokens",
},
],
}
print_json(payload)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|