模型的generate方法

[jingyaogong]

收到了几次这样的问题反馈：

代码里调用的 model.generate() 的这个generate函数到底在哪？

generated_ids = model.generate(
    inputs=inputs["input_ids"], 
    attention_mask=inputs["attention_mask"],
    ...
)

之前版本为了精简代码，移除了自定义的 generate 实现，默认依赖 transformers 的 GenerationMixin

如果想自己掌控采样的逻辑，可在 MiniMindForCausalLM 类中添加如下方法：

class MiniMindForCausalLM:
# others...
    # 以下仅供参考，写法不固定（不确保永远适配）
    @torch.inference_mode()
    def generate(self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=None, streamer=None, use_cache=True, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs):
        input_ids = kwargs.pop("input_ids", inputs).repeat(num_return_sequences, 1)
        attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
        past_key_values = kwargs.pop("past_key_values", None)
        finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device)
        if streamer: streamer.put(input_ids.cpu())
        for _ in range(max_new_tokens):
            past_len = past_key_values[0][0].shape[1] if past_key_values else 0
            outputs = self.forward(input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, **kwargs)
            attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
            logits = outputs.logits[:, -1, :] / temperature
            if repetition_penalty != 1.0:
                for i in range(input_ids.shape[0]): logits[i, torch.unique(input_ids[i])] /= repetition_penalty
            if top_k > 0: 
                logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
                mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
                logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')
            next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
            if eos_token_id is not None: next_token = torch.where(finished.unsqueeze(-1), next_token.new_full((next_token.shape[0], 1), eos_token_id), next_token)
            input_ids = torch.cat([input_ids, next_token], dim=-1)
            past_key_values = outputs.past_key_values if use_cache else None
            if streamer: streamer.put(next_token.cpu())
            if eos_token_id is not None:
                finished |= next_token.squeeze(-1).eq(eos_token_id)
                if finished.all(): break
        if streamer: streamer.end()
        if kwargs.get("return_kv"): return {'generated_ids': input_ids, 'past_kv': past_key_values}
        return input_ids

好处是

• 完全可控：生成流程清晰透明，调试和定制都很方便
• 零依赖：不依赖 GenerationMixin 的生成框架（例如transformers>=4.48升级后，past_key_values初始化发生变动，这边必须同步适配）
• 易于理解：每行代码都能看懂，适合学习 temperature / top_p / top_k 的工作原理
• 灵活修改：想调整采样策略随时改，不用研究复杂的继承体系

坏处是

• 功能有限：不支持 beam search、stopping_criteria、diversity penalty 等很多高级特性（但目前基本都不会用到）
• 维护成本：transformers 未来更新后很可能需要手动适配
• 兼容性：部分第三方库可能依赖标准的 GenerationMixin 接口

后续

目前这个实现足以覆盖日常使用场景。长期来看有两条路：

1. 保持现状：继续维护精简版 generate，保持代码简洁
2. 回归标准：移除自定义实现，直接继承 GenerationMixin

适配本身不难（无非是循环调用 forward 进行逐 token 的自回归采样），但会让 model 结构变得臃肿（甚至比模型全部代码还长，具体可以参考 https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py#L2252 ），不利于维护和学习。

未来会根据具体情况选择保留与否——尽可能维护一个自己实现的 generate，但当更多特性加入、臃肿程度超过权衡点时，会直接继承 GenerationMixin。

保留此条 discussion 用于解释这个问题