A simple scheduler for masking policies. Args: dim: int, the dimensionality of the state space. schedule: a DiffusionSchedule object for scheduling rates.
(self,
dim,
schedule,
tokenizer,
use_fast_inference=True,
sample_cls=None,
word_freq=None,
word_freq_lambda=0.,
history_per_term=10,
device='cuda:0'
)
| 305 | |
| 306 | class MaskDiffusion(DiscreteDiffusionMatrixBase): |
| 307 | def __init__(self, |
| 308 | dim, |
| 309 | schedule, |
| 310 | tokenizer, |
| 311 | use_fast_inference=True, |
| 312 | sample_cls=None, |
| 313 | word_freq=None, |
| 314 | word_freq_lambda=0., |
| 315 | history_per_term=10, |
| 316 | device='cuda:0' |
| 317 | ): |
| 318 | """A simple scheduler for masking policies. |
| 319 | Args: |
| 320 | dim: int, the dimensionality of the state space. |
| 321 | schedule: a DiffusionSchedule object for scheduling rates. |
| 322 | """ |
| 323 | |
| 324 | self.num_steps = schedule.num_steps |
| 325 | self.sample_cls=sample_cls |
| 326 | self.schedule = schedule |
| 327 | self.use_fast_inference = use_fast_inference |
| 328 | self.dim = dim # allow mask |
| 329 | self.tokenizer = tokenizer |
| 330 | self.mask = torch.nn.functional.one_hot(torch.tensor(self.tokenizer.mask_token_id), num_classes=self.dim).unsqueeze(1).repeat(1, self.dim).float() |
| 331 | self.device = device |
| 332 | self.state = self._create_state() |
| 333 | self.word_freq = word_freq.to(device) |
| 334 | self.history_per_term = history_per_term |
| 335 | self._loss_history = np.zeros( |
| 336 | [self.num_steps, history_per_term], dtype=np.float64 |
| 337 | ) |
| 338 | self._loss_counts = np.zeros([self.num_steps], dtype=np.int) |
| 339 | |
| 340 | import math |
| 341 | self.word_freq_lambda = word_freq_lambda * torch.sin(torch.arange(schedule.num_steps + 1, device=device) / schedule.num_steps * math.pi) |
| 342 | |
| 343 | def _create_state(self): |
| 344 | """Initializes values used by the get function.""" |
nothing calls this directly
no test coverage detected