Space for all-prefix-optimal coupling

.gitignore

@@ -0,0 +1 @@
+__pycache__

README.md

@@ -1,12 +1,18 @@
 ---
-title: All Prefix Optimal Coupling Demo 1
-emoji: 📊
-colorFrom: yellow
-colorTo: red
+title: All-Prefix-Optimal Coupling
+emoji: 🔗
+colorFrom: gray
+colorTo: gray
 sdk: gradio
+python_version: 3.11
 sdk_version: 4.41.0
 app_file: app.py
+short_description: Tightly pair LLM responses
+models:
+  - Qwen/Qwen2-0.5B-Instruct
 pinned: false
+preload_from_hub:
+  - Qwen/Qwen2-0.5B-Instruct
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

algorithms.py

@@ -0,0 +1,542 @@
+# Device-independent algorithms for LLM.
+
+import logging
+import time
+import torch
+import torch.nn.functional as F
+import transformers
+logger = logging.getLogger(__name__)
+
+# Helper to pull out the response tokens.
+def isolate_responses_BL(output_BL, prompt_len, eos_token_id):
+    responses_BL = []
+    resp_BL = output_BL[:, prompt_len:]
+    for i in range(resp_BL.shape[0]):
+        resp_L = resp_BL[i]
+        resplen = resp_L.shape[0]
+        for j in range(resplen):
+            if resp_L[j] == eos_token_id:
+                resplen = j+1
+                break
+        response_L = resp_L[:resplen].cpu().detach().numpy()
+        responses_BL.append(response_L)
+    return responses_BL
+
+def tokenize_prompt(device, tokenizer, chat, quiet=False):
+    # Tokenize the prompt.
+    prompt_BL = tokenizer.apply_chat_template(
+        [chat],
+        tokenizer=True, add_generation_prompt=True, return_tensors='pt'
+    ).to(device)
+    if not quiet:
+        print('PROMPT:')
+        print(tokenizer.decode(prompt_BL[0]))
+    return prompt_BL    
+
+def generate(device, model, tokenizer, chat):
+    """Generate a response using huggingface's generation."""
+
+    prompt_BL = tokenize_prompt(device, tokenizer, chat)
+    prompt_len = prompt_BL.shape[1]
+
+    # Generate response.
+
+    # Unfortunately, huggingface's generation code uses 'cumsum',
+    # which doesn't have a deterministic implementation.
+    torch.use_deterministic_algorithms(False)
+    generation_output = model.generate(
+        inputs=prompt_BL, max_new_tokens=512, do_sample=True,
+        return_dict_in_generate=True,
+    )
+    torch.use_deterministic_algorithms(True)
+
+    output_BL = generation_output.sequences
+
+    for response_L in isolate_responses_BL(output_BL, prompt_len, tokenizer.eos_token_id):
+        print(f'RESPONSE:')
+        print(tokenizer.decode(response_L))
+
+    return tokenizer.decode(response_L)
+
+def generate_with_logits(device, model, tokenizer, chat, seed=None):
+    # Huggingface generation that returns logits too.
+
+    if seed is not None:
+        transformers.set_seed(seed)
+
+    prompt_BL = tokenize_prompt(device, tokenizer, chat)
+    prompt_len = prompt_BL.shape[1]
+
+    # Unfortunately, huggingface's generation code uses 'cumsum',
+    # which doesn't have a deterministic implementation.
+    torch.use_deterministic_algorithms(False)
+    generation_output = model.generate(
+        inputs=prompt_BL, max_new_tokens=16, do_sample=True,
+        return_dict_in_generate=True, output_logits=True,
+    )
+    torch.use_deterministic_algorithms(True)
+
+    output_BL = generation_output.sequences
+    logits_BLV = torch.stack(generation_output.logits, axis=1)
+
+    for response_L in isolate_responses_BL(output_BL, prompt_len, tokenizer.eos_token_id):
+        print(f'RESPONSE:')
+        print(tokenizer.decode(response_L))
+
+    return response_L, logits_BLV[0]
+
+def response_logits(device, model, tokenizer, chat, response_L):
+    # Calculate logits using a single pass.
+
+    prompt_BL = tokenize_prompt(device, tokenizer, chat)
+    response_pt_L = torch.from_numpy(response_L[:-1])
+
+    # Concatenate along axis 1.
+    input_ids = torch.cat((prompt_BL, response_pt_L[None]), dim=1)
+
+    outputs = model(input_ids)
+    logits_BLV = outputs.logits
+
+    return logits_BLV[0][-len(response_L):]
+
+def generate_custom(device, model, tokenizer, chat, max_tokens=512, seed=None, return_tokens=False, quiet=False, return_lnprobs=False):
+    """Generate a response using custom generation."""
+
+    if seed is not None:
+        torch.manual_seed(seed)
+
+    prompt_BL = tokenize_prompt(device, tokenizer, chat, quiet=quiet)
+
+    # Generate response.
+
+    input_ids = prompt_BL
+    past_key_values = None
+
+    response_L = []
+    lnprobs_L = []
+    while 1:
+        outputs = model(input_ids, past_key_values=past_key_values, use_cache=True)
+        past_key_values = outputs.past_key_values
+
+        logits = outputs.logits[0, -1, :]
+        p_V = F.softmax(logits, dim=-1)
+        token = torch.multinomial(p_V, num_samples=1).item()
+        response_L.append(token)
+        lnprobs_L.append(F.log_softmax(logits, dim=-1)[token].item())
+
+        p_token = p_V[token].item()
+        if not quiet:
+            print(f'  Sampled token {format_token(tokenizer, token)}   ({p_token*100.0:.3f}%)')
+
+        if token == tokenizer.eos_token_id or len(response_L) >= max_tokens:
+            break
+        
+        input_ids = torch.tensor([[token]], device=device)
+    
+    if not quiet:
+        print(f'RESPONSE:')
+        print(tokenizer.decode(response_L))
+
+    if return_tokens:
+        if return_lnprobs:
+            return response_L, lnprobs_L
+        else:
+            return response_L
+    else:
+        assert not return_lnprobs
+        return tokenizer.decode(response_L)
+
+def format_token(tokenizer, token_id):
+    return repr(tokenizer.decode(token_id))
+
+def apoc(device, model_x, model_y, tokenizer, chat_x, chat_y, max_tokens=512, seed=None, return_tokens=False, quiet=False):
+    """Generate a response using APOC unconditional sampling."""
+
+    if seed is not None:
+        torch.manual_seed(seed)
+
+    # This early implementation of the algorithm is numerically non-robust,
+    # so reduce problems by using high-precision floating-point.
+    logit_dtype = torch.float64
+
+    prompt_x_BL = tokenize_prompt(device, tokenizer, chat_x, quiet=quiet)
+    prompt_y_BL = tokenize_prompt(device, tokenizer, chat_y, quiet=quiet)
+
+    # Four variables are needed, since in the first iteration it depends on prompt (X vs Y),
+    # whereas in later iterations it depends on response (a vs b).
+    input_ids_xa = prompt_x_BL
+    input_ids_ya = prompt_y_BL
+    input_ids_xb = prompt_x_BL
+    input_ids_yb = prompt_y_BL
+
+    past_key_values_xa = None
+    past_key_values_ya = None
+    past_key_values_xb = None
+    past_key_values_yb = None
+
+    equal = True
+    a_eos = False
+    b_eos = False
+    def zero():
+        return torch.zeros(1, dtype=logit_dtype, device=device)
+    ln_pya_m_ln_pxa = zero()
+    ln_pxb_m_ln_pyb = zero()
+
+    response_a_L = []
+    response_b_L = []
+    i = 0
+    while 1:
+        if i >= max_tokens or (a_eos and b_eos): break
+        if not quiet:
+            print(f'Generating response token {i}')
+        i += 1
+
+        forward_passes_start = time.perf_counter()
+
+        if not a_eos:
+            outputs = model_x(input_ids_xa, past_key_values=past_key_values_xa, use_cache=True)
+            past_key_values_xa = outputs.past_key_values
+            logits = outputs.logits[0, -1, :].to(logit_dtype)
+            ln_pxa_V = F.log_softmax(logits, dim=-1)
+
+            outputs = model_y(input_ids_ya, past_key_values=past_key_values_ya, use_cache=True)
+            past_key_values_ya = outputs.past_key_values
+            logits = outputs.logits[0, -1, :].to(logit_dtype)
+            ln_pya_V = F.log_softmax(logits, dim=-1)
+        
+        if not b_eos:
+            if equal:
+                # In equal mode, neither input_ids nor past_key_values depends on a vs b,
+                # so we can reuse the forward pass results for a 50% time savings.
+                assert not a_eos
+                past_key_values_xb = past_key_values_xa
+                ln_pxb_V = ln_pxa_V
+                past_key_values_yb = past_key_values_ya
+                ln_pyb_V = ln_pya_V
+            else:
+                outputs = model_x(input_ids_xb, past_key_values=past_key_values_xb, use_cache=True)
+                past_key_values_xb = outputs.past_key_values
+                logits = outputs.logits[0, -1, :].to(logit_dtype)
+                ln_pxb_V = F.log_softmax(logits, dim=-1)
+
+                outputs = model_y(input_ids_yb, past_key_values=past_key_values_yb, use_cache=True)
+                past_key_values_yb = outputs.past_key_values
+                logits = outputs.logits[0, -1, :].to(logit_dtype)
+                ln_pyb_V = F.log_softmax(logits, dim=-1)
+        
+        forward_passes_end = time.perf_counter()
+        if not quiet:
+            print(f'  Forward passes took {(forward_passes_end - forward_passes_start)*1000:.0f} ms')
+        
+        if equal:
+            ln_pmeet_V = torch.minimum(
+                ln_pxa_V + torch.maximum(zero(), -ln_pya_m_ln_pxa),
+                ln_pya_V + torch.maximum(zero(), ln_pya_m_ln_pxa),
+            )
+            pmeet_V = torch.exp(ln_pmeet_V)
+            pmeet = torch.sum(pmeet_V)
+
+            if not quiet:
+                print(f'  Equal mode (pmeet={pmeet*100.0:.3f}%)')
+
+            if torch.rand_like(pmeet) < pmeet:
+                token_a = token_b = torch.multinomial(pmeet_V, num_samples=1).item()
+
+                p_token_a = (pmeet_V[token_a] / pmeet).item()
+                if not quiet:
+                    print(f'  Sampled {format_token(tokenizer, token_a)}   ({p_token_a*100.0:.1f}%)')
+            else:
+                if not quiet:
+                    print('  Exited equal mode')
+                equal = False
+
+        if not equal:
+            if not a_eos:
+                wxt_V = torch.maximum(zero(), torch.exp(ln_pxa_V) - torch.exp(ln_pya_V + ln_pya_m_ln_pxa))
+                token_a = torch.multinomial(wxt_V, num_samples=1).item()
+
+                p_token_a = (wxt_V[token_a] / torch.sum(wxt_V)).item()
+                if not quiet:
+                    print(f'  Sampled token_a {format_token(tokenizer, token_a)}   ({p_token_a*100.0:.3f}%)')
+            if not b_eos:
+                wyt_V = torch.maximum(zero(), torch.exp(ln_pyb_V) - torch.exp(ln_pxb_V + ln_pxb_m_ln_pyb))
+                token_b = torch.multinomial(wyt_V, num_samples=1).item()
+
+                p_token_b = (wyt_V[token_b] / torch.sum(wyt_V)).item()
+                if not quiet:
+                    print(f'  Sampled token_b {format_token(tokenizer, token_b)}   ({p_token_b*100.0:.3f}%)')
+        
+        if not a_eos:
+            response_a_L.append(token_a)
+            input_ids_xa = input_ids_ya = torch.tensor([[token_a]], device=device)
+            ln_pya_m_ln_pxa += ln_pya_V[token_a] - ln_pxa_V[token_a]
+            if token_a == tokenizer.eos_token_id:
+                a_eos = True
+        if not b_eos:
+            response_b_L.append(token_b)
+            input_ids_xb = input_ids_yb = torch.tensor([[token_b]], device=device)
+            ln_pxb_m_ln_pyb += ln_pxb_V[token_b] - ln_pyb_V[token_b]
+            if token_b == tokenizer.eos_token_id:
+                b_eos = True
+
+    if not quiet:
+        print(f'RESPONSE X:')
+        print(tokenizer.decode(response_a_L))
+        print(f'RESPONSE Y:')
+        print(tokenizer.decode(response_b_L))
+
+    if return_tokens:
+        return response_a_L, response_b_L
+    else:
+        return tokenizer.decode(response_a_L), tokenizer.decode(response_b_L)
+
+# Alternative implementation.
+@torch.no_grad()
+def apoc_alt(device, model_x, model_y, tokenizer, chat_x, chat_y, max_tokens=512, seed=None):
+    if seed is not None:
+        torch.manual_seed(seed)
+
+    prompt_x_BL = tokenize_prompt(device, tokenizer, chat_x, quiet=True)
+    prompt_y_BL = tokenize_prompt(device, tokenizer, chat_y, quiet=True)
+    model_pair = ModelPair(model_x, model_y, prompt_x_BL, prompt_y_BL)
+
+    logger.debug('PROMPT X:')
+    logger.debug(tokenizer.decode(prompt_x_BL[0]))
+    logger.debug('PROMPT Y:')
+    logger.debug(tokenizer.decode(prompt_y_BL[0]))
+
+    return _apoc_impl(model_pair, tokenizer, max_tokens)
+
+LOGIT_DTYPE = torch.float64
+
+class ModelPair:
+    def __init__(self, model_x, model_y, prompt_x_BL, prompt_y_BL):
+        self._model_x = model_x
+        self._model_y = model_y
+        self._prompt_x_BL = prompt_x_BL
+        self._prompt_y_BL = prompt_y_BL
+        self._is_swapped = False
+
+    def start(self):
+        # Return logprobs for the initial token.
+
+        outputs = self._model_x(self._prompt_x_BL, use_cache=True)
+        self._past_key_values_x = outputs.past_key_values
+        logits = outputs.logits[0, -1, :].to(LOGIT_DTYPE)
+        lnpx_V = F.log_softmax(logits, dim=-1)
+
+        outputs = self._model_y(self._prompt_y_BL, use_cache=True)
+        self._past_key_values_y = outputs.past_key_values
+        logits = outputs.logits[0, -1, :].to(LOGIT_DTYPE)
+        lnpy_V = F.log_softmax(logits, dim=-1)
+
+        return self._maybe_swap(lnpx_V, lnpy_V)
+
+    def step(self, token):
+        # Append the given token, then return logprobs for the next token.
+
+        forward_passes_start = time.perf_counter()
+
+        input_ids = torch.tensor([[token]], device=self._prompt_x_BL.device)
+
+        outputs = self._model_x(input_ids, past_key_values=self._past_key_values_x, use_cache=True)
+        self._past_key_values_x = outputs.past_key_values
+        logits = outputs.logits[0, -1, :].to(LOGIT_DTYPE)
+        lnpx_V = F.log_softmax(logits, dim=-1)
+
+        outputs = self._model_y(input_ids, past_key_values=self._past_key_values_y, use_cache=True)
+        self._past_key_values_y = outputs.past_key_values
+        logits = outputs.logits[0, -1, :].to(LOGIT_DTYPE)
+        lnpy_V = F.log_softmax(logits, dim=-1)
+
+        forward_passes_end = time.perf_counter()
+        logger.debug(f'Incremental forward passes took {(forward_passes_end - forward_passes_start)*1000:.0f} ms')
+
+        return self._maybe_swap(lnpx_V, lnpy_V)
+    
+    def get_position(self):
+        # Return a position that can be rewound to.
+        return self._past_key_values_x, self._past_key_values_y
+    
+    def rewind_to(self, position):
+        # Rewind the KV cache.
+        self._past_key_values_x, self._past_key_values_y = position
+    
+    def swap_models(self):
+        # Exchange the order of the models.
+        self._is_swapped = not self._is_swapped
+    
+    def _maybe_swap(self, a, b):
+        if self._is_swapped:
+            return b, a
+        else:
+            return a, b
+
+def _apoc_impl(model_pair, tokenizer, max_tokens):
+    prefix = []
+    lnpx_V, lnpy_V = model_pair.start()
+    lnpy_m_lnpx = torch.zeros(1, dtype=lnpx_V.dtype, device=lnpx_V.device)
+
+    while 1:
+        ln_peq_V = torch.minimum(
+            lnpx_V + F.relu(-lnpy_m_lnpx),
+            lnpy_V + F.relu(lnpy_m_lnpx),
+        )
+        peq_V = torch.exp(ln_peq_V)
+        peq = torch.sum(peq_V)
+
+        if torch.rand_like(peq) > peq:
+            logger.debug(f'Completed common prefix ({(1-peq)*100.:.3f}%)')
+            break
+        logger.debug(f'Extending common prefix ({peq*100.:.3f}%)')
+        
+        token = torch.multinomial(peq_V, 1).item()
+        prefix.append(token)
+        lnpy_m_lnpx += lnpy_V[token] - lnpx_V[token]
+
+        p_token = (peq_V[token] / peq).item()
+        logger.debug(f'Sampled prefix token {format_token(tokenizer, token)}   ({p_token*100.0:.3f}%)')
+
+        if token == tokenizer.eos_token_id or len(prefix) >= max_tokens:
+            return prefix, prefix
+        
+        lnpx_V, lnpy_V = model_pair.step(token)
+    
+    remaining_tokens = max_tokens - len(prefix)
+    split_pos = model_pair.get_position()
+    response_a = prefix + _apoc_gen_suffix(model_pair, tokenizer, remaining_tokens, lnpx_V, lnpy_V, lnpy_m_lnpx)
+    logger.debug('First suffix complete; rewinding')
+    model_pair.rewind_to(split_pos)
+    model_pair.swap_models()
+    response_b = prefix + _apoc_gen_suffix(model_pair, tokenizer, remaining_tokens, lnpy_V, lnpx_V, -lnpy_m_lnpx)
+
+    return response_a, response_b
+
+def _apoc_gen_suffix(model_pair, tokenizer, max_tokens, lnpx_V, lnpy_V, lnpy_m_lnpx):
+    lnpy_m_lnpx = lnpy_m_lnpx.clone()
+    suffix = []
+    while 1:
+        wx_V = F.relu(torch.exp(lnpx_V) - torch.exp(lnpy_V + lnpy_m_lnpx))
+        token = torch.multinomial(wx_V, 1).item()
+        suffix.append(token)
+        lnpy_m_lnpx += lnpy_V[token] - lnpx_V[token]
+
+        p_token = (wx_V[token] / torch.sum(wx_V)).item()
+        logger.debug(f'Sampled suffix token {format_token(tokenizer, token)}   ({p_token*100.0:.3f}%)')
+
+        if token == tokenizer.eos_token_id or len(suffix) >= max_tokens:
+            return suffix
+        
+        lnpx_V, lnpy_V = model_pair.step(token)
+
+def generate_streaming(device, model, tokenizer, chat, max_tokens=512, seed=None):
+    """Stream a response using custom generation."""
+
+    prompt_BL = tokenize_prompt(device, tokenizer, chat, quiet=True)
+    logger.debug('PROMPT:')
+    logger.debug(tokenizer.decode(prompt_BL[0]))
+
+    if seed is not None:
+        torch.manual_seed(seed)
+
+    return _generate_streaming_impl(device, model, tokenizer, prompt_BL, max_tokens)
+
+def _generate_streaming_impl(device, model, tokenizer, prompt_BL, max_tokens):
+    input_ids = prompt_BL
+    past_key_values = None
+
+    n_tokens = 0
+    while 1:
+        outputs = model(input_ids, past_key_values=past_key_values, use_cache=True)
+        past_key_values = outputs.past_key_values
+
+        logits = outputs.logits[0, -1, :]
+        p_V = F.softmax(logits, dim=-1)
+        token = torch.multinomial(p_V, num_samples=1).item()
+
+        p_token = p_V[token].item()
+        logger.debug(f'  Sampled token {format_token(tokenizer, token)}   ({p_token*100.0:.3f}%)')
+
+        yield token
+        n_tokens += 1
+
+        if token == tokenizer.eos_token_id or n_tokens >= max_tokens:
+            break
+        
+        input_ids = torch.tensor([[token]], device=device)
+
+# APOC unconditional streaming
+@torch.no_grad()
+def apoc_streaming(device, model_x, model_y, tokenizer, chat_x, chat_y, max_tokens=512, seed=None):
+    if seed is not None:
+        torch.manual_seed(seed)
+
+    prompt_x_BL = tokenize_prompt(device, tokenizer, chat_x, quiet=True)
+    prompt_y_BL = tokenize_prompt(device, tokenizer, chat_y, quiet=True)
+    model_pair = ModelPair(model_x, model_y, prompt_x_BL, prompt_y_BL)
+
+    logger.debug('PROMPT X:')
+    logger.debug(tokenizer.decode(prompt_x_BL[0]))
+    logger.debug('PROMPT Y:')
+    logger.debug(tokenizer.decode(prompt_y_BL[0]))
+
+    return _apoc_streaming_impl(model_pair, tokenizer, max_tokens)
+
+def _apoc_streaming_impl(model_pair, tokenizer, max_tokens):
+    remaining_tokens = max_tokens
+    lnpx_V, lnpy_V = model_pair.start()
+    lnpy_m_lnpx = torch.zeros(1, dtype=lnpx_V.dtype, device=lnpx_V.device)
+
+    while 1:
+        ln_peq_V = torch.minimum(
+            lnpx_V + F.relu(-lnpy_m_lnpx),
+            lnpy_V + F.relu(lnpy_m_lnpx),
+        )
+        peq_V = torch.exp(ln_peq_V)
+        peq = torch.sum(peq_V)
+
+        if torch.rand_like(peq) > peq:
+            logger.debug(f'Completed common prefix ({(1-peq)*100.:.3f}%)')
+            break
+        logger.debug(f'Extending common prefix ({peq*100.:.3f}%)')
+        
+        token = torch.multinomial(peq_V, 1).item()
+        remaining_tokens -= 1
+        yield token, token
+        lnpy_m_lnpx += lnpy_V[token] - lnpx_V[token]
+
+        p_token = (peq_V[token] / peq).item()
+        logger.debug(f'Sampled prefix token {format_token(tokenizer, token)}   ({p_token*100.0:.3f}%)')
+
+        if token == tokenizer.eos_token_id or remaining_tokens == 0:
+            return
+        
+        lnpx_V, lnpy_V = model_pair.step(token)
+    
+    split_pos = model_pair.get_position()
+    for token_a in _apoc_streaming_gen_suffix(model_pair, tokenizer, remaining_tokens, lnpx_V, lnpy_V, lnpy_m_lnpx):
+        yield token_a, None
+    logger.debug('Suffix a complete; rewinding')
+    model_pair.rewind_to(split_pos)
+    model_pair.swap_models()
+    for token_b in _apoc_streaming_gen_suffix(model_pair, tokenizer, remaining_tokens, lnpy_V, lnpx_V, -lnpy_m_lnpx):
+        yield None, token_b
+    logger.debug('Suffix b complete')
+
+def _apoc_streaming_gen_suffix(model_pair, tokenizer, max_tokens, lnpx_V, lnpy_V, lnpy_m_lnpx):
+    remaining_tokens = max_tokens
+    lnpy_m_lnpx = lnpy_m_lnpx.clone()
+    while 1:
+        wx_V = F.relu(torch.exp(lnpx_V) - torch.exp(lnpy_V + lnpy_m_lnpx))
+        token = torch.multinomial(wx_V, 1).item()
+        remaining_tokens -= 1
+        yield token
+        lnpy_m_lnpx += lnpy_V[token] - lnpx_V[token]
+
+        p_token = (wx_V[token] / torch.sum(wx_V)).item()
+        logger.debug(f'Sampled suffix token {format_token(tokenizer, token)}   ({p_token*100.0:.3f}%)')
+
+        if token == tokenizer.eos_token_id or remaining_tokens == 0:
+            return
+        
+        lnpx_V, lnpy_V = model_pair.step(token)

