Something went wrong. Try again.
Anonymize your writing style. Zig WASM engine detects authorship markers, fine-tuned LLM rewrites to remove them. Runs entirely in-browser. fantasma.qstorage.quilibrium.com
wasm privacy qwen zig
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140const std = @import("std");
/// Temperature + top-k sampling for autoregressive generation.pub const Sampler = struct { temperature: f32, top_k: u32, xoshiro: std.Random.Xoshiro256,
pub fn init(temperature: f32, top_k: u32, seed: u64) Sampler { return .{ .temperature = temperature, .top_k = top_k, .xoshiro = std.Random.Xoshiro256.init(seed), }; }
/// Sample a token index from logits. /// Applies temperature scaling, then top-k filtering, then categorical sampling. pub fn sample(self: *Sampler, logits: []f32) u32 { const n = logits.len; if (n == 0) return 0;
// Temperature scaling if (self.temperature != 1.0 and self.temperature > 0.0) { const inv_temp = 1.0 / self.temperature; for (logits) |*l| { l.* *= inv_temp; } }
// Greedy (temperature ~0) if (self.temperature <= 0.0 or self.top_k == 1) { return argmax(logits); }
// Top-k: find the k-th largest value and zero out everything below const k = @min(self.top_k, @as(u32, @intCast(n)));
// Find k-th largest using partial sort via repeated max finding // For small k this is efficient enough var threshold: f32 = -std.math.inf(f32); if (k < n) { // Use a simple approach: find the k-th largest value // by maintaining the minimum of the top-k seen so far var top_k_min: f32 = std.math.inf(f32); var top_k_vals: [256]f32 = undefined; const actual_k: usize = @min(k, 256); var count: usize = 0;
for (logits) |l| { if (count < actual_k) { top_k_vals[count] = l; count += 1; if (count == actual_k) { // Find min of top_k_vals top_k_min = top_k_vals[0]; for (top_k_vals[1..actual_k]) |v| { if (v < top_k_min) top_k_min = v; } } } else if (l > top_k_min) { // Replace the min for (&top_k_vals, 0..) |*v, i| { _ = i; if (v.* == top_k_min) { v.* = l; break; } } top_k_min = top_k_vals[0]; for (top_k_vals[1..actual_k]) |v| { if (v < top_k_min) top_k_min = v; } } } threshold = top_k_min; }
// Softmax with top-k filtering var max_val: f32 = -std.math.inf(f32); for (logits) |l| { if (l >= threshold and l > max_val) max_val = l; }
var sum: f32 = 0.0; for (logits) |*l| { if (l.* >= threshold) { l.* = @exp(l.* - max_val); sum += l.*; } else { l.* = 0.0; } }
if (sum <= 0.0) return argmax(logits);
// Normalize const inv_sum = 1.0 / sum; for (logits) |*l| { l.* *= inv_sum; }
// Categorical sampling var r = self.xoshiro.random().float(f32); for (logits, 0..) |p, i| { r -= p; if (r <= 0.0) return @intCast(i); }
return @intCast(n - 1); }};
fn argmax(vals: []const f32) u32 { var best: u32 = 0; var best_val: f32 = -std.math.inf(f32); for (vals, 0..) |v, i| { if (v > best_val) { best_val = v; best = @intCast(i); } } return best;}
test "sampler greedy" { var logits = [_]f32{ 1.0, 3.0, 2.0, 0.5 }; var s = Sampler.init(0.0, 1, 42); const token = s.sample(&logits); try std.testing.expectEqual(@as(u32, 1), token);}
test "sampler top_k" { var logits = [_]f32{ 1.0, 10.0, 9.0, 0.5 }; var s = Sampler.init(1.0, 2, 42); const token = s.sample(&logits); // Should be either 1 or 2 (top-2) try std.testing.expect(token == 1 or token == 2);}