File size: 5,580 Bytes
a4019dd | 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 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | from __future__ import annotations
import json
from dataclasses import asdict
from pathlib import Path
import torch
from .checkpoints import load_token_checkpoint
from .config import PipelineConfig, VerbalizationTokenSetConfig
from .modeling import ModelBundle
def apply_token_pair(bundle: ModelBundle, token_dir: Path) -> None:
input_emb = bundle.model.get_input_embeddings()
ai_ckpt = load_token_checkpoint(token_dir / "ai_token.pt")
human_ckpt = load_token_checkpoint(token_dir / "human_token.pt")
ai_id = bundle.tokenizer.convert_tokens_to_ids(bundle.config.model.ai_token)
human_id = bundle.tokenizer.convert_tokens_to_ids(bundle.config.model.human_token)
with torch.no_grad():
input_emb.weight[ai_id].copy_(ai_ckpt.embedding.to(input_emb.weight.device, dtype=input_emb.weight.dtype))
input_emb.weight[human_id].copy_(
human_ckpt.embedding.to(input_emb.weight.device, dtype=input_emb.weight.dtype)
)
@torch.inference_mode()
def sample_text(
bundle: ModelBundle,
prompt: str,
*,
max_new_tokens: int,
do_sample: bool,
temperature: float,
top_p: float,
) -> str:
input_ids = bundle.tokenizer(prompt, return_tensors="pt", add_special_tokens=False)["input_ids"].to(
bundle.model.device
)
output = bundle.model.generate(
input_ids=input_ids,
max_new_tokens=max_new_tokens,
do_sample=do_sample,
temperature=temperature,
top_p=top_p,
pad_token_id=bundle.tokenizer.eos_token_id,
)
return bundle.tokenizer.decode(output[0, input_ids.shape[1] :], skip_special_tokens=True).strip()
def _chat_prompt(bundle: ModelBundle, content: str) -> str:
return bundle.tokenizer.apply_chat_template(
[{"role": "user", "content": content}],
tokenize=False,
add_generation_prompt=True,
)
@torch.inference_mode()
def verbalize_token(bundle: ModelBundle, token: str, *, max_new_tokens: int, do_sample: bool) -> str:
prompt = _chat_prompt(bundle, f"Describe what qualities of text would be implied by {token}.")
prompt += f"{token} text typically describes text that"
return sample_text(
bundle,
prompt,
max_new_tokens=max_new_tokens,
do_sample=do_sample,
temperature=1.0,
top_p=0.95 if do_sample else 1.0,
)
@torch.inference_mode()
def verbalize_difference(bundle: ModelBundle, *, max_new_tokens: int, do_sample: bool) -> str:
ai_token = bundle.config.model.ai_token
human_token = bundle.config.model.human_token
prompt = _chat_prompt(
bundle,
(
f"What is the difference between {ai_token} text and {human_token} text? "
f"Refer to {ai_token} text as A-type text and {human_token} text as B-type text. "
"Do not discuss the literal token strings; describe the passage qualities they accompany."
),
)
prompt += (
f"{ai_token}, which I will refer to as A-type text, and {human_token}, "
"B-type text, have many similarities and differences."
)
return sample_text(
bundle,
prompt,
max_new_tokens=max_new_tokens,
do_sample=do_sample,
temperature=1.0,
top_p=0.95 if do_sample else 1.0,
)
def verbalize_token_set(
bundle: ModelBundle,
token_set: VerbalizationTokenSetConfig,
config: PipelineConfig,
) -> dict:
apply_token_pair(bundle, token_set.token_dir)
payload = {
"name": token_set.name,
"token_dir": str(token_set.token_dir),
"model_name": config.model.model_name,
"ai_token": config.model.ai_token,
"human_token": config.model.human_token,
"greedy": {
"ai": verbalize_token(
bundle,
config.model.ai_token,
max_new_tokens=config.verbalization.max_new_tokens,
do_sample=False,
),
"human": verbalize_token(
bundle,
config.model.human_token,
max_new_tokens=config.verbalization.max_new_tokens,
do_sample=False,
),
"difference": verbalize_difference(
bundle,
max_new_tokens=config.verbalization.max_new_tokens,
do_sample=False,
),
},
"samples": [],
}
for _ in range(config.verbalization.n_samples):
payload["samples"].append(
{
"ai": verbalize_token(
bundle,
config.model.ai_token,
max_new_tokens=config.verbalization.max_new_tokens,
do_sample=True,
),
"human": verbalize_token(
bundle,
config.model.human_token,
max_new_tokens=config.verbalization.max_new_tokens,
do_sample=True,
),
"difference": verbalize_difference(
bundle,
max_new_tokens=config.verbalization.max_new_tokens,
do_sample=True,
),
}
)
return payload
def save_verbalizations(output_dir: Path, results: list[dict]) -> None:
output_dir.mkdir(parents=True, exist_ok=True)
(output_dir / "verbalizations.json").write_text(json.dumps(results, indent=2), encoding="utf-8")
for result in results:
(output_dir / f"{result['name']}.json").write_text(json.dumps(result, indent=2), encoding="utf-8")
|