確率的プログラミングとLLM

SMCステアリングを用いたFeynman-Kac Transformer modelsの生成過程:

1. 初期トークン／プロンプト ---> [Transformerモデル] ---+---> 候補トークン群
   (確率的制約は適用されていない)

2. 候補トークン群 ---------> [SMCステアリング] ------> 制約に基づいて選ばれたトークン
   (SMCが確率的プログラムとして提供される制約に基づき、候補トークンを精査・選択)

3. 選ばれたトークン --------> [Transformerモデル] ---+---> 新たな候補トークン群
   (選択されたトークンが次の入力として用いられる)

4. 新たな候補トークン群 ---> [SMCステアリング] ------> 制約に基づいて選ばれる次のトークン
   ...

async def main():
    constraint_model = ConstraintModel(prompt, 50)
    particles = await smc_standard(constraint_model, 40)
    for p in particles:
        print(f"{p.context}")

asyncio.run(main())

LLM = CachedCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", auth_token=HF_AUTH_TOKEN)
LLM.batch_size = 40

particles = [copy.deepcopy(model) for _ in range(n_particles)]

await asyncio.gather(*[p.step() for p in particles if not p.done_stepping()])

class ConstraintModel(Model):
    def __init__(self, prompt, max_tokens):
        super().__init__()
        self.context = LMContext(LLM, prompt)
        self.q       = LMContext(LLM, prompt)
        self.max_tokens = max_tokens

    async def step(self):
        # ここで使えるトークンを絞ってる(例の中では5文字以内)
        mask = self.active_constraint_mask()
        
        # next_token() が次のトークン単一、ではなくトークンの distribution を返してる
        # proposal が mask に基づいた distribution(ここちょっと intervene 内で dist.log_prob(x) が何してるのか謎。return されてる x 使ってないし)
        token = await self.sample(self.context.next_token(), proposal = await self.proposal(mask))

        # Condition on constraint — a no-op since proposal already guarantees the constraint
        self.condition(token.token_id in mask)
        
        # Reduce number of max tokens remaining
        self.max_tokens -= 1
        
        print(f"{self.context}")

        # Check if done
        if token == LLM.tokenizer.eos_token_id or self.max_tokens == 0:
            self.finish()
            return

class LMContext:    
    def __init__(self, lm, prompt, temp=1.0):
        """Create a new `LMContext` with a given prompt and temperature.
        
        Args:
            lm (hfppl.llms.CachedCausalLM): the language model for which this is a context.
            prompt (str): a string with which to initialize the context. Will be tokenized using `lm.tokenizer`.
            temp (float): temeprature for next-token distribution (0 < temp < float('inf'))"""
        self.lm                  = lm
        self.s                   = TokenSequence(lm, prompt)
        self.next_token_logprobs = log_softmax(lm.next_token_logprobs_unbatched(self.s.seq) / temp)
        self.temp                = temp
        self.NO_MASK    = set(range(len(self.lm.vocab)))
        self.model_mask = self.NO_MASK
        self.prompt_string_length = len(str(self.s))
        self.show_prompt = False
        
    def next_token(self):
        """Distribution over the next token.
        
        Sampling or observing from this distribution advances the state of this `LMContext` instance."""
        return LMNextToken(self)

    async def intervene(self, dist, x):
        """Force the distribution to take on the value `x`, but do not _condition_ on this result.
        
        This is useful primarily with distributions that have side effects (e.g., modifying some state).
        For example, a model with the code
        
        ```python
        token_1 = await self.sample(self.stateful_lm.next_token())
        await self.observe(self.stateful_lm.next_token(), token_2)
        ```
        
        encodes a posterior inference problem, to find `token_1` values that *likely preceded* `token_2`. By contrast,
        
        ```python
        token_1 = await self.sample(stateful_lm.next_token())
        await self.intervene(self.stateful_lm.next_token(), token_2)
        ```
        
        encodes a much easier task: freely generate `token_1` and then force-feed `token_2` as the following token.
        
        Args:
            dist (hfppl.distributions.distribution.Distribution): the distribution on which to intervene.
            x: the value to intervene with.
        """
        await dist.log_prob(x)
        return x

await self.intervene(self.q.mask_dist(mask), True)

    def mask_dist(self, mask):
        """Bernoulli distribution, with probability of True equal to the probability that the next token of this `LMContext` belongs
        to the given mask.
        
        Sampling or observing from this distribution modifies the state of this `LMContext` instance, so that
        the `next_token()` distribution either *will* (if True) or *will not* (if False) generate a token from
        the given mask.
        
        Args:
            mask: a `set(int)` specifying which token ids are included within the mask."""
        return LMTokenMask(self, mask)     

class LMTokenMask(Distribution):
    def __init__(self, ctx, mask):
        self.ctx  = ctx
        self.mask = mask
        
    async def log_prob(self, v):
        good_tokens  = self.ctx.model_mask.intersection(self.mask) if v else self.ctx.model_mask - self.mask
        bad_tokens   = [i for i in self.ctx.model_mask if i not in good_tokens]
        logprob_good = logsumexp(self.ctx.next_token_logprobs[list(good_tokens)])
        self.ctx.next_token_logprobs[bad_tokens] = float('-inf')
        self.ctx.next_token_logprobs -= logprob_good
        self.ctx.model_mask = good_tokens
        return logprob_good