124 lines
4.2 KiB
Python
124 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import asdict, dataclass
|
|
from math import ceil
|
|
from pathlib import Path
|
|
|
|
|
|
BLOCK_TOKEN_BUDGET = 24
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SmallAppendTraceConfig:
|
|
output_path: Path
|
|
session_count: int = 8
|
|
turns_per_session: int = 3
|
|
initial_input_length: int = 10_000
|
|
append_input_length: int = 1_000
|
|
output_length: int = 1_000
|
|
inter_turn_gap_s: float = 1.0
|
|
session_stagger_s: float = 0.1
|
|
request_type: str = "coder"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SmallAppendTraceSummary:
|
|
output_path: str
|
|
session_count: int
|
|
turns_per_session: int
|
|
request_count: int
|
|
initial_input_length: int
|
|
append_input_length: int
|
|
output_length: int
|
|
inter_turn_gap_s: float
|
|
session_stagger_s: float
|
|
|
|
|
|
def write_small_append_trace(config: SmallAppendTraceConfig) -> SmallAppendTraceSummary:
|
|
if config.session_count <= 0:
|
|
raise ValueError("session_count must be > 0")
|
|
if config.turns_per_session <= 0:
|
|
raise ValueError("turns_per_session must be > 0")
|
|
if config.initial_input_length < 0:
|
|
raise ValueError("initial_input_length must be >= 0")
|
|
if config.append_input_length < 0:
|
|
raise ValueError("append_input_length must be >= 0")
|
|
if config.output_length < 0:
|
|
raise ValueError("output_length must be >= 0")
|
|
|
|
config.output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
records: list[dict[str, object]] = []
|
|
next_chat_id = 1_000_000
|
|
|
|
for session_idx in range(config.session_count):
|
|
root_chat_id = next_chat_id
|
|
previous_chat_id = -1
|
|
session_base_time = session_idx * config.session_stagger_s
|
|
base_block_count = ceil(config.initial_input_length / BLOCK_TOKEN_BUDGET)
|
|
base_hash_ids = [
|
|
_hash_id_for(session_idx=session_idx, block_idx=block_idx)
|
|
for block_idx in range(base_block_count)
|
|
]
|
|
|
|
for turn_idx in range(config.turns_per_session):
|
|
chat_id = root_chat_id if turn_idx == 0 else next_chat_id
|
|
if turn_idx > 0:
|
|
next_chat_id += 1
|
|
|
|
input_length = config.initial_input_length + turn_idx * (
|
|
config.append_input_length + config.output_length
|
|
)
|
|
total_block_count = ceil(input_length / BLOCK_TOKEN_BUDGET)
|
|
hash_ids = base_hash_ids + [
|
|
_hash_id_for(
|
|
session_idx=session_idx,
|
|
block_idx=base_block_count + append_block_idx,
|
|
)
|
|
for append_block_idx in range(max(0, total_block_count - base_block_count))
|
|
]
|
|
|
|
records.append(
|
|
{
|
|
"chat_id": chat_id,
|
|
"parent_chat_id": previous_chat_id,
|
|
"timestamp": session_base_time
|
|
+ turn_idx * config.inter_turn_gap_s,
|
|
"input_length": input_length,
|
|
"output_length": config.output_length,
|
|
"type": config.request_type,
|
|
"turn": turn_idx + 1,
|
|
"hash_ids": hash_ids,
|
|
}
|
|
)
|
|
previous_chat_id = chat_id
|
|
|
|
next_chat_id += 1
|
|
|
|
records.sort(key=lambda item: float(item["timestamp"]))
|
|
with config.output_path.open("w", encoding="utf-8") as handle:
|
|
for record in records:
|
|
handle.write(json.dumps(record, sort_keys=True) + "\n")
|
|
|
|
summary = SmallAppendTraceSummary(
|
|
output_path=str(config.output_path),
|
|
session_count=config.session_count,
|
|
turns_per_session=config.turns_per_session,
|
|
request_count=len(records),
|
|
initial_input_length=config.initial_input_length,
|
|
append_input_length=config.append_input_length,
|
|
output_length=config.output_length,
|
|
inter_turn_gap_s=config.inter_turn_gap_s,
|
|
session_stagger_s=config.session_stagger_s,
|
|
)
|
|
summary_path = config.output_path.with_suffix(
|
|
config.output_path.suffix + ".summary.json"
|
|
)
|
|
with summary_path.open("w", encoding="utf-8") as handle:
|
|
json.dump(asdict(summary), handle, indent=2, sort_keys=True)
|
|
return summary
|
|
|
|
|
|
def _hash_id_for(*, session_idx: int, block_idx: int) -> int:
|
|
return session_idx * 1_000_000 + block_idx
|