Update modeling_quiet.py

modeling_quiet.py

@@ -1169,6 +1169,59 @@ def nonzero_mean(x, axis=None):
 def loss_mean(x):
     return x.sum() / (x != 0).sum()
 
+class QuietGenerationMixin(GenerationMixin):
+    def generate(self, input_ids, attention_mask=None, **generate_kwargs):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        max_length = generate_kwargs.get("max_length", 20)
+        temp = generate_kwargs.get("temperature", 1.0)
+
+        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
+
+        streamer = generate_kwargs.get("streamer")
+        if streamer is not None:
+            streamer.put(input_ids)
+            streamer.end()
+
+        return input_ids
+
 class QuietForCausalLM(QuietPreTrainedModel, QuietGenerationMixin):
     _tied_weights_keys = ["lm_head.weight"]
 
@@ -2228,58 +2281,6 @@ class QuietForCausalLM(QuietPreTrainedModel, QuietGenerationMixin):
     QUIET_START_DOCSTRING,
 )
 
-class QuietGenerationMixin(GenerationMixin):
-    def generate(self, input_ids, attention_mask=None, **generate_kwargs):
-        if attention_mask is None:
-            attention_mask = torch.ones_like(input_ids)
-
-        max_length = generate_kwargs.get("max_length", 20)
-        temp = generate_kwargs.get("temperature", 1.0)
-
-        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
-
-        streamer = generate_kwargs.get("streamer")
-        if streamer is not None:
-            streamer.put(input_ids)
-            streamer.end()
-
-        return input_ids
 # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->Quiet, LLAMA->QUIET
 class QuietForSequenceClassification(QuietPreTrainedModel):
     def __init__(self, config):