import argparse import json import time import numpy as np import torch from common import ROOT, inside, save_json, permit_compute def main(): parser = argparse.ArgumentParser() parser.add_argument('--model', default=str(ROOT / 'models/original')) parser.add_argument('--request', required=True) parser.add_argument('--out', required=True) args = parser.parse_args() permit_compute() from transformers import AutoModelForImageTextToText, AutoTokenizer request = json.loads(inside(args.request).read_text()) out = inside(args.out) if out.exists(): raise ValueError('Используй новый каталог активаций') out.mkdir(parents=True) save_json(out / 'request.json', request) wanted = request.get('return', {}) attention_layers = wanted.get('attn_layers', []) if any(x not in (3, 7, 11, 15, 19, 23) for x in attention_layers): raise ValueError('Attention доступен только для full-attention слоёв') tokenizer = AutoTokenizer.from_pretrained(args.model, local_files_only=True) if 'token_ids' in request: inputs = request['token_ids'] else: texts = request.get('texts') if texts is None: texts = [tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=False, enable_thinking=request.get('enable_thinking', True)) for chat in request['chats']] inputs = [tokenizer.encode(text, add_special_tokens=False) for text in texts] if not inputs or max(map(len, inputs)) > 36864 or sum(map(len, inputs)) > 524288: raise ValueError('Превышен размер teacher-forced запроса') windows = request.get('store_windows') if windows is not None: if len(windows) != len(inputs): raise ValueError('Нужны окна для каждого текста') for tokens, selections in zip(inputs, windows): if not selections or any(not 0 <= left < right <= len(tokens) for left, right in selections): raise ValueError('Некорректные окна сохранения') if windows is not None and attention_layers: raise ValueError('Оконное сохранение attention не поддерживается') if attention_layers and (len(inputs) > 2 or max(map(len, inputs)) > 2048): raise ValueError('Attention ограничен двумя текстами по 2048 токенов') model = AutoModelForImageTextToText.from_pretrained(args.model, local_files_only=True, dtype=torch.bfloat16, attn_implementation='eager' if attention_layers else 'sdpa').eval().to('cuda') model.requires_grad_(False) captures = {} capture_chunk = True handles = [] def capture(name, pre=False, attention=False): def hook(module, arguments, output=None): if not capture_chunk: return if pre: value = arguments[0] elif attention: value = output[1] if isinstance(output, tuple) and len(output) > 1 else None if value is None: raise RuntimeError('Рантайм не вернул attention: нельзя подменять его пустым результатом') else: value = output[0] if isinstance(output, tuple) else output captures[name] = value.detach().to('cpu', dtype=torch.float32).numpy() return hook language = model.model.language_model for layer in wanted.get('hidden_states', []): if not 0 <= layer <= 24: raise ValueError('Некорректный номер residual stream') module = language.embed_tokens if layer == 0 else language.layers[layer - 1] handles.append(module.register_forward_hook(capture(f'hidden_{layer}'))) for layer in wanted.get('mlp_act', []): handles.append(language.layers[layer].mlp.down_proj.register_forward_pre_hook(capture(f'mlp_act_{layer}', pre=True))) for layer in attention_layers: handles.append(language.layers[layer].self_attn.register_forward_hook(capture(f'attention_{layer}', attention=True))) metadata = {'model': str(inside(args.model)), 'lengths': list(map(len, inputs)), 'files': [], 'hidden_definition': '0=embedding, l=raw residual after block l-1 (before final norm)', 'attention_note': 'Only full-attention layers; complete context, no attention chunks', 'logprob_definition': 'entry p scores token at position p+1', 'runtime': 'transformers'} if windows is not None: metadata['store_windows'] = windows metadata['window_note'] = 'Full context is forwarded and cached; only overlapping chunks are saved, with one preceding token for requested log-probabilities.' with torch.inference_mode(): for text_index, tokens in enumerate(inputs): ids = torch.tensor([tokens], dtype=torch.long, device='cuda') np.save(out / f'input_ids_{text_index}.npy', np.asarray(tokens, dtype=np.int64)) cache = None step = len(tokens) if attention_layers else request.get('chunk_size', 256) if not 1 <= step <= 2048: raise ValueError('Некорректный chunk_size') for start in range(0, len(tokens), step): end = min(start + step, len(tokens)) captures.clear() margin = int(wanted.get('token_logprobs', True)) capture_chunk = windows is None or any(start < right and end > max(0, left-margin) for left, right in windows[text_index]) result = model(input_ids=ids[:, start:end], past_key_values=cache, use_cache=True, output_attentions=bool(attention_layers), logits_to_keep=0 if capture_chunk else 1) cache = result.past_key_values n = min(end, len(tokens) - 1) - start if capture_chunk and wanted.get('token_logprobs', True) and n > 0: logits = result.logits[0, :n].float() labels = ids[0, start + 1:start + 1 + n] logp = logits.gather(1, labels[:, None]).squeeze(1) - torch.logsumexp(logits, dim=-1) captures['token_logprobs'] = logp.cpu().numpy() k = int(wanted.get('topk', 0)) if not 0 <= k <= 50: raise ValueError('topk должен быть от 0 до 50') if k: values, indices = logits.topk(k, dim=-1) captures['topk_logprobs'] = (values - torch.logsumexp(logits, dim=-1, keepdim=True)).cpu().numpy() captures['topk_ids'] = indices.cpu().numpy() del logits for name, value in captures.items(): filename = f'text{text_index}_chunk{start}_{name}.npy' np.save(out / filename, value) metadata['files'].append({'name': filename, 'signal': name, 'text': text_index, 'token_start': start, 'token_end': end, 'shape': list(value.shape)}) del result save_json(out / 'meta.json', metadata) for handle in handles: handle.remove() metadata['complete'] = True save_json(out / 'meta.json', metadata) if __name__ == '__main__': main()