app.py

@@ -0,0 +1,159 @@
+# Gradio demo of streaming generation of multiple LLM response pairs.
+
+import logging
+import time
+import html
+import numpy as np
+import gradio as gr
+import util
+
+# gr.DataFrame is currently bugged for updating values,
+# so we must use raw HTML.
+# https://github.com/gradio-app/gradio/issues/8160
+def make_html_table(headers, data):
+    rows = ['<tr>' + ''.join(f'<th style="width: 50%">{h}</th>' for h in headers) + '</tr>\n']
+    for row in data:
+        rows.append('<tr>' + ''.join(f'<td style="width: 50%; font-family: monospace; white-space: pre-wrap;">{v}</td>' for v in row) + '</tr>\n')
+    return '<table style="width: 100%; table-layout: fixed">\n' + ''.join(rows) + '</table>\n'
+
+def highlight_prefix(tokens, prefix_len):
+    prefix_tokens = tokens[:prefix_len]
+
+    s = tokenizer.decode(tokens, skip_special_tokens=True)
+    prefix_s = tokenizer.decode(prefix_tokens, skip_special_tokens=True)
+
+    s_lcp_len = util.longest_common_prefix(np.array(list(s)), np.array(list(prefix_s)))
+
+    prefix_html = html.escape(s[:s_lcp_len])
+    suffix_html = html.escape(s[s_lcp_len:])
+
+    #highlight_style = 'background-color: #FFFFAE;'
+    #highlight_style = 'text-decoration: underline;'
+    highlight_style = 'background-color: #90FF90;'
+
+    return f'<span style="{highlight_style}">{prefix_html}</span>{suffix_html}'
+
+def format_response_pair(tokens_a, tokens_b):
+    # This is slightly convoluted, so as to properly handle grapheme clusters that span token boundaries.
+    token_lcp_len = util.longest_common_prefix(tokens_a, tokens_b)
+    return highlight_prefix(tokens_a, token_lcp_len), highlight_prefix(tokens_b, token_lcp_len)
+
+HEADERS = ['Response (Left)', 'Response (Right)']
+repo_id = "Qwen/Qwen2-0.5B-Instruct"
+
+DRY_RUN = False
+
+if DRY_RUN:
+    from load import load_tokenizer
+
+    tokenizer = load_tokenizer(repo_id)
+
+    def fn(max_tokens, num_responses, prompt_x, prompt_y):
+        rows = [['']*2 for i in range(num_responses)]
+        
+        yield make_html_table(HEADERS, rows)
+
+        for j in range(num_responses):
+            response_raw_a = f'Sure!\n\n1 2 3 4 & 5.'
+            response_raw_b = f'Sure!\n\n1 2 3 4 5 & 6.'
+
+            response_tok_a = tokenizer.encode(response_raw_a, add_special_tokens=False, return_tensors='np')[0]
+            response_tok_b = tokenizer.encode(response_raw_b, add_special_tokens=False, return_tensors='np')[0]
+
+            steps = 1 + max(len(response_tok_a), len(response_tok_b))
+
+            for i in range(steps):
+                time.sleep(0.1)
+                prefix_tok_a = response_tok_a[:i]
+                prefix_tok_b = response_tok_b[:i]
+
+                content_a, content_b = format_response_pair(prefix_tok_a, prefix_tok_b)
+
+                rows[j][0] = content_a
+                rows[j][1] = content_b
+
+                yield make_html_table(HEADERS, rows)
+else:
+    from load import load_model
+    import algorithms
+
+    logging.basicConfig(format='%(levelname)s:%(name)s: %(message)s')
+    algorithms.logger.setLevel(logging.INFO)
+
+    model, tokenizer = load_model(repo_id)
+
+    def make_chat(system_msg, prompt):
+        chat = [
+                {
+                    'role': 'system',
+                    'content': system_msg,
+                },
+                {
+                    'role': 'user',
+                    'content': prompt,
+                },
+        ]
+        return chat
+
+    def fn(max_tokens, num_responses, prompt_x, prompt_y):
+        rows = [['']*2 for i in range(num_responses)]
+        yield make_html_table(HEADERS, rows)
+
+        for j in range(num_responses):
+            system_msg = "You are a helpful assistant."
+
+            chat_x = make_chat(system_msg, prompt_x)
+            chat_y = make_chat(system_msg, prompt_y)
+
+            gen = algorithms.apoc_streaming(
+                'cpu',
+                model,
+                model,
+                tokenizer,
+                chat_x,
+                chat_y,
+                max_tokens=max_tokens,
+            )
+            response_a_L = []
+            response_b_L = []
+            for token_a, token_b in gen:
+                dirty = False
+                if token_a is not None:
+                    response_a_L.append(token_a)
+                    dirty = True
+                if token_b is not None:
+                    response_b_L.append(token_b)
+                    dirty = True
+                
+                if dirty:
+                    content_a, content_b = format_response_pair(np.array(response_a_L), np.array(response_b_L))
+
+                    rows[j][0] = content_a
+                    rows[j][1] = content_b
+                
+                yield make_html_table(HEADERS, rows)
+
+demo = gr.Interface(
+    fn=fn,
+    inputs=[
+        gr.Slider(1, 512, label='Max Tokens', value=48),
+        gr.Slider(1, 16, step=1, label='Num Responses', value=8),
+        gr.Textbox(label='Prompt (Left)'),
+        gr.Textbox(label='Prompt (Right)'),
+        ],
+    outputs=[
+        gr.HTML(),
+        ],
+    title='All-Prefix-Optimal Coupling',
+    description='Try similar prompts to see the effect of the difference between them. '
+        f'Model: `{repo_id}`.'
+        ,
+    examples=[
+        [48, 8, 'Count from 1 to 5.', 'Count from 1 to 6.'],
+        [48, 8, 'Tell me a joke.', 'Tell me a funny joke.'],
+        [48, 8, 'Calculate 3 + 4', 'Calculate 3 + 5'],
+        [48, 8, "What's the capital of Canada?", "What's the capital of France?"],
+    ],
+    )
+
+demo.launch()

