Update generate.py

generate.py

@@ -1,15 +1,83 @@
 import torch
 from transformers.utils import logging
+from transformers.generation.utils import (
+    GenerationMixin,
+    validate_stopping_criteria,
+    StoppingCriteriaList,
+)
 
 logger = logging.get_logger(__name__)
 
 @torch.no_grad()
+def custom_generate(model, input_ids, attention_mask, max_new_tokens, streamer, **kwargs):
+    finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=input_ids.device)
+    for cur_token_idx in range(max_new_tokens):
+        # Sample the next token
+        new_ids = model(
+            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[:, :, model.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 != model.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] / kwargs.get("temperature", 1.0), 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), model.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 == model.tokenizer.eos_token_id or new_ids_sampled == model.tokenizer.bos_token_id or new_ids_sampled == model.tokenizer.pad_token_id:
+                finished_generating[answer_idx] = 1
+            # Check if the end token is generated
+            if new_ids_sampled == model.tokenizer.convert_tokens_to_ids("<|/assistant|>"):
+                finished_generating[answer_idx] = 1
+        if finished_generating.all():
+            break
+        streamer.put(new_ids_sampled)
+    return input_ids, attention_mask
+
 def generate(
     self,
     input_ids,
     attention_mask=None,
     max_length=None,
+    min_length=None,
+    do_sample=None,
+    early_stopping=None,
+    num_beams=None,
     temperature=1.0,
+    top_k=None,
+    top_p=None,
+    repetition_penalty=None,
+    bad_words_ids=None,
+    bos_token_id=None,
+    pad_token_id=None,
+    eos_token_id=None,
+    length_penalty=None,
+    no_repeat_ngram_size=None,
+    num_return_sequences=None,
+    decoder_start_token_id=None,
+    use_cache=None,
+    num_beam_groups=None,
+    diversity_penalty=None,
+    prefix_allowed_tokens_fn=None,
+    output_attentions=None,
+    output_hidden_states=None,
+    output_scores=None,
+    return_dict_in_generate=None,
+    forced_bos_token_id=None,
+    forced_eos_token_id=None,
+    remove_invalid_values=None,
+    synced_gpus=None,
     n_ahead=4,
     n_ahead_talk=4,
     merged_talk_heads=True,
@@ -23,10 +91,8 @@ def generate(
     use_weighted_talk_head=True,
     trust_remote_code=True,
     torch_dtype=torch.bfloat16,
-    **kwargs
+    **model_kwargs,
 ):
-    batch_size, seq_length = input_ids.shape
-
     # Set model attributes
     self.max_thoughts = n_ahead + n_ahead_talk + 1
     self.merged_talk_heads = merged_talk_heads
@@ -51,38 +117,18 @@ def generate(
     self.rm_initialized = True
     self.original_mode = False
 
-    # Keep track of which sequences have finished generating
-    finished_generating = torch.zeros(batch_size, dtype=torch.bool, device=input_ids.device)
-
-    while input_ids.shape[-1] < max_length:
-        # Get the model outputs
-        model_outputs = self(
-            input_ids[~finished_generating],
-            attention_mask=attention_mask[~finished_generating] if attention_mask is not None else None,
-            **kwargs
-        )
-        logits = model_outputs.logits[:, -1, :]
-
-        # Apply temperature scaling
-        logits = logits / temperature
-
-        # Sample the next token
-        next_token_logits = logits
-        next_token_id = torch.multinomial(torch.softmax(next_token_logits, dim=-1), num_samples=1).squeeze(-1)
-
-        # Assign the sampled token to the sequences that are still generating
-        input_ids = torch.cat([input_ids, next_token_id.unsqueeze(-1)], dim=-1)
+    # Initialize a TextStreamer for streaming the generated text
+    streamer = TextStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-        # Update the attention mask if provided
-        if attention_mask is not None:
-            attention_mask = torch.cat([attention_mask, torch.ones_like(next_token_id.unsqueeze(-1))], dim=-1)
-
-        # Mark sequences as finished if the end token is generated
-        finished_generating = finished_generating | (next_token_id == self.tokenizer.eos_token_id)
-
-        # Stop generation if all sequences are finished
-        if finished_generating.all():
-            break
+    # Generate using the custom generate function
+    input_ids, attention_mask = custom_generate(
+        self,
+        input_ids,
+        attention_mask,
+        max_length,
+        streamer,
+        temperature=temperature,
+        **model_kwargs,
+    )
 
-    # Return the generated token IDs and attention mask
     return input_ids, attention_mask