| """Shared paths, the live arena definition, jsonl helpers, the Granite tokenizer and stage markers.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| import sys |
| import time |
| import urllib.request |
| from collections.abc import Iterable, Iterator |
| from functools import lru_cache |
| from pathlib import Path |
| from typing import Any |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| ARENA_ID = "ctx_compress" |
| ARENA_URL = f"https://api.microtensor.cloud/v2/arenas/{ARENA_ID}" |
| BASE = "ibm-granite/granite-4.0-h-micro" |
| BASE_REVISION = "d5f01a3ea75f088947be3aae039f4ad52837dfde" |
| PROBE_HF = "Qwen/Qwen2.5-1.5B-Instruct" |
| PROBE_HF_REVISION = "989aa7980e4cf806f80c7fef2b1adb7bc71aa306" |
| HEADINGS = ( |
| "Primary request and intent", |
| "Key technical concepts", |
| "Files and code sections", |
| "Errors and fixes", |
| "Problem solving", |
| "All user messages", |
| "Pending tasks", |
| "Current work", |
| "Next step", |
| ) |
| |
| EXPECTED = { |
| "allowed_bases": [BASE], |
| "slm_prompt": ( |
| "<|start_of_role|>system<|end_of_role|>{system}<|end_of_text|>\n" |
| "<|start_of_role|>user<|end_of_role|>{user}<|end_of_text|>\n" |
| "<|start_of_role|>assistant<|end_of_role|>" |
| ), |
| "stop": ["<|end_of_text|>"], |
| "max_tokens": 3072, |
| "pass_threshold": 0.6, |
| } |
|
|
|
|
| def out_dir() -> Path: |
| path = Path(os.environ.get("CCSLM_OUT") or ROOT / "out") |
| path.mkdir(parents=True, exist_ok=True) |
| return path |
|
|
|
|
| def log(message: str) -> None: |
| stamp = time.strftime("%H:%M:%S") |
| print(f"[{stamp}] {message}", file=sys.stderr, flush=True) |
|
|
|
|
| def read_jsonl(path: Path | str) -> list[dict[str, Any]]: |
| with Path(path).open(encoding="utf-8") as handle: |
| return [json.loads(line) for line in handle if line.strip()] |
|
|
|
|
| def iter_jsonl(path: Path | str) -> Iterator[dict[str, Any]]: |
| with Path(path).open(encoding="utf-8") as handle: |
| for line in handle: |
| if line.strip(): |
| yield json.loads(line) |
|
|
|
|
| def write_jsonl(path: Path | str, rows: Iterable[dict[str, Any]]) -> int: |
| target = Path(path) |
| target.parent.mkdir(parents=True, exist_ok=True) |
| count = 0 |
| temp = target.with_suffix(target.suffix + ".tmp") |
| with temp.open("w", encoding="utf-8") as handle: |
| for row in rows: |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") |
| count += 1 |
| temp.replace(target) |
| return count |
|
|
|
|
| def write_json(path: Path | str, value: Any) -> None: |
| target = Path(path) |
| target.parent.mkdir(parents=True, exist_ok=True) |
| target.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") |
|
|
|
|
| def read_json(path: Path | str, default: Any = None) -> Any: |
| target = Path(path) |
| if not target.is_file(): |
| return default |
| return json.loads(target.read_text(encoding="utf-8")) |
|
|
|
|
| def definition(refresh: bool = False) -> dict[str, Any]: |
| """The arena definition, cached under out/arena.json; refresh re-reads the live API.""" |
| cache = out_dir() / "arena.json" |
| if cache.is_file() and not refresh: |
| return json.loads(cache.read_text(encoding="utf-8"))["definition"] |
| with urllib.request.urlopen(ARENA_URL, timeout=30) as response: |
| document = json.loads(response.read().decode("utf-8")) |
| write_json(cache, document) |
| return document["definition"] |
|
|
|
|
| def check_definition(found: dict[str, Any]) -> list[str]: |
| """Differences between the live definition and what this project was built for.""" |
| grading = found.get("grading") or {} |
| generation = found.get("generation") or {} |
| live = { |
| "allowed_bases": found.get("allowed_bases"), |
| "slm_prompt": found.get("slm_prompt"), |
| "stop": generation.get("stop"), |
| "max_tokens": generation.get("max_tokens"), |
| "pass_threshold": grading.get("pass_threshold"), |
| } |
| problems = [f"{key}: built for {EXPECTED[key]!r}, live {value!r}" for key, value in live.items() if value != EXPECTED[key]] |
| headings = grading.get("headings") |
| if headings and tuple(headings) != HEADINGS: |
| problems.append(f"headings changed: {headings}") |
| return problems |
|
|
|
|
| def render_prompt(template: str, task: dict[str, Any]) -> str: |
| """The validator's render_prompt: replace {key} with each task field, in task order.""" |
| rendered = template |
| for key, value in task.items(): |
| rendered = rendered.replace("{" + str(key) + "}", str(value)) |
| return rendered |
|
|
|
|
| def slm_prompt(task: dict[str, Any]) -> str: |
| return render_prompt(EXPECTED["slm_prompt"], task) |
|
|
|
|
| @lru_cache(maxsize=4) |
| def tokenizer(name: str = BASE, revision: str = BASE_REVISION) -> Any: |
| from tokenizers import Tokenizer |
|
|
| local = Path(name) |
| if (local / "tokenizer.json").is_file(): |
| return Tokenizer.from_file(str(local / "tokenizer.json")) |
| return Tokenizer.from_pretrained(name, revision=revision) |
|
|
|
|
| def count_tokens(text: str) -> int: |
| """Token count as the grader counts the output: Granite tokenizer, default special token handling.""" |
| return len(tokenizer().encode(text).ids) |
|
|
|
|
| def base_dir() -> Path: |
| """The local snapshot of the pinned base model, downloaded on first use. |
| |
| CCSLM_BASE_DIR points somewhere else; the tests use it for a tiny model of the same architecture. |
| """ |
| override = os.environ.get("CCSLM_BASE_DIR") |
| if override: |
| return Path(override) |
| from huggingface_hub import snapshot_download |
|
|
| return Path(snapshot_download(BASE, revision=BASE_REVISION)) |
|
|
|
|
| MAMBA_TYPES = ("mamba", "linear_attention") |
| ATTENTION_TYPES = ("attention", "full_attention") |
|
|
|
|
| def portable_config(model_config: dict[str, Any], reference: dict[str, Any]) -> dict[str, Any]: |
| """The base's own config.json with this model's layer list, in the base's naming. |
| |
| transformers 5 renames Granite's layer types (mamba -> linear_attention, attention -> |
| full_attention), writes time_step_limit as {"__float__": "Infinity"} and swaps torch_dtype for |
| dtype when it saves. The serving engine and the validators load the file with their own |
| versions, so a shipped model keeps the base's config and changes only what pruning changed |
| (num_hidden_layers, layer_types) and what quantisation adds (quantization_config). |
| """ |
| out = dict(reference) |
| types = model_config.get("layer_types") or reference["layer_types"] |
| out["layer_types"] = [ |
| "mamba" if t in MAMBA_TYPES else "attention" if t in ATTENTION_TYPES else t for t in types |
| ] |
| out["num_hidden_layers"] = len(out["layer_types"]) |
| if model_config.get("quantization_config"): |
| out["quantization_config"] = model_config["quantization_config"] |
| return out |
|
|
|
|
| def write_portable_config(model_dir: Path, reference_dir: Path | None = None) -> None: |
| reference = json.loads(((reference_dir or base_dir()) / "config.json").read_text(encoding="utf-8")) |
| found = json.loads((model_dir / "config.json").read_text(encoding="utf-8")) |
| write_json(model_dir / "config.json", portable_config(found, reference)) |
|
|
|
|
| class Stages: |
| """Done markers under out/stages so the pipeline can resume.""" |
|
|
| def __init__(self, root: Path | None = None) -> None: |
| self.root = (root or out_dir()) / "stages" |
| self.root.mkdir(parents=True, exist_ok=True) |
|
|
| def done(self, name: str) -> bool: |
| return (self.root / f"{name}.done").is_file() |
|
|
| def mark(self, name: str, info: dict[str, Any] | None = None) -> None: |
| write_json(self.root / f"{name}.done", {"finished": time.time(), **(info or {})}) |
|
|
| def clear(self, name: str) -> None: |
| (self.root / f"{name}.done").unlink(missing_ok=True) |
|
|
