Spaces:

ChandraPrakashBathula
/

Speculative-Decoding

Paused

Speculative-Decoding / frontend /src /api.ts

Add live SSE streaming of generation and speculation rounds; remove em dashes

7055229 about 2 months ago

5.92 kB

export interface Round {
  proposed: number;
  accepted: number;
}

export interface BaselineResult {
  mode: "baseline";
  text: string;
  total_time_ms: number;
  target_calls: number;
  tokens_generated: number;
  rounds: [];
}

export interface SpeculativeResult {
  mode: "speculative";
  text: string;
  total_time_ms: number;
  target_calls: number;
  tokens_generated: number;
  rounds: Round[];
  round_texts: string[];
  gamma: number;
  /** empirical per-token acceptance rate, accepted/proposed over all rounds */
  alpha_hat: number | null;
  /** measured (draft cost per token) / (target cost per forward pass) */
  c_hat: number | null;
  /** the 2023 speedup formula evaluated at this run's own alpha_hat and c_hat */
  predicted_speedup: number | null;
  draft_time_ms: number;
  target_time_ms: number;
}

export interface HealthResponse {
  status: "ok" | "loading";
  device: string;
  draft_model: string;
  target_model: string;
}

export async function getHealth(): Promise<HealthResponse> {
  const res = await fetch("/health");
  if (!res.ok) throw new Error(`/health returned ${res.status}`);
  return res.json();
}

async function generate<T>(body: Record<string, unknown>): Promise<T> {
  const res = await fetch("/generate", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(body),
  });
  if (!res.ok) {
    throw new Error(`/generate failed: ${res.status} ${await res.text()}`);
  }
  return res.json();
}

export function runBaseline(prompt: string, maxNewTokens: number) {
  return generate<BaselineResult>({
    prompt,
    K: 1,
    mode: "baseline",
    max_new_tokens: maxNewTokens,
  });
}

export function runSpeculative(prompt: string, K: number, maxNewTokens: number) {
  return generate<SpeculativeResult>({
    prompt,
    K,
    mode: "speculative",
    max_new_tokens: maxNewTokens,
  });
}

/* ---------- live streaming ---------- */

export type StreamEvent =
  | { type: "start"; mode: string; K: number; max_new_tokens: number }
  | { type: "tokens"; text: string; count: number }
  | { type: "round"; proposed: number; accepted: number }
  | { type: "error"; message: string }
  | ({ type: "done" } & SpeculativeResult);

/**
 * Streams a generation over server-sent events.
 *
 * Uses fetch + ReadableStream rather than EventSource because the request is a
 * POST with a JSON body, which EventSource cannot send.
 */
export async function streamGenerate(
  body: { prompt: string; K: number; mode: "baseline" | "speculative"; max_new_tokens: number },
  onEvent: (event: StreamEvent) => void,
  signal?: AbortSignal,
): Promise<void> {
  const res = await fetch("/generate/stream", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(body),
    signal,
  });
  if (!res.ok || !res.body) {
    throw new Error(`/generate/stream failed: ${res.status} ${await res.text()}`);
  }

  const reader = res.body.getReader();
  const decoder = new TextDecoder();
  let buffer = "";

  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    buffer += decoder.decode(value, { stream: true });

    // SSE frames are separated by a blank line
    const frames = buffer.split("\n\n");
    buffer = frames.pop() ?? "";
    for (const frame of frames) {
      const line = frame.split("\n").find((l) => l.startsWith("data: "));
      if (!line) continue;
      try {
        onEvent(JSON.parse(line.slice(6)) as StreamEvent);
      } catch {
        // a partial or malformed frame is skipped rather than killing the stream
      }
    }
  }
}

/* ---------- derived metrics computed from measured rounds ---------- */

/**
 * Per-position acceptance rate within a speculation round.
 *
 * A round logged as {proposed: 4, accepted: 3} means draft positions 1..3 were
 * accepted and position 4 was rejected; positions beyond the rejection are
 * never evaluated, so they are excluded from that position's denominator
 * rather than counted as rejections. This is what makes the curve a real
 * conditional acceptance rate P(accept position i | reached position i), and
 * it is the direct empirical test of the i.i.d. alpha assumption the 2023
 * speedup formula makes.
 */
