Skip to content

smc

SequenceModel

A single SMC particle: one candidate sequence and its weight.

clone copies the mutable state (context, weight, twist_amount, max_tokens, done) and shares the sampler, critic and configuration.

Parameters:

Name Type Description Default
unit_sampler TokenSampler

Draws one unit per step via sample.

 required
critic Potential

Twists and reweights the particle.

 None
max_tokens int

Token budget; EOS is forced at the boundary.

 float('inf')
twist_with_critic bool

Whether the critic twists at every step, rather than scoring once at termination.

 True
terminate_when callable

A context -> bool stop condition; when it fires, EOS closes the sequence in that same step.

 None
verbosity int

0 is silent, 1 prints the particle at each step.

 0
Source code in genlm/control/sampler/smc.py 
 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
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
class SequenceModel:
    """A single SMC particle: one candidate sequence and its weight.

    `clone` copies the mutable state (`context`, `weight`, `twist_amount`,
    `max_tokens`, `done`) and shares the sampler, critic and configuration.

    Args:
        unit_sampler (TokenSampler): Draws one unit per step via `sample`.
        critic (Potential, optional): Twists and reweights the particle.
        max_tokens (int): Token budget; EOS is forced at the boundary.
        twist_with_critic (bool): Whether the critic twists at every step, rather
            than scoring once at termination.
        terminate_when (callable, optional): A `context -> bool` stop condition;
            when it fires, EOS closes the sequence in that same step.
        verbosity (int): 0 is silent, 1 prints the particle at each step.
    """

    def __init__(
        self,
        unit_sampler,
        critic=None,
        max_tokens=float("inf"),
        twist_with_critic=True,
        terminate_when=None,
        verbosity=0,
    ):
        self.unit_sampler = unit_sampler
        self.critic = critic
        self.max_tokens = max_tokens
        self.twist_with_critic = twist_with_critic
        self.terminate_when = terminate_when
        self.verbosity = verbosity

        self.context = []
        self.weight = 0.0
        self.twist_amount = 0.0
        self.done = False

    def clone(self):
        """Return a particle with copied state and a shared sampler and critic."""
        new = SequenceModel(
            unit_sampler=self.unit_sampler,
            critic=self.critic,
            max_tokens=self.max_tokens,
            twist_with_critic=self.twist_with_critic,
            terminate_when=self.terminate_when,
            verbosity=self.verbosity,
        )
        new.context = list(self.context)
        new.weight = self.weight
        new.twist_amount = self.twist_amount
        new.done = self.done
        return new

    def _add_weight(self, amt):
        """Add `amt` to the weight, mapping a NaN result to `-inf`.

        Raises:
            ValueError: If `amt` is `+inf`, which violates the potential contract.
        """
        if amt == float("inf"):
            raise ValueError(
                "A potential returned a log-weight of +inf, which violates the "
                "potential contract."
            )
        w = self.weight + amt
        self.weight = float("-inf") if np.isnan(w) else w

    def score(self, amt):
        self._add_weight(amt)

    def twist(self, amt):
        """Add `amt` to the weight provisionally; `untwist` takes it back."""
        self.twist_amount += amt
        self._add_weight(amt)

    def untwist(self):
        self._add_weight(-self.twist_amount)
        self.twist_amount = 0.0

    def finish(self):
        self.untwist()
        self.done = True

    async def start(self):
        """Score the empty sequence's prefix weight."""
        start_w = await self.unit_sampler.start_weight()
        if start_w == float("-inf"):
            raise ValueError(
                "Start weight is -inf (log(0)). This is likely because a potential "
                "assigns zero weight to the empty sequence under `prefix`, which "
                "violates the potential contract."
            )
        self.score(start_w)

    async def step(self):
        """Advance the particle by one unit, forcing EOS at the token budget."""
        self.untwist()

        if self.max_tokens == 1:
            logw = await self.unit_sampler.logw_eos(self.context)
            unit = EOS
        else:
            unit, logw, _ = await self.unit_sampler.sample(self.context)

        self.score(logw)
        self._append(unit)

        if self.weight == float("-inf"):
            self.finish()
            return

        twist_amt = None
        if self.critic is not None and self.twist_with_critic:
            twist_amt = float(await self.critic.score(self.context))
            if twist_amt == float("-inf"):
                self.score(twist_amt)
                self.finish()
                return
            self.twist(twist_amt)

        if self.verbosity > 0:
            print(self.__repr__())

        self.max_tokens -= 1
        if self.max_tokens == 0 or self.context[-1] is EOS:
            self.finish()
            if self.critic is None:
                return
            if twist_amt is None:
                # Terminal-only critic: reweight once, at termination.
                self.score(float(await self.critic.score(self.context)))
            else:
                # At termination the twist becomes real weight.
                self.score(twist_amt)

    def _append(self, unit):
        """Extend the context by one unit, then by EOS if `terminate_when` fires.

        A multi-token unit ending in EOS is split, keeping `context[-1] is EOS` for
        every terminal sequence.
        """
        if isinstance(unit, list) and unit and unit[-1] is EOS:
            if len(unit) > 1:
                self.context.append(unit[:-1])
            self.context.append(EOS)
        else:
            self.context.append(unit)
        if (
            self.terminate_when is not None
            and self.context[-1] is not EOS
            and self.terminate_when(self.context)
        ):
            self.context.append(EOS)

    def __repr__(self):
        return (
            f"{self.weight:.2f}:\t"
            + colors.magenta % "["
            + (colors.magenta % "|").join(escape(y) for y in self.context)
            + colors.magenta % "]"
        )

