| """Compile Qwen question branches or a shared sequence with per-question readouts.""" |
| from copy import deepcopy |
| from dataclasses import dataclass |
| from pathlib import Path |
| import hashlib |
| import torch |
| from .schema import candidates, target_distribution |
| from .data_types import QuestionSpec |
|
|
|
|
| @dataclass |
| class Branch: |
| inputs: dict |
| candidate_positions: list |
| decision_position: int |
| context_span: tuple | None = None |
| instruction_span: tuple | None = None |
| candidate_spans: list | None = None |
|
|
|
|
| @dataclass(init=False) |
| class Question(QuestionSpec): |
| branches: list |
|
|
| def __init__(self, qid, kind, keys, descriptions, branches, target, is_text): |
| |
| super().__init__(qid, kind, keys, descriptions, target, is_text) |
| self.branches = branches |
|
|
|
|
| @dataclass |
| class CompiledState: |
| questions: list |
| logical_tokens: int |
| compute_tokens: int |
| media: list |
| inputs: dict | None = None |
|
|
|
|
| class Compiler: |
| def __init__(self, processor, media_root=".", max_length=8192, media_kwargs=None, verified_media_hashes=None, |
| execution="question"): |
| if execution not in {"question", "shared_state"}: |
| raise ValueError(f"Unknown qwen_execution: {execution}") |
| self.execution = execution |
| self.processor = processor |
| self.tokenizer = processor.tokenizer |
| self.media_root = Path(media_root) |
| self.max_length = max_length |
| self.media_kwargs = media_kwargs or {} |
| |
| |
| self.verified_media_hashes = verified_media_hashes or {} |
|
|
| def _messages(self, state): |
| messages = [{"role": "user", "content": state}] if isinstance(state, str) else deepcopy(state["messages"]) |
| media = [] |
| for message in messages: |
| if message["role"] not in {"user", "assistant", "system"}: |
| raise ValueError("Unsupported state role") |
| if isinstance(message["content"], str): |
| message["content"] = [{"type": "text", "text": message["content"]}] |
| for item in message["content"]: |
| kind = item["type"] |
| if kind == "text": |
| |
| if any(t in item["text"] for t in self.tokenizer.all_special_tokens): |
| raise ValueError("Reserved tokenizer control token in state") |
| continue |
| if kind not in {"image_url", "video_url"}: |
| raise ValueError(f"Unsupported content type: {kind}") |
| url = item[kind]["url"] |
| if "://" in url: |
| raise ValueError("Materialize remote media locally with a hash before training") |
| path = Path(url) |
| path = path if path.is_absolute() else self.media_root / path |
| path = path.resolve(strict=True) |
| digest = self.verified_media_hashes.get(str(path)) |
| if digest is None: |
| digest = hashlib.sha256(path.read_bytes()).hexdigest() |
| media.append({"path": str(path), "sha256": digest, "type": kind}) |
| item.clear() |
| item.update(type=kind.removesuffix("_url"), path=str(path)) |
| return messages, media |
|
|
| def compile(self, record, rng=None, labeled_only=False): |
| selected = [] |
| for qid, question in record["request"]["questions"].items(): |
| pairs = candidates(question) |
| target = record.get("targets", {}).get(qid) |
| if labeled_only and self.execution == "question" and target_distribution(target, [k for k, _ in pairs]) is None: |
| continue |
| if rng is not None and question["type"] == "choice": |
| rng.shuffle(pairs) |
| selected.append((qid, question, pairs, target)) |
| |
| if not selected or (labeled_only and not any( |
| target_distribution(target, [k for k, _ in pairs]) is not None |
| for _, _, pairs, target in selected)): |
| return CompiledState([], 0, 0, []) |
| messages, media = self._messages(record["request"]["state"]) |
| expected_assets = {str((self.media_root / a["path"]).resolve()): a["sha256"] for a in record.get("assets", [])} |
| for entry in media: |
| if entry["path"] in expected_assets and expected_assets[entry["path"]] != entry["sha256"]: |
| raise ValueError(f"Media hash mismatch: {entry['path']}") |
| base = dict(self.processor.apply_chat_template( |
| messages, tokenize=True, add_generation_prompt=False, return_dict=True, |
| return_tensors="pt", processor_kwargs=dict(self.media_kwargs, return_metadata=True))) |
| |
| metadata = base.pop("video_metadata", None) |
| if metadata is not None: |
| for entry, meta in zip([m for m in media if m["type"] == "video_url"], metadata): |
| entry["sampling"] = dict(meta) if isinstance(meta, dict) else vars(meta) |
| sampling = entry["sampling"] |
| indices = sampling.get("frames_indices") |
| if indices is not None and sampling.get("fps"): |
| sampling["frames_indices"] = [int(i) for i in indices] |
| sampling["timestamps_seconds"] = [int(i) / sampling["fps"] for i in indices] |
| sampling["visible_range_seconds"] = [min(sampling["timestamps_seconds"]), max(sampling["timestamps_seconds"])] |
| for name in ("image_grid_thw", "video_grid_thw"): |
| if name in base: |
| grid = base[name] |
| media.append({name: grid.tolist(), "visual_tokens": int(grid.prod(-1).sum()) // 4}) |
| base_length = base["input_ids"].shape[1] |
| if self.execution == "shared_state": |
| return self._compile_shared(selected, base, media, labeled_only) |
| logical_tokens, compute_tokens = base_length, 0 |
| compiled = [] |
| for qid, question, pairs, target in selected: |
| keys = [k for k, _ in pairs] |
| branches = [] |
| groups = [[pair] for pair in pairs] if question["type"] == "score" else [pairs] |
| for group in groups: |
| suffix = [] |
| def append(text, char_span=None): |
| start = base_length + len(suffix) |
| if char_span is None: |
| suffix.extend(self.tokenizer.encode(text, add_special_tokens=False)) |
| return None |
| |
| encoded = self.tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) |
| suffix.extend(encoded["input_ids"]) |
| indices = [i for i, (a, b) in enumerate(encoded["offset_mapping"]) |
| if a < char_span[1] and b > char_span[0]] |
| if not indices: |
| raise ValueError("Empty token span for decision role") |
| return (start + indices[0], start + indices[-1] + 1) |
| |
| suffix.extend(self.tokenizer.encode("<|im_start|>user\n", add_special_tokens=False)) |
| for value in [question["instructions"]] + [v for pair in group for v in pair]: |
| if any(t in value for t in self.tokenizer.all_special_tokens): |
| raise ValueError("Reserved tokenizer control token in question/criteria") |
| prefix = "Task: " + question["type"] + "\nQuestion: " |
| instruction_span = append(prefix + question["instructions"] + "\nCandidates:\n", |
| (len(prefix), len(prefix) + len(question["instructions"]))) |
| positions, candidate_spans = [], [] |
| for key, description in group: |
| |
| prefix = key + ": " if question["type"] != "score" else "" |
| candidate_spans.append(append(prefix + description, (len(prefix), len(prefix) + len(description)))) |
| positions.append(base_length + len(suffix) - 1) |
| append("\n") |
| append("Decision:") |
| length = base_length + len(suffix) |
| if length > self.max_length: |
| raise ValueError(f"{qid}: {length} tokens exceed max_length={self.max_length}; no truncation") |
| inputs = dict(base) |
| inputs["input_ids"] = torch.cat([base["input_ids"], torch.tensor([suffix], dtype=torch.long)], dim=1) |
| inputs["attention_mask"] = torch.ones_like(inputs["input_ids"]) |
| if "mm_token_type_ids" in base: |
| inputs["mm_token_type_ids"] = torch.cat([base["mm_token_type_ids"], torch.zeros((1, len(suffix)), dtype=torch.long)], dim=1) |
| branches.append(Branch(inputs, positions, length - 1, (0, base_length), |
| instruction_span, candidate_spans)) |
| compute_tokens += length |
| logical_tokens += len(suffix) |
| compiled.append(Question(qid, question["type"], keys, [v for _, v in pairs], branches, |
| target_distribution(target, keys), not media)) |
| return CompiledState(compiled, logical_tokens, compute_tokens, media) |
|
|
| def _compile_shared(self, selected, base, media, labeled_only=False): |
| """All schemas precede all readouts; labels and external IDs never enter inputs. |
| |
| This is ordinary causal attention, not isolated question/Score branches. |
| No answer tokens or LM output layer are required by the decision heads. |
| """ |
| base_length = base["input_ids"].shape[1] |
| suffix, compiled = [], [] |
|
|
| def append(text, span=None): |
| start = base_length + len(suffix) |
| if span is None: |
| suffix.extend(self.tokenizer.encode(text, add_special_tokens=False)) |
| return None |
| encoded = self.tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) |
| suffix.extend(encoded["input_ids"]) |
| indices = [i for i, (a, b) in enumerate(encoded["offset_mapping"]) if a < span[1] and b > span[0]] |
| if not indices: |
| raise ValueError("Empty token span for decision role") |
| return start + indices[0], start + indices[-1] + 1 |
|
|
| append("<|im_start|>user\n") |
| for index, (qid, question, pairs, target) in enumerate(selected): |
| for value in [question["instructions"]] + [v for pair in pairs for v in pair]: |
| if any(t in value for t in self.tokenizer.all_special_tokens): |
| raise ValueError("Reserved tokenizer control token in question/criteria") |
| prefix = f"Question {index + 1}\nTask: {question['type']}\nQuestion: " |
| instruction = append(prefix + question["instructions"] + "\nCandidates:\n", |
| (len(prefix), len(prefix) + len(question["instructions"]))) |
| positions, spans = [], [] |
| for key, description in pairs: |
| prefix = key + ": " if question["type"] != "score" else "" |
| spans.append(append(prefix + description, (len(prefix), len(prefix) + len(description)))) |
| positions.append(base_length + len(suffix) - 1) |
| append("\n") |
| |
| branch = Branch({}, positions, -1, (0, base_length), instruction, spans) |
| keys = [key for key, _ in pairs] |
| compiled.append(Question(qid, question["type"], keys, [v for _, v in pairs], [branch], |
| target_distribution(target, keys), not media)) |
| append("<|im_end|>\n<|im_start|>assistant\n") |
| for index, question in enumerate(compiled): |
| append(f"Question {index + 1} Decision:") |
| question.branches[0].decision_position = base_length + len(suffix) - 1 |
| append("\n") |
| length = base_length + len(suffix) |
| if length > self.max_length: |
| raise ValueError(f"Shared state: {length} tokens exceed max_length={self.max_length}; no truncation") |
| inputs = dict(base) |
| inputs["input_ids"] = torch.cat([base["input_ids"], base["input_ids"].new_tensor([suffix])], dim=1) |
| inputs["attention_mask"] = torch.cat([base["attention_mask"], |
| base["attention_mask"].new_ones((1, len(suffix)))], dim=1) |
| if "mm_token_type_ids" in base: |
| inputs["mm_token_type_ids"] = torch.cat([base["mm_token_type_ids"], |
| base["mm_token_type_ids"].new_zeros((1, len(suffix)))], dim=1) |
| for question in compiled: |
| question.branches[0].inputs = inputs |
| if labeled_only: |
| compiled = [question for question in compiled if question.target is not None] |
| return CompiledState(compiled, length, length, media, inputs) |
|
|