load.py

@@ -0,0 +1,61 @@
+# Code to load a model.
+
+import os
+import warnings
+import torch
+import transformers
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+def load_model(repo_id, device_map=None, bnb=None, torch_dtype='auto'):
+    # Try our best to get deterministic results.
+    if device_map is not None:
+        # For determinism with CUDA >= 10.2, PyTorch says to use one of these.
+        os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
+        #os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':16:8'
+
+    torch.use_deterministic_algorithms(True)
+
+    # Ignore a spurious warning from huggingface_hub:
+    # https://github.com/huggingface/transformers/issues/30618
+    warnings.filterwarnings('ignore', message="`resume_download` is deprecated")
+
+    # Ignore a spurious warning from bitsandbytes.
+    warnings.filterwarnings('ignore', message="MatMul8bitLt: inputs will be cast from")
+
+    print(f'Loading model "{repo_id}" (bnb = "{bnb}")...')
+
+    # Ignore a spurious warning "Special tokens have been added..."
+    transformers.logging.set_verbosity_error()
+    tokenizer = AutoTokenizer.from_pretrained(repo_id, use_fast=True)
+    transformers.logging.set_verbosity_warning()
+
+    bnb_config = None
+    if bnb == 'nf8':
+        bnb_config = BitsAndBytesConfig(load_in_8bit=True)
+    if bnb == 'nf4':
+        bnb_config = BitsAndBytesConfig(load_in_4bit=True)
+
+    model = AutoModelForCausalLM.from_pretrained(
+        repo_id,
+        torch_dtype=torch_dtype,
+        device_map=device_map,
+        quantization_config=bnb_config,
+    )
+
+    # Disable gradients to save memory.
+    for param in model.parameters():
+        param.requires_grad = False
+
+    # Try our best to get deterministic results.
+    model.eval()
+
+    print('Done loading model.')
+
+    return model, tokenizer
+
+def load_tokenizer(repo_id):
+    # Ignore a spurious warning "Special tokens have been added..."
+    transformers.logging.set_verbosity_error()
+    tokenizer = AutoTokenizer.from_pretrained(repo_id, use_fast=True)
+    transformers.logging.set_verbosity_warning()
+    return tokenizer

requirements.txt

@@ -0,0 +1,4 @@
+huggingface_hub==0.22.2
+numpy==1.26.4
+torch==2.2.2
+transformers==4.40.2

util.py

@@ -0,0 +1,15 @@
+import numpy as np
+
+def longest_common_prefix(xs, ys):
+    min_len = min(len(xs), len(ys))
+    idxs = (xs[:min_len] != ys[:min_len]).nonzero()[0]
+    if len(idxs) > 0:
+        return idxs[0]
+    else:
+        return min_len
+    
+# Like np.cumsum, but with a leading zero.
+def cumsum0(x, axis):
+    pad_width = len(x.shape) * [(0,0)]
+    pad_width[axis] = (1,0)
+    return np.cumsum(np.pad(x, pad_width), axis=axis)