|
| 1 | +import time |
| 2 | + |
| 3 | +import torch |
| 4 | +from jaxtyping import Float, Int |
| 5 | +from torch import Tensor |
| 6 | + |
| 7 | + |
| 8 | +def uniform_sample( |
| 9 | + prob_x: Float[Tensor, "bs vocab"] | Float[Tensor, "bs L vocab"], |
| 10 | + n_samples: int = 128, |
| 11 | + cls_token_id: int | None = None, |
| 12 | + device: torch.device | None = None, |
| 13 | +) -> Int[Tensor, "bs n_samples"] | Int[Tensor, "bs n_samples L"]: |
| 14 | + """Uniform sampling over the vocabulary for virtual do-interventions. |
| 15 | +
|
| 16 | + Supports both single-step distributions and full trajectory distributions. |
| 17 | + Internal logic automatically detects dimensionality to return consistent shapes. |
| 18 | +
|
| 19 | + Args: |
| 20 | + prob_x: Next token probabilities over the vocabulary. |
| 21 | + Can be [batch_size, vocab] or [batch_size, seq_len, vocab]. |
| 22 | + n_samples: Number of samples (particles) to generate per batch element. |
| 23 | + cls_token_id: If provided, forces the first token of every sample to this ID. |
| 24 | + device: Target device for sampled tensors. |
| 25 | +
|
| 26 | + Returns: |
| 27 | + sampled_tokens: The discrete samples. [bs, n_samples] for 2D input |
| 28 | + or [bs, n_samples, L] for 3D. |
| 29 | + """ |
| 30 | + |
| 31 | + device = device or prob_x.device |
| 32 | + |
| 33 | + if prob_x.dim() == 2: |
| 34 | + # ---- Single-step intervention ---- |
| 35 | + bs, vocab = prob_x.shape |
| 36 | + |
| 37 | + sampled_tokens = torch.randint(low=0, high=vocab, size=(bs,), device=device) |
| 38 | + |
| 39 | + if cls_token_id is not None: |
| 40 | + sampled_tokens[:] = cls_token_id |
| 41 | + |
| 42 | + return sampled_tokens |
| 43 | + |
| 44 | + elif prob_x.dim() == 3: |
| 45 | + # ---- Trajectory intervention ---- |
| 46 | + bs, L, vocab = prob_x.shape |
| 47 | + n = n_samples |
| 48 | + |
| 49 | + sampled = torch.randint(low=0, high=vocab, size=(bs, n, L), device=device) |
| 50 | + |
| 51 | + # Force CLS token if needed |
| 52 | + if cls_token_id is not None: |
| 53 | + sampled[:, :, 0] = cls_token_id |
| 54 | + |
| 55 | + return sampled |
| 56 | + else: |
| 57 | + raise ValueError("prob_x must be 2D or 3D tensor") |
| 58 | + |
| 59 | + |
| 60 | +def multinomial_sample( |
| 61 | + prob_x: Float[Tensor, "bs vocab"] | Float[Tensor, "bs L vocab"], |
| 62 | + n_samples: int = 128, |
| 63 | + cls_token_id: int | None = None, |
| 64 | + **kwargs, |
| 65 | +): |
| 66 | + """ |
| 67 | + Multinomial sampling from prob_x. |
| 68 | +
|
| 69 | + Args: |
| 70 | + prob_x: Next token probabilities over the vocabulary. |
| 71 | + Can be [batch_size, vocab] or [batch_size, seq_len, vocab]. |
| 72 | + n_samples: Number of samples (particles) to generate per batch element. |
| 73 | + cls_token_id: If provided, forces the first token of every sample to this ID. |
| 74 | + device: Target device for sampled tensors. |
| 75 | +
|
| 76 | + Returns: |
| 77 | + sampled_tokens: The discrete samples. [bs, n_samples] |
| 78 | + for 2D input or [bs, n_samples, L] for 3D. |
| 79 | + """ |
| 80 | + |
| 81 | + if prob_x.dim() == 2: |
| 82 | + # ---- Single-step sampling ---- |
| 83 | + # prob_x: [bs, vocab] |
| 84 | + sampled_tokens = torch.multinomial(prob_x, 1) |
| 85 | + return sampled_tokens |
| 86 | + |
| 87 | + elif prob_x.dim() == 3: |
| 88 | + # ---- Trajectory sampling ---- |
| 89 | + bs, L, vocab = prob_x.shape |
| 90 | + n = n_samples |
| 91 | + |
| 92 | + # Expand for n samples |
| 93 | + probs = prob_x.unsqueeze(1).expand(bs, n, L, vocab) |
| 94 | + probs = probs.reshape(-1, vocab) # [(bs*n*L), vocab] |
| 95 | + |
| 96 | + # Sample |
| 97 | + sampled = torch.multinomial(probs, 1).squeeze(-1) |
| 98 | + sampled = sampled.view(bs, n, L) |
| 99 | + |
| 100 | + # Force CLS token if needed |
| 101 | + if cls_token_id is not None: |
| 102 | + sampled[:, :, 0] = cls_token_id |
| 103 | + |
| 104 | + return sampled |
| 105 | + |
| 106 | + else: |
| 107 | + raise ValueError("prob_x must be 2D or 3D tensor") |
| 108 | + |
| 109 | + |
| 110 | +def ancestral_sampling( |
| 111 | + model: any, |
| 112 | + encoded_input: dict[str, Tensor], |
| 113 | + value: int = 64, |
| 114 | + guidance: int = 2, |
| 115 | + context: int = 10, |
| 116 | + proposal=multinomial_sample, |
| 117 | + **kwargs, |
| 118 | +): |
| 119 | + """ |
| 120 | + Standard Ancestral Sampling using a proposal function |
| 121 | +
|
| 122 | + Args: |
| 123 | + model: The autoregressive model to sample from. |
| 124 | + encoded_input: A dictionary containing 'input_ids' and 'attention_mask' tensors. |
| 125 | + value: Number of samples (particles) to generate per batch element. |
| 126 | + guidance: Number of initial tokens to use as conditioning context. |
| 127 | + context: Total length of the generated sequence (including guidance). |
| 128 | + proposal: The sampling function to use for generating tokens (e.g., multinomial_sample). |
| 129 | + Returns: |
| 130 | + sampled_tokens: The generated token sequences. Shape [bs*value, context]. |
| 131 | + """ |
| 132 | + |
| 133 | + torch.cuda.synchronize() |
| 134 | + start_time = time.time() |
| 135 | + N = value |
| 136 | + |
| 137 | + with torch.no_grad(): |
| 138 | + # ---- Step 1: initialize ---- |
| 139 | + start_tokens = encoded_input["input_ids"][:, :guidance].to(model.device).clone() |
| 140 | + attn_mask = encoded_input["attention_mask"].to(model.device) |
| 141 | + |
| 142 | + # upsample (repeat N times) |
| 143 | + start_tokens = start_tokens.unsqueeze(1).repeat(1, N, 1).reshape(-1, guidance) |
| 144 | + attn_mask = attn_mask.unsqueeze(1).repeat(1, N, 1).reshape(-1, attn_mask.size(-1)) |
| 145 | + |
| 146 | + for i in range(0, context - guidance): |
| 147 | + output = model( |
| 148 | + input_ids=start_tokens, |
| 149 | + attention_mask=attn_mask[:, : guidance + i].to(model.device), |
| 150 | + ) |
| 151 | + # we take the last digit. Be careful if padded |
| 152 | + prob_x = torch.nn.functional.softmax(output["logits"][:, -1, :], dim=-1) |
| 153 | + random_token = proposal(prob_x) |
| 154 | + start_tokens = torch.cat([start_tokens, random_token], dim=-1) |
| 155 | + |
| 156 | + torch.cuda.synchronize() |
| 157 | + elapsed = time.time() - start_time |
| 158 | + print("Ancestral Sampling - Elapsed time: ", elapsed) |
| 159 | + return start_tokens |
0 commit comments