File size: 6,004 Bytes
22ca93a | 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 | """Native-tokenization input for a diagnostic experiment: is the Qwen-token input the bottleneck?
The Qwen answer bytes decode losslessly to text; the encoder re-tokenizes that text with its OWN
tokenizer (the segmentation it was pretrained on). Labels move from Qwen tokens to native tokens by
character overlap: a native token gets the label of the first positive Qwen token it overlaps
(B only on the first native token of that Qwen token, I after), 0 if it only overlaps O tokens,
-100 otherwise. Predictions go back to Qwen tokens as the mean over the native tokens that overlap
each Qwen token, so every metric is computed on exactly the same Qwen tokens as the other taggers.
Interface mirrors byt5_adapter.QwenBytes (prompt / answer / special ids).
"""
from __future__ import annotations
import numpy as np
import torch
from transformers import AutoTokenizer
from byt5_adapter import QwenBytes
class NativeText:
def __init__(self, qwen_tokenizer_dir: str, base: str):
self.qb = QwenBytes(qwen_tokenizer_dir)
self.tok = AutoTokenizer.from_pretrained(base)
self.cls = [self.tok.cls_token_id]
self.sep = [self.tok.sep_token_id]
self.pad_id = self.tok.pad_token_id
def prompt(self, query: str) -> list[int]:
return self.cls + self.tok(query, add_special_tokens=False)["input_ids"][:128] + self.sep
def answer(self, token_ids, token_labels=None):
chunks = [self.qb.tb[t] for t in token_ids]
full = b"".join(chunks)
text = full.decode("utf-8", errors="replace")
# byte offset -> char index
char_of_byte, ci = [], 0
for ch in text:
n = len(ch.encode("utf-8")) if ch != "�" else 1
char_of_byte += [ci] * n
ci += 1
char_of_byte.append(ci)
qspan, b = [], 0
for c in chunks:
qspan.append((char_of_byte[b], char_of_byte[b + len(c)] if len(c) else char_of_byte[b]))
b += len(c)
enc = self.tok(text, add_special_tokens=False, return_offsets_mapping=True)
ids, offs = enc["input_ids"], enc["offset_mapping"]
# native token -> overlapping qwen tokens (both sorted by position: two pointers)
per_native, k0 = [], 0
for s, e in offs:
while k0 < len(qspan) and qspan[k0][1] <= s:
k0 += 1
ks, k = [], k0
while k < len(qspan) and qspan[k][0] < max(e, s + 1):
if qspan[k][1] > qspan[k][0]:
ks.append(k)
k += 1
per_native.append(ks)
spans = [[] for _ in token_ids]
for j, ks in enumerate(per_native):
for k in ks:
spans[k].append(j)
labs = []
if token_labels is not None:
seen_bad = set()
for j, ks in enumerate(per_native):
pos = [k for k in ks if token_labels[k] > 0]
if pos:
k = pos[0]
l = token_labels[k]
if l % 2 == 1 and k in seen_bad:
l += 1
if l % 2 == 1:
seen_bad.add(k)
labs.append(l)
elif ks and all(token_labels[k] == 0 for k in ks):
labs.append(0)
else:
labs.append(-100)
return ids, labs, spans
def windows(n: int, a0: int, max_len: int, stride: int):
span = max_len - a0 - 1
starts, s = [], 0
while True:
starts.append(s)
if s + span >= n:
break
s += span - stride
return [(s, min(n, s + span)) for s in starts]
def expand_windows_native(nt: NativeText, rows, max_len, stride, synth_weight, labels_fn, include_synth=True,
include_corrected=True):
out = {"input_ids": [], "labels": [], "weight": []}
for r in rows:
v = r["variant"]
if (v.startswith("synthetic") and not include_synth) or (v == "corrected" and not include_corrected):
continue
a0 = r["answer_start"]
lab = labels_fn(r)
ids, nl, _ = nt.answer(r["input_ids"][a0:], lab[a0:])
if not ids:
continue
p = nt.prompt(r["query"])
w = synth_weight if v.startswith("synthetic") else 1.0
for s, e in windows(len(ids), len(p), max_len, stride):
out["input_ids"].append(p + ids[s:e] + nt.sep)
out["labels"].append([-100] * len(p) + nl[s:e] + [-100])
out["weight"].append(w)
return out
@torch.no_grad()
def predict_row_native(nt: NativeText, model, row, max_len, stride, device, batch=8):
a0 = row["answer_start"]
ids, _, spans = nt.answer(row["input_ids"][a0:])
C = model.config.num_labels
out = np.zeros((len(spans), C), dtype=np.float32)
out[:, 0] = 1.0
if not ids:
return out
p0 = nt.prompt(row["query"])
ws = windows(len(ids), len(p0), max_len, stride)
probs = np.zeros((len(ids), C), dtype=np.float32)
best = np.full(len(ids), -1.0)
for b in range(0, len(ws), batch):
chunk = ws[b:b + batch]
seqs = [p0 + ids[s:e] + nt.sep for s, e in chunk]
L = max(len(x) for x in seqs)
x = torch.full((len(seqs), L), nt.pad_id, dtype=torch.long)
m = torch.zeros((len(seqs), L), dtype=torch.long)
for i, sq in enumerate(seqs):
x[i, :len(sq)] = torch.as_tensor(sq)
m[i, :len(sq)] = 1
logits = model(input_ids=x.to(device), attention_mask=m.to(device)).logits.float()
p = torch.softmax(logits, -1).cpu().numpy()
for i, (s, e) in enumerate(chunk):
pos = np.arange(s, e)
centr = np.minimum(pos - s, e - 1 - pos).astype(float)
upd = centr > best[pos]
probs[pos[upd]] = p[i, len(p0) + (pos[upd] - s)]
best[pos[upd]] = centr[upd]
for k, js in enumerate(spans):
if js:
out[k] = probs[js].mean(0)
return out
|