Method generate

demo/memDec.py:112–172

Greedy decoding with **shared** stopping criteria. We keep two independent KV caches (one per sub‑model) and extend them step‑by‑step.

(  # type: ignore[override]
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.LongTensor] = None,
        max_new_tokens: int = 20,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        do_sample: bool = False,             # must be False (greedy) for now
        generation_config: Optional[GenerationConfig] = None,
        **kwargs,
    )

    # ------------------------------------------------------------------ #
    @torch.no_grad()
    def generate(  # type: ignore[override]
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.LongTensor] = None,
        max_new_tokens: int = 20,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        do_sample: bool = False,             # must be False (greedy) for now
        generation_config: Optional[GenerationConfig] = None,
        **kwargs,
    ):
        """
        Greedy decoding with **shared** stopping criteria.
        We keep two independent KV caches (one per sub‑model) and extend them
        step‑by‑step.
        """
        if do_sample:
            raise ValueError("MemoryDecoder.generate only supports greedy decoding (do_sample=False).")

        device = input_ids.device
        batch_size = input_ids.shape[0]

        # Initialise caches with a single forward.
        outputs = self.forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            **kwargs,
        )
        next_token_logits = outputs["logits"][:, -1, :]            # (B, V)

        base_past = outputs["past_key_values"]
        knn_past = outputs["knn_past_key_values"]

        # Greedy select
        next_tokens = torch.argmax(next_token_logits, dim=-1).unsqueeze(-1)  # (B,1)
        generated = torch.cat([input_ids, next_tokens], dim=-1)              # (B,T+1)

        # --- main loop -------------------------------------------------- #
        num_new_token = 0
        while True:
            if stopping_criteria is not None and False not in stopping_criteria(generated, None):
                break
            if num_new_token >= max_new_tokens:
                break

            outputs = self.forward(
                input_ids=next_tokens,
                attention_mask=None,          # past manages causal masking
                past_key_values=base_past,
                knn_past_key_values=knn_past,
                use_cache=True,
                **kwargs,
            )
            next_token_logits = outputs["logits"][:, -1, :]
            base_past = outputs["past_key_values"]
            knn_past = outputs["knn_past_key_values"]

            next_tokens = torch.argmax(next_token_logits, dim=-1).unsqueeze(-1)
            generated = torch.cat([generated, next_tokens], dim=-1)