File size: 5,084 Bytes
cbc33fe | 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 | """
vLLM generation wrapper shared by pool building, eval, and smoke tests.
Captures per-story cumulative logprob, which E4's `divpo-prob` variant needs:
DivPO's probability criterion picks the LOWEST-logprob passing story as the
diverse chosen and the HIGHEST-logprob failing story as the common rejected.
The highest-probability sample in a temperature pool is, by construction, the
near-greedy one -- which is why this is the principled version of "reject the
greedy decode".
Length normalization matters here. Raw cumulative logprob scales with token
count, so ranking by it would just rank by length (short stories win). We keep
both `cumlogprob` and `mean_logprob = cumlogprob / n_tokens` and use the mean
for DivPO ranking.
"""
from __future__ import annotations
import json
from dataclasses import dataclass, asdict
from pathlib import Path
# NEVER truncate a story. The prior run died of a max_new_tokens wall, so this
# is sized for real headroom, not for the target length: a 500-word story is
# ~650-700 Qwen tokens, so 768 would clip anything that runs long and would
# reproduce the exact bug we are fixing. 1280 lets the model finish and overrun
# a little; overrun is then DETECTED by the >600-word gate rather than silently
# chopped. Measure the failure, don't manufacture it.
#
# Nothing downstream truncates either: the judge receives the complete story
# text, and embeddings are computed over the complete story text.
MAX_NEW_TOKENS = 1280
MAX_MODEL_LEN = 2048
@dataclass
class Gen:
prompt_id: str
prompt: str
idx: int
text: str
n_tokens: int
cumlogprob: float
mean_logprob: float
finish_reason: str
def as_dict(self) -> dict:
return asdict(self)
def build_llm(
model: str,
gpu_mem_util: float = 0.85,
max_model_len: int = MAX_MODEL_LEN,
seed: int = 0,
enable_lora: bool = False,
):
from vllm import LLM
kw = dict(
model=model,
dtype="bfloat16",
gpu_memory_utilization=gpu_mem_util,
max_model_len=max_model_len,
seed=seed,
enforce_eager=False,
disable_log_stats=True,
)
if enable_lora:
kw.update(enable_lora=True, max_lora_rank=32)
return LLM(**kw)
def sampling_params(
n: int, temperature: float, top_p: float, seed: int | None,
max_tokens: int = MAX_NEW_TOKENS,
):
from vllm import SamplingParams
return SamplingParams(
n=n,
temperature=temperature,
top_p=top_p,
max_tokens=max_tokens,
seed=seed,
logprobs=0, # cumulative_logprob only; per-token table not needed
skip_special_tokens=True,
)
def render_chat(tokenizer, prompt: str) -> str:
"""Render the chat template, with thinking OFF where the model has it.
Qwen3-8B is a hybrid-thinking model: left at its default it emits a <think>
block before the story, which would (a) eat the token budget, (b) pollute
the embedding with reasoning text, and (c) make the 8B arm incomparable to
the 4B-Instruct arms, which have no thinking mode at all. Qwen3-4B-Instruct
ignores the kwarg, so the same call is correct for both.
"""
from data import chat_messages
msgs = chat_messages(prompt)
try:
return tokenizer.apply_chat_template(
msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False,
)
except TypeError:
return tokenizer.apply_chat_template(
msgs, tokenize=False, add_generation_prompt=True,
)
def generate(
llm,
tokenizer,
prompts: list[dict],
n: int,
temperature: float,
top_p: float,
seed: int | None = None,
max_tokens: int = MAX_NEW_TOKENS,
lora_path: str | None = None,
) -> list[Gen]:
"""prompts = [{'id':..., 'prompt':...}] -> flat list of n*len(prompts) Gens."""
texts = [render_chat(tokenizer, p["prompt"]) for p in prompts]
sp = sampling_params(n, temperature, top_p, seed, max_tokens)
kw = {}
if lora_path:
from vllm.lora.request import LoRARequest
kw["lora_request"] = LoRARequest("adapter", 1, lora_path)
outs = llm.generate(texts, sp, **kw)
gens: list[Gen] = []
for p, out in zip(prompts, outs):
for j, o in enumerate(out.outputs):
ntok = len(o.token_ids)
cum = float(o.cumulative_logprob) if o.cumulative_logprob is not None else 0.0
gens.append(Gen(
prompt_id=p["id"], prompt=p["prompt"], idx=j,
text=o.text.strip(), n_tokens=ntok,
cumlogprob=cum,
mean_logprob=(cum / ntok) if ntok else 0.0,
finish_reason=o.finish_reason or "",
))
return gens
def save_gens(gens: list[Gen], path: str | Path) -> None:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "w") as f:
for g in gens:
f.write(json.dumps(g.as_dict()) + "\n")
def load_gens(path: str | Path) -> list[Gen]:
return [Gen(**json.loads(l)) for l in open(path) if l.strip()]
|