export function perPositionAcceptance(rounds: Round[], gamma: number) {
  const reached = new Array(gamma).fill(0);
  const accepted = new Array(gamma).fill(0);

  for (const round of rounds) {
    for (let i = 0; i < Math.min(round.proposed, gamma); i++) {
      reached[i] += 1;
      if (i < round.accepted) accepted[i] += 1;
      else break; // positions after the first rejection were never evaluated
    }
  }

  return reached.map((n, i) => ({
    position: i + 1,
    reached: n,
    accepted: accepted[i],
    rate: n > 0 ? accepted[i] / n : null,
  }));
}

/**
 * Distribution of "how many drafted tokens were accepted" per round, compared
 * against the geometric distribution the theory assumes with the same alpha.
 */
export function acceptedLengthDistribution(rounds: Round[], gamma: number, alpha: number) {
  const scored = rounds.filter((r) => r.proposed > 0);
  const counts = new Array(gamma + 1).fill(0);
  for (const round of scored) counts[Math.min(round.accepted, gamma)] += 1;

  const total = scored.length || 1;
  return counts.map((count, n) => {
    // P(exactly n accepted) = alpha^n * (1-alpha) for n < gamma; alpha^gamma for n = gamma
    const theoretical = n === gamma ? Math.pow(alpha, gamma) : Math.pow(alpha, n) * (1 - alpha);
    return {
      accepted: n,
      empirical: count / total,
      theoretical,
      count,
    };
  });
}