clone()

Return a particle with copied state and a shared sampler and critic.

Source code in genlm/control/sampler/smc.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def clone(self):
    """Return a particle with copied state and a shared sampler and critic."""
    new = SequenceModel(
        unit_sampler=self.unit_sampler,
        critic=self.critic,
        max_tokens=self.max_tokens,
        twist_with_critic=self.twist_with_critic,
        terminate_when=self.terminate_when,
        verbosity=self.verbosity,
    )
    new.context = list(self.context)
    new.weight = self.weight
    new.twist_amount = self.twist_amount
    new.done = self.done
    return new

twist(amt)

Add amt to the weight provisionally; untwist takes it back.

Source code in genlm/control/sampler/smc.py 
83
84
85
86
def twist(self, amt):
    """Add `amt` to the weight provisionally; `untwist` takes it back."""
    self.twist_amount += amt
    self._add_weight(amt)

start() async

Score the empty sequence's prefix weight.

Source code in genlm/control/sampler/smc.py 
 96
 97
 98
 99
100
101
102
103
104
105
async def start(self):
    """Score the empty sequence's prefix weight."""
    start_w = await self.unit_sampler.start_weight()
    if start_w == float("-inf"):
        raise ValueError(
            "Start weight is -inf (log(0)). This is likely because a potential "
            "assigns zero weight to the empty sequence under `prefix`, which "
            "violates the potential contract."
        )
    self.score(start_w)

step() async

Advance the particle by one unit, forcing EOS at the token budget.

Source code in genlm/control/sampler/smc.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
async def step(self):
    """Advance the particle by one unit, forcing EOS at the token budget."""
    self.untwist()

    if self.max_tokens == 1:
        logw = await self.unit_sampler.logw_eos(self.context)
        unit = EOS
    else:
        unit, logw, _ = await self.unit_sampler.sample(self.context)

    self.score(logw)
    self._append(unit)

    if self.weight == float("-inf"):
        self.finish()
        return

    twist_amt = None
    if self.critic is not None and self.twist_with_critic:
        twist_amt = float(await self.critic.score(self.context))
        if twist_amt == float("-inf"):
            self.score(twist_amt)
            self.finish()
            return
        self.twist(twist_amt)

    if self.verbosity > 0:
        print(self.__repr__())

    self.max_tokens -= 1
    if self.max_tokens == 0 or self.context[-1] is EOS:
        self.finish()
        if self.critic is None:
            return
        if twist_amt is None:
            # Terminal-only critic: reweight once, at termination.
            self.score(float(await self.critic.score(self.context)))
        else:
            # At termination the twist becomes real weight.
            self.score(twist_amt)

smc_standard(model, n_particles, ess_threshold=0.5, resampling_method='multinomial', json_path=None) async

Run standard SMC over clones of model.

Parameters:

Name Type Description Default
model SequenceModel

The particle template, cloned n_particles times.

 required
n_particles int

Number of particles.

 required
ess_threshold float

Resample when ESS falls below this fraction of n_particles.

 0.5
resampling_method str

One of 'multinomial', 'stratified', 'systematic', 'residual'.

 'multinomial'
json_path str

Path to write the inference record to, for use with the InferenceVisualizer.

 None

Returns:

Type Description
list[SequenceModel]

The completed particles.
Source code in genlm/control/sampler/smc.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
async def smc_standard(
    model,
    n_particles,
    ess_threshold=0.5,
    resampling_method="multinomial",
    json_path=None,
):
    """Run standard SMC over clones of `model`.

    Args:
        model (SequenceModel): The particle template, cloned `n_particles` times.
        n_particles (int): Number of particles.
        ess_threshold (float): Resample when ESS falls below this fraction of
            `n_particles`.
        resampling_method (str): One of 'multinomial', 'stratified', 'systematic',
            'residual'.
        json_path (str, optional): Path to write the inference record to, for use
            with the `InferenceVisualizer`.

    Returns:
        (list[SequenceModel]): The completed particles.
    """
    resample_fn = get_resampling_fn(resampling_method)
    particles = [model.clone() for _ in range(n_particles)]
    await asyncio.gather(*[p.start() for p in particles])

    record = SMCRecord(n_particles) if json_path is not None else None
    ancestor_indices = None

    while any(not p.done for p in particles):
        await asyncio.gather(*[p.step() for p in particles if not p.done])

        if record is not None:
            if not record.history:
                record.add_init(particles)
            elif ancestor_indices is not None:
                record.add_resample(ancestor_indices, particles)
            else:
                record.add_smc_step(particles)

        ancestor_indices = None
        W = np.array([p.weight for p in particles])
        if np.all(W == -np.inf):
            continue
        w_sum = logsumexp(W)
        nw = W - w_sum
        with np.errstate(divide="ignore"):
            log_ess = -logsumexp(nw * 2)
            if log_ess >= np.log(ess_threshold) + np.log(n_particles):
                continue

        probs = np.exp(nw)
        probs /= probs.sum()  # np.random.choice is strict on sum==1
        ancestor_indices = list(resample_fn(probs))
        if record is not None:
            ancestor_indices.sort()  # reproducible record
        avg_weight = w_sum - np.log(n_particles)
        particles = [particles[i].clone() for i in ancestor_indices]
        for p in particles:
            p.weight = avg_weight

    if json_path is not None:
        with open(json_path, "w") as f:
            f.write(record.to_json())
        print(f"Saved record to {json_path}")

    return particles