Correct way to do infilling with the diffusion model

[marinegor]

Hi everyone,

I wonder, what's the correct way to do infilling with the model?

As far as I can see, the model only provides an interface to do an unconditional sample, i.e. diffusion.restore_model_and_sample(...). But what if I, for instance, what to fill in the blanks, what'd be the correct (i.e. intended by the authors) way to do that?

For instance, I can do something like this:

diffusion = Diffusion.from_checkpoint(...)
sentence = 'London is the capital of _'
assert '_' == diffusion.tokenizer.mask_token

tokenized = diffusion.tokenizer(sentence)
restored = diffusion.tokenizer.batch_decode(diffusion(tokenized))
...

but I feel that it's a wrong way, since we're not really doing the denoising.

In turn, doing something similar to diffusion._sample(...) feels like a more correct way to do that, but then it's unclear at which stage should we incorporate the already known tokens.