已经遮住补位，为什么短句批量生成还是取错分数？

"""Inspect first-output selection in a tiny, fixed-parameter causal computation.

Standard library only. No trained model, Transformers execution, KV cache,
generation loop or real-language quality measurement is performed.
"""
import csv
import json
from copy import deepcopy
from pathlib import Path

EMBEDDINGS = {"PAD": [0.0, 0.0], "A": [4.0, 0.0], "B": [0.0, 4.0]}
CANDIDATES = ("A", "B")
PROMPTS = {"short": ["A", "B"], "long": ["A", "B", "A", "B"]}


def validate_mask(mask):
    if not mask or any(type(bit) is not int or bit not in (0, 1) for bit in mask):
        raise ValueError("nonempty integer 0/1 mask required")
    if not any(mask):
        raise ValueError("all-PAD prompts have no final content position")


def logical_positions(mask):
    validate_mask(mask)
    count = 0
    result = []
    for bit in mask:
        result.append(count if bit else 0)  # PAD position 0 is this toy's convention.
        count += bit
    return result


def last_valid_index(mask):
    validate_mask(mask)
    return max(i for i, bit in enumerate(mask) if bit == 1)


def forward(tokens, mask, position_policy="logical"):
    validate_mask(mask)
    if len(tokens) != len(mask) or any(token not in EMBEDDINGS for token in tokens):
        raise ValueError("known tokens with matching mask length required")
    # In this fixture PAD has its own identity; real PAD/EOS sharing needs position metadata.
    if any((token != "PAD") != bool(bit) for token, bit in zip(tokens, mask)):
        raise ValueError("toy PAD identity and mask disagree")
    if position_policy not in ("logical", "physical"):
        raise ValueError("unknown position policy")
    positions = logical_positions(mask) if position_policy == "logical" else list(range(len(tokens)))
    x = [[EMBEDDINGS[token][0] + p, EMBEDDINGS[token][1] - p]
         for token, p in zip(tokens, positions)]
    rows = []
    for i, token in enumerate(tokens):
        allowed = [j for j in range(i + 1) if mask[j] == 1]
        # Q=K=0, so every allowed attention score is zero and weights are uniform.
        weights = [1.0 / len(allowed) if j in allowed else 0.0 for j in range(len(tokens))]
        # Empty-key PAD queries are inspection-only rows; do not normalize an empty set.
        context = [sum(x[j][d] for j in allowed) / len(allowed) if allowed else 0.0
                   for d in range(2)]
        logits = [x[i][d] + context[d] for d in range(2)]  # Residual plus identity readout.
        rows.append({"index": i, "token": token, "query_is_content": bool(mask[i]),
                     "position_id": positions[i], "input_vector": x[i],
                     "allowed_key_indices": allowed, "attention_weights": weights,
                     "empty_key_query": not allowed, "context": context, "logits": logits})
    return {"tokens": list(tokens), "mask": list(mask), "position_policy": position_policy,
            "position_ids": positions, "rows": rows}


def pack(prompts, side):
    if side not in ("left", "right") or not prompts:
        raise ValueError("nonempty prompts and left/right padding required")
    if any(not tokens or any(token not in CANDIDATES for token in tokens) for tokens in prompts.values()):
        raise ValueError("original prompts must contain one or more A/B tokens")
    width = max(map(len, prompts.values()))
    result = {}
    for name, tokens in prompts.items():
        pads = ["PAD"] * (width - len(tokens))
        content, empty = [1] * len(tokens), [0] * len(pads)
        result[name] = {"tokens": pads + tokens if side == "left" else tokens + pads,
                        "mask": empty + content if side == "left" else content + empty}
    return result


def inspect_selection(trace, selector):
    # Incorrect selectors are retained only to make their failures visible in the report.
    if selector == "last_valid":
        index = last_valid_index(trace["mask"])
    elif selector == "last_physical":
        index = len(trace["mask"]) - 1
    elif selector == "count_minus_one":
        index = sum(trace["mask"]) - 1
    else:
        raise ValueError("unknown selector")
    row = trace["rows"][index]
    winner = max(range(len(CANDIDATES)), key=lambda i: row["logits"][i])
    return {"selected_index": index, "selected_is_content": row["query_is_content"],
            "logits": row["logits"], "next_candidate": CANDIDATES[winner]}


def rejected(action):
    try:
        action()
    except ValueError:
        return
    raise AssertionError("invalid teaching input was accepted")


