Method completion

model.py:351–371

only accepts latent-space data.

(self, data: T.Tensor, temps=(0.8, (0.5, 0.1)), gen_len=None, use_cache=True)

    @T.no_grad()
    def completion(self, data: T.Tensor, temps=(0.8, (0.5, 0.1)), gen_len=None, use_cache=True) -> T.Tensor:
        """
        only accepts latent-space data.
        """
        if use_cache:
            self.init_cache(data.shape[0], data.device, T.bfloat16)

        next_input = generated = data

        target_len = min(data.shape[1] + default(gen_len, data.shape[1]), self.c.stack_config.seq_len)
        
        for _ in tqdm0(range(data.shape[1], target_len)):
            model_input = next_input if use_cache else generated

            next_input = self.next_latent(model_input, temps)
    
            generated = T.cat([generated, next_input], dim=1)

        if use_cache:
            self.deinit_cache()
        return generated