Update modeling_quiet.py

modeling_quiet.py

@@ -1328,46 +1328,42 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
     def generate(self, input_ids, attention_mask=None, **kwargs):
         if attention_mask is None:
             attention_mask = torch.ones_like(input_ids)
-
+        
         max_length = kwargs.get("max_length", 20)
         temp = kwargs.get("temperature", 1.0)
 
-        batch_size = input_ids.shape[0]
-        device = input_ids.device
-
-        finished_generating = torch.zeros(batch_size, dtype=torch.bool, device=device)
-
-        for cur_token_idx in range(max_length):
-            new_ids = self(
-                input_ids[~finished_generating],
-                attention_mask=attention_mask[~finished_generating]
-            )['logits']
-
-            new_ids[:, :, self.tokenizer.vocab_size:] = -float("inf")
-
-            for list_idx, answer_idx in enumerate((~finished_generating).nonzero(as_tuple=True)[0]):
-                base_answer_ids = input_ids[answer_idx]
-                new_answer_ids = new_ids[list_idx]
-                last_token_idx = (base_answer_ids != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0].max()
-
-                new_ids_sampled = torch.multinomial(
-                    torch.nn.functional.softmax(new_answer_ids[last_token_idx] / temp, dim=-1), 1)
-
-                if last_token_idx + 1 >= len(base_answer_ids):
-                    new_padding = torch.full((batch_size, 1), self.tokenizer.pad_token_id, dtype=torch.long, device=device)
-                    input_ids = torch.cat([input_ids, new_padding], dim=-1)
-                    attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
-
-                attention_mask[answer_idx, last_token_idx + 1] = 1
-                input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
-
-                if new_ids_sampled in [self.tokenizer.eos_token_id, self.tokenizer.bos_token_id, self.tokenizer.pad_token_id]:
-                    finished_generating[answer_idx] = 1
-
-            if finished_generating.all():
-                break
-
-        return input_ids
+        with torch.no_grad():
+            finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=input_ids.device)
+            for cur_token_idx in range(max_length):
+                # Sample the next token
+                new_ids = self(
+                    input_ids[~finished_generating],
+                    attention_mask=attention_mask[~finished_generating]
+                )['logits']
+                # Mask out the start and end thought tokens so we don't accidentally sample them
+                new_ids[:, :, self.tokenizer.vocab_size:] = -float("inf")
+                for list_idx, answer_idx in enumerate((~finished_generating).nonzero(as_tuple=True)[0]):
+                    # Find the index of the last token that is not padding
+                    base_answer_ids = input_ids[answer_idx]
+                    new_answer_ids = new_ids[list_idx]
+                    last_token_idx = (base_answer_ids != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0].max()
+
+                    new_ids_sampled = torch.multinomial(
+                            torch.nn.functional.softmax(new_answer_ids[last_token_idx] / temp, dim=-1), 1)
+                    # Assign the new id to the last token
+                    if last_token_idx + 1 >= len(base_answer_ids):
+                        # Add padding everywhere
+                        new_padding = torch.full((len(input_ids), 1), self.tokenizer.pad_token_id, dtype=torch.long,
+                                                 device=input_ids.device)
+                        input_ids = torch.cat([input_ids, new_padding], dim=-1)
+                        attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
+                    attention_mask[answer_idx, last_token_idx + 1] = 1
+                    input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
+                    if new_ids_sampled == self.tokenizer.eos_token_id or new_ids_sampled == self.tokenizer.bos_token_id or new_ids_sampled == self.tokenizer.pad_token_id:
+                        finished_generating[answer_idx] = 1
+                if finished_generating.all():
+                    break
+        return input_ids, attention_mask
 
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)