def main():
    frozen = deepcopy(PROMPTS)
    traces, results, baselines, csv_rows = {}, {}, {}, []
    for name, tokens in PROMPTS.items():
        trace = forward(tokens, [1] * len(tokens))
        result = inspect_selection(trace, "last_valid")
        traces[f"single/{name}"] = trace
        baselines[name] = result
        csv_rows.append(["single", name, "none", "logical", "last_valid", result["selected_index"],
                         True, *result["logits"], result["next_candidate"], True, 0.0])

    specs = [
        ("right_wrong_last", "right", "logical", "last_physical"),
        ("right_correct_last_valid", "right", "logical", "last_valid"),
        ("left_correct_last", "left", "logical", "last_physical"),
        ("left_wrong_count", "left", "logical", "count_minus_one"),
        ("left_wrong_positions", "left", "physical", "last_physical"),
    ]
    for case, side, position_policy, selector in specs:
        results[case] = {}
        for name, item in pack(PROMPTS, side).items():
            trace = forward(item["tokens"], item["mask"], position_policy)
            result = inspect_selection(trace, selector)
            difference = max(abs(a - b) for a, b in zip(result["logits"], baselines[name]["logits"]))
            result["matches_single_logits"] = difference <= 1e-12
            result["max_logit_difference"] = difference
            traces[f"{case}/{name}"] = trace
            results[case][name] = result
            csv_rows.append([case, name, side, position_policy, selector, result["selected_index"],
                             result["selected_is_content"], *result["logits"], result["next_candidate"],
                             result["matches_single_logits"], difference])

    assert baselines["short"]["logits"] == [3.5, 4.5]
    assert baselines["short"]["next_candidate"] == "B"
    assert results["right_wrong_last"]["short"]["logits"] == [2.5, 1.5]
    assert results["left_wrong_positions"]["short"]["logits"] == [7.5, 0.5]
    assert not results["left_wrong_count"]["short"]["selected_is_content"]
    for case in ("right_wrong_last", "left_wrong_count", "left_wrong_positions"):
        assert not results[case]["short"]["matches_single_logits"]
    for case in ("right_correct_last_valid", "left_correct_last"):
        assert all(r["matches_single_logits"] and r["selected_is_content"] for r in results[case].values())
    assert all(results[case]["long"]["matches_single_logits"] for case, *_ in specs)
    assert last_valid_index([0, 0, 1, 1]) == 3 and sum([0, 0, 1, 1]) - 1 == 1
    assert last_valid_index([1, 1, 0, 0]) == sum([1, 1, 0, 0]) - 1 == 1
    rejected(lambda: forward(["PAD", "PAD"], [0, 0]))
    rejected(lambda: forward(["A"], [True]))
    rejected(lambda: forward(["A"], [1, 0]))
    rejected(lambda: forward(["A", "PAD"], [1, 1]))
    assert PROMPTS == frozen and len(csv_rows) == 12

    parameters = {"prompts": PROMPTS, "embeddings": EMBEDDINGS, "output_candidates": CANDIDATES,
                  "position_vector": "[p, -p]", "Q_and_K": "zero", "V": "input vector",
                  "readout": "input vector plus uniform causal context; identity two-score output",
                  "pad_position_id_in_logical_mode": 0, "argmax_tie_rule": "first candidate A"}
    report = {"scope": "fixed-parameter synthetic first-output calculation, not a full Transformer, framework run or generate()",
              "baselines": baselines, "cases": results, "csv_data_rows": len(csv_rows),
              "limits": ["Only the first next-output selection is inspected; no generation loop or cache.",
                         "Absolute position vector [p,-p] is an explicit toy design, not RoPE.",
                         "All-PAD prompts are rejected; empty-key PAD query rows use zero context for inspection only.",
                         "Real frameworks may have other positions, output shapes and padding contracts."]}
    for filename, value in (("batch-inputs.json", parameters), ("position-traces.json", traces),
                            ("batch-generation-report.json", report)):
        Path(filename).write_text(json.dumps(value, indent=2) + "\n", encoding="utf-8")
    with Path("selection-results.csv").open("w", encoding="utf-8", newline="") as stream:
        writer = csv.writer(stream)
        writer.writerow(["case", "prompt", "padding_side", "position_policy", "selector", "selected_index",
                         "selected_is_content", "logit_A", "logit_B", "next_candidate",
                         "matches_single_logits", "max_logit_difference"])
        writer.writerows(csv_rows)
    print("single short:", baselines["short"]["logits"], baselines["short"]["next_candidate"])
    for case, *_ in specs:
        row = results[case]["short"]
        print(case, row["logits"], row["next_candidate"], "matches_single=" + str(row["matches_single_logits"]))
    print("PASS: two valid batch controls, three short-prompt counterexamples, all-PAD rejected; four files saved")


if __name__ == "__main__":
    main()