File size: 2,156 Bytes
39371ea | 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 | // Transformers.js 4.2.0 exposes top_p but does not apply it in generate().
// Apply nucleus filtering after its temperature/repetition processors and
// before its normal multinomial sampler. No min-p or top-k filter is added.
export function nucleusProcessor(topP = 0.95) {
if (!(topP > 0 && topP <= 1)) throw Error('topP must be in (0, 1].');
let sorted;
return (_inputIds, logits) => {
if (topP === 1) return logits;
const size = logits.dims.at(-1);
sorted ??= new Float32Array(size);
if (sorted.length !== size) sorted = new Float32Array(size);
for (let offset = 0; offset < logits.data.length; offset += size) {
const scores = logits.data.subarray(offset, offset + size);
let max = -Infinity;
for (let i = 0; i < size; i++) max = Math.max(max, scores[i]);
if (!Number.isFinite(max)) throw Error('The model produced invalid sampling scores.');
let total = 0;
for (let i = 0; i < size; i++) total += Math.exp(scores[i] - max);
const tailMass = (1 - topP) * total;
// Most vocabulary entries have negligible mass. Exclude a tail only
// after measuring its entire mass; this is exact, not a top-k shortcut.
let floor = max - 8, count, mass;
do {
count = 0; mass = 0;
for (let i = 0; i < size; i++) {
if (scores[i] >= floor) sorted[count++] = scores[i];
else mass += Math.exp(scores[i] - max);
}
if (mass <= tailMass) break;
floor -= 8;
} while (true);
sorted.subarray(0, count).sort();
let removed = 0;
// Always retain at least one token, including when logits have ties.
while (removed < count - 1) {
const next = mass + Math.exp(sorted[removed] - max);
if (next > tailMass) break;
mass = next; removed++;
}
const cutoff = removed ? sorted[removed - 1] : floor;
let tied = 0;
for (let i = removed - 1; i >= 0 && sorted[i] === cutoff; i--) tied++;
for (let i = 0; i < size; i++) {
if (scores[i] < cutoff || (scores[i] === cutoff && tied-- > 0)) scores[i] = -Infinity;
}
}
return logits;
};
}
|