", "", "", "", "", "", "", "=", ";", "out", "inp"] + OP_NAMES + VAR_NAMES
ST_STOI = {w: i for i, w in enumerate(STRAIGHT_WORDS)}
ST_MAX_LEN = 48
def random_straight(rng, n_lines=None):
"""
一个随机的直线程序（词列表），例如 v3 = rot inp ; v7 = rev v3 ; v1 = dup v3 ; out v1 ;
【思路】每一行定义一个**新的**变量（名字从 v0 … v9 里随机挑没用过的），参数从已经定义的变量（含 inp）里随机挑；
最后 out 最后定义的那个变量。变量名是随机的，所以"下一行定义哪个变量"从前缀里猜不出来。
【参数】n_lines：赋值语句的行数（默认 3–5 随机）
"""
n_lines = n_lines or rng.randint(3, 5)
names = rng.sample(VAR_NAMES, n_lines)
defined, toks = ["inp"], []
for v in names:
toks += [v, "=", rng.choice(OP_NAMES), rng.choice(defined), ";"]
defined.append(v)
return toks + ["out", names[-1], ";"]
def check_straight(toks):
"""
检查一个直线程序能不能运行：每行格式对、每个用到的变量都已定义、没有重复定义、最后有 out。
【返回】(True, None) 或 (False, 原因)
"""
lines, cur = [], []
for t in toks:
if t == ";":
lines.append(cur)
cur = []
else:
cur.append(t)
if cur:
return False, "最后一行没有以 ; 结束"
if not lines or lines[-1][:1] != ["out"]:
return False, "没有 out"
defined = {"inp"}
for ln in lines[:-1]:
if len(ln) != 4 or ln[1] != "=" or ln[2] not in OPS or ln[0] not in VAR_NAMES:
return False, f"语法错误：{' '.join(ln)}"
if ln[3] not in defined:
return False, f"NameError：{ln[3]} 还没有定义"
if ln[0] in defined:
return False, f"{ln[0]} 重复定义"
defined.add(ln[0])
out = lines[-1]
if len(out) != 2 or out[1] not in defined:
return False, "out 的变量没有定义"
return True, None
def encode_l2r(toks):
"""从左到右的普通写法： 程序 。"""
return [ST_STOI[""]] + [ST_STOI[t] for t in toks] + [ST_STOI[""]]
def encode_fim(prefix, middle, suffix):
"""
Bavarian et al. (2022) 的 PSM 格式： 前缀 后缀 中间 。
【思路】把文档切成三段、交换后两段的顺序：模型仍然是从左到右生成，但生成"中间"时已经看过了"后缀"。
【返回】(ids, 中间开始的位置)——训练时只在"中间"和 上计损失也可以，这里按论文对整个序列计损失。
"""
ids = ([ST_STOI[""]] + [ST_STOI[t] for t in prefix] + [ST_STOI[""]] + [ST_STOI[t] for t in suffix]
+ [ST_STOI[""]])
return ids + [ST_STOI[t] for t in middle] + [ST_STOI[""]], len(ids)
def split_lines(toks):
"""把直线程序按 ";" 切成行（每行保留结尾的 ";"）。"""
out, cur = [], []
for t in toks:
cur.append(t)
if t == ";":
out.append(cur)
cur = []
return out
def fim_example(rng, toks, mode="line"):
"""
训练用的随机切分：前缀 / 中间 / 后缀。
mode = "token"：在 token 序列里随机挑两个位置 i ≤ j，中间 = [i:j]（论文是在字符级随机切，思路相同）；
mode = "line" ：在**行边界**上随机挑 a < b，中间 = 第 a … b−1 行（至少一行）。
【返回】(prefix, middle, suffix)
【两种切法的取舍】按行切，训练和"补一整行"的测试完全一致；按 token 随机切，模型见过各种切点，
但很少恰好切在行边界上。Bavarian et al. (2022) 的表 2：按行切的模型补单行最好，但遇到随机切点几乎完全失败；
字符级随机切的模型在随机切点上好得多，补单行也只差一点。03 号脚本在小模型上比较这两种切法。
"""
if mode == "token":
i, j = sorted(rng.randint(0, len(toks)) for _ in range(2))
return toks[:i], toks[i:j], toks[j:]
lines = split_lines(toks)
a = rng.randrange(len(lines))
b = rng.randint(a + 1, len(lines))
return sum(lines[:a], []), sum(lines[a:b], []), sum(lines[b:], [])
def straight_true_nll(toks):
"""
一个直线程序在**生成它的真实分布**下的负对数概率（纳特）和被预测的 token 数（含 ）。
【思路】按 random_straight 的生成过程逐步累加：行数在 3–5 里均匀，所以第 4、5 行开头"继续写还是 out"各有概率；
每行的变量名在没用过的名字里均匀（第 i 行有 10 − i 种）、操作 1/4、参数在已定义的 i + 1 个变量里均匀；
"="、";"、out 后面的变量（就是最后定义的那个）和 都是确定的。
【用途】任何模型从左到右的平均损失都不可能低于"真实分布的熵"——用它判断模型离最优还有多远。
"""
lines = split_lines(toks)
n = len(lines) - 1
nll = 0.0
for i in range(n):
if i == 3:
nll -= math.log(2 / 3)
elif i == 4:
nll -= math.log(1 / 2)
nll -= math.log(1 / (len(VAR_NAMES) - i)) + math.log(1 / len(OP_NAMES)) + math.log(1 / (i + 1))
if n == 3:
nll -= math.log(1 / 3)
elif n == 4:
nll -= math.log(1 / 2)
return nll, len(toks) + 1
def straight_batch(seed, fim_rate, mode="line", batch_size=64):
"""一批语言二的训练数据：每个程序以 fim_rate 的概率写成 FIM 格式（切法见 fim_example），否则写成从左到右的格式；对整个序列计损失。"""
rng = random.Random(seed)
seqs = []
for _ in range(batch_size):
toks = random_straight(rng)
if rng.random() < fim_rate:
seqs.append(encode_fim(*fim_example(rng, toks, mode))[0])
else:
seqs.append(encode_l2r(toks))
T = max(len(s) for s in seqs) - 1
inp = torch.zeros(batch_size, T, dtype=torch.long)
tgt = torch.full((batch_size, T), IGNORE, dtype=torch.long)
for i, s in enumerate(seqs):
n = len(s) - 1
inp[i, :n] = torch.tensor(s[:-1])
tgt[i, :n] = torch.tensor(s[1:])
return inp, tgt
@torch.no_grad()
def greedy_until(model, prompt, stop_id, max_new=12):
"""从一个提示开始贪心生成，直到生成 stop_id 或达到 max_new 个 token。返回生成的 id 列表（不含 stop_id）。"""
ids = torch.tensor([prompt])
out = []
for _ in range(max_new):
nxt = model(ids)[0][0, -1].argmax().item()
if nxt == stop_id:
break
out.append(nxt)
ids = torch.cat([ids, torch.tensor([[nxt]])], dim=1)
return out