-
-class Gang(nn.Module):
- def __init__(self, models, nb_models_for_generation, mode="groupthink"):
- super().__init__()
- self.models = nn.ModuleList(models)
- self.nb_models_for_generation = nb_models_for_generation
- self.mode = mode
-
- def forward(self, bs):
- # If first = 0, we are re-starting an auto-regressive process,
- # that's the right moment to randomize who gonna do it
- if bs.first == 0:
- self.models_to_use = [
- self.models[k]
- for k in torch.randperm(len(self.models))[
- : self.nb_models_for_generation
- ]
- ]
-
- all_the_logits = torch.cat(
- [model(bs).x[None] for model in self.models_to_use], dim=0
- )
-
- if self.mode == "groupthink":
- y = all_the_logits.mean(dim=0)
- elif self.mode == "groupwork":
- m = torch.rand(all_the_logits.size(), device=all_the_logits.device)
- m = (m.sort(dim=0).indices == 0).long()
- y = (y * m).sum(dim=0)
- else:
- raise ValueError(f"Invalid mode {self.mode}")
-
- return BracketedSequence(y, bs.first, bs.nb)
-
-
-######################################################################
-