/** The 2023 closed form, used to draw the theoretical curve across K. */
export function theoreticalSpeedup(alpha: number, gamma: number, c: number): number {
  if (alpha > 0.999999) return (gamma + 1) / (gamma * c + 1);
  return (1 - Math.pow(alpha, gamma + 1)) / ((1 - alpha) * (gamma * c + 1));
}

	export interface Round {
	proposed: number;
	accepted: number;
	}

	export interface BaselineResult {
	mode: "baseline";
	text: string;
	total_time_ms: number;
	target_calls: number;
	tokens_generated: number;
	rounds: [];
	}

	export interface SpeculativeResult {
	mode: "speculative";
	text: string;
	total_time_ms: number;
	target_calls: number;
	tokens_generated: number;
	rounds: Round[];
	round_texts: string[];
	gamma: number;
	/** empirical per-token acceptance rate, accepted/proposed over all rounds */
	alpha_hat: number \| null;
	/** measured (draft cost per token) / (target cost per forward pass) */
	c_hat: number \| null;
	/** the 2023 speedup formula evaluated at this run's own alpha_hat and c_hat */
	predicted_speedup: number \| null;
	draft_time_ms: number;
	target_time_ms: number;
	}

	export interface HealthResponse {
	status: "ok" \| "loading";
	device: string;
	draft_model: string;
	target_model: string;
	}

	export async function getHealth(): Promise<HealthResponse> {
	const res = await fetch("/health");
	if (!res.ok) throw new Error(`/health returned ${res.status}`);
	return res.json();
	}

	async function generate<T>(body: Record<string, unknown>): Promise<T> {
	const res = await fetch("/generate", {
	method: "POST",
	headers: { "Content-Type": "application/json" },
	body: JSON.stringify(body),
	});
	if (!res.ok) {
	throw new Error(`/generate failed: ${res.status} ${await res.text()}`);
	}
	return res.json();
	}

	export function runBaseline(prompt: string, maxNewTokens: number) {
	return generate<BaselineResult>({
	prompt,
	K: 1,
	mode: "baseline",
	max_new_tokens: maxNewTokens,
	});
	}

	export function runSpeculative(prompt: string, K: number, maxNewTokens: number) {
	return generate<SpeculativeResult>({
	prompt,
	K,
	mode: "speculative",
	max_new_tokens: maxNewTokens,
	});
	}

	/* ---------- live streaming ---------- */

	export type StreamEvent =
	\| { type: "start"; mode: string; K: number; max_new_tokens: number }
	\| { type: "tokens"; text: string; count: number }
	\| { type: "round"; proposed: number; accepted: number }
	\| { type: "error"; message: string }
	\| ({ type: "done" } & SpeculativeResult);

	/**
	* Streams a generation over server-sent events.
	*
	* Uses fetch + ReadableStream rather than EventSource because the request is a
	* POST with a JSON body, which EventSource cannot send.
	*/
	export async function streamGenerate(
	body: { prompt: string; K: number; mode: "baseline" \| "speculative"; max_new_tokens: number },
	onEvent: (event: StreamEvent) => void,
	signal?: AbortSignal,
	): Promise<void> {
	const res = await fetch("/generate/stream", {
	method: "POST",
	headers: { "Content-Type": "application/json" },
	body: JSON.stringify(body),
	signal,
	});
	if (!res.ok \|\| !res.body) {
	throw new Error(`/generate/stream failed: ${res.status} ${await res.text()}`);
	}

	const reader = res.body.getReader();
	const decoder = new TextDecoder();
	let buffer = "";

	for (;;) {
	const { done, value } = await reader.read();
	if (done) break;
	buffer += decoder.decode(value, { stream: true });

	// SSE frames are separated by a blank line
	const frames = buffer.split("\n\n");
	buffer = frames.pop() ?? "";
	for (const frame of frames) {
	const line = frame.split("\n").find((l) => l.startsWith("data: "));
	if (!line) continue;
	try {
	onEvent(JSON.parse(line.slice(6)) as StreamEvent);
	} catch {
	// a partial or malformed frame is skipped rather than killing the stream
	}
	}
	}
	}

	/* ---------- derived metrics computed from measured rounds ---------- */

	/**
	* Per-position acceptance rate within a speculation round.
	*
	* A round logged as {proposed: 4, accepted: 3} means draft positions 1..3 were
	* accepted and position 4 was rejected; positions beyond the rejection are
	* never evaluated, so they are excluded from that position's denominator
	* rather than counted as rejections. This is what makes the curve a real
	* conditional acceptance rate P(accept position i \| reached position i), and
	* it is the direct empirical test of the i.i.d. alpha assumption the 2023
	* speedup formula makes.
	*/
	export function perPositionAcceptance(rounds: Round[], gamma: number) {
	const reached = new Array(gamma).fill(0);
	const accepted = new Array(gamma).fill(0);

	for (const round of rounds) {
	for (let i = 0; i < Math.min(round.proposed, gamma); i++) {
	reached[i] += 1;
	if (i < round.accepted) accepted[i] += 1;
	else break; // positions after the first rejection were never evaluated
	}
	}

	return reached.map((n, i) => ({
	position: i + 1,
	reached: n,
	accepted: accepted[i],
	rate: n > 0 ? accepted[i] / n : null,
	}));
	}

	/**
	* Distribution of "how many drafted tokens were accepted" per round, compared
	* against the geometric distribution the theory assumes with the same alpha.
	*/
	export function acceptedLengthDistribution(rounds: Round[], gamma: number, alpha: number) {
	const scored = rounds.filter((r) => r.proposed > 0);
	const counts = new Array(gamma + 1).fill(0);
	for (const round of scored) counts[Math.min(round.accepted, gamma)] += 1;

	const total = scored.length \|\| 1;
	return counts.map((count, n) => {
	// P(exactly n accepted) = alpha^n * (1-alpha) for n < gamma; alpha^gamma for n = gamma
	const theoretical = n === gamma ? Math.pow(alpha, gamma) : Math.pow(alpha, n) * (1 - alpha);
	return {
	accepted: n,
	empirical: count / total,
	theoretical,
	count,
	};
	});
	}

	/** The 2023 closed form, used to draw the theoretical curve across K. */
	export function theoreticalSpeedup(alpha: number, gamma: number, c: number): number {
	if (alpha > 0.999999) return (gamma + 1) / (gamma * c + 1);
	return (1 - Math.pow(alpha, gamma + 1)) / ((1 - alpha) * (gamma * c + 1));
	}