def sample(
        self,
        prompts: str | list[str] | list[int] | list[list[int]],
        num_tokens: int = 5,
        temperature: float = 1.0,
        top_p: float | None = None,
        hooks: TransformerHooks = TransformerHooks(),
    ) -> dict[str, Any]:
        """
        Sampling with the transformer!

        If top_p is set, then nucleus sampling is used.
        Otherwise, the sampling will be Categorical.
        If temperature=0, sampling is deterministic (and top_p is ignored).

        (Warning: when using torch.use_deterministic_algorithms(True),
        nucleus sampling will throw an error. It depends on torch.cumsum,
        which unfortunately has no deterministic implementation in torch.)

        Output is a dict {'tokens': list[list[int]], 'strings': list[str]}
        """
        prompts = [prompts] if isinstance(prompts, str) else prompts
        if isinstance(prompts[0], str):
            X: list[list[int]] = [self.enc.encode(prompt) for prompt in prompts]
        elif isinstance(prompts[0], int):
            X = [prompts]
        else:
            X = prompts
        X, pad = prep_input_and_pad(X, "left", self.device)
        H = None
        beta = 1 / max(temperature, 1e-10)
        out = {
            "tokens": [[] for _ in prompts],
            "strings": ["" for _ in prompts],
        }

        # sampling loop
        for _ in range(num_tokens):
            with torch.no_grad():
                # get logits
                Y, H = self.forward(X, H, pad, hooks=hooks)
                logits = Y[:, -1] * beta

                # sampling only works if logits are floats
                logits = logits.float()

                # perform sampling
                if temperature == 0:
                    tokens = torch.argmax(logits, dim=-1)
                elif top_p is not None:
                    tokens = nucleus_sample(logits, top_p)
                else:
                    tokens = Categorical(logits=logits).sample()
                X, pad = tokens.unsqueeze(-1), None

            for batch_idx, token in enumerate(tokens):
                out["tokens"][batch_idx].append(token.item())
                out["strings"][batch_idx] += self.enc.decode([token.item()])

        return out

        Sampling with the transformer!

        If top_p is set, then nucleus sampling is used.
        Otherwise, the sampling will be Categorical.
        If temperature=0, sampling is deterministic (and top_p is ignored).

        (Warning: when using torch.use_deterministic_algorithms(True),
        nucleus sampling will throw an error. It depends on torch.cumsum,
        which unfortunately has no deterministic implementation in torch.)

        Output is a dict {'tokens': list[list[int]], 'strings': list[str]}
        
neuron_explainer/models/transformer.py