Update generate.py

generate.py

@@ -1,11 +1,3 @@
-import torch
-from transformers.generation.utils import (
-    GenerationMixin,
-    validate_stopping_criteria,
-    StoppingCriteriaList,
-)
-from transformers import TextStreamer
-
 def custom_generate(
     self,
     input_ids,
@@ -42,6 +34,11 @@ def custom_generate(
     synced_gpus=None,
     **kwargs,
 ):
+    print("Input IDs shape:", input_ids.shape)
+    print("Input IDs:", input_ids)
+    print("Attention Mask shape:", attention_mask.shape if attention_mask is not None else None)
+    print("Attention Mask:", attention_mask)
+
     if input_ids is None or input_ids.nelement() == 0:
         # If input_ids is None or an empty tensor, create a default input tensor
         input_ids = torch.LongTensor([[self.tokenizer.bos_token_id]]).to(self.device)
@@ -61,6 +58,8 @@ def custom_generate(
                 **kwargs
             )['logits']
 
+            print(f"Step {cur_token_idx + 1}: New IDs shape: {new_ids.shape}")
+
             # Mask out the start and end thought tokens so we don't accidentally sample them
             new_ids[:, :, self.tokenizer.vocab_size:] = -float("inf")
 
@@ -73,6 +72,8 @@ def custom_generate(
                 new_ids_sampled = torch.multinomial(
                     torch.nn.functional.softmax(new_answer_ids[last_token_idx] / temperature, dim=-1), 1)
 
+                print(f"Step {cur_token_idx + 1}: New IDs sampled: {new_ids_sampled}")
+
                 # Assign the new id to the last token
                 if last_token_idx + 1 >= len(base_answer_ids):
                     # Add padding everywhere
@@ -100,128 +101,7 @@ def custom_generate(
             if streamer is not None:
                 streamer.put(new_ids_sampled)
 
-        return generated_token_ids
-
-
-def generate(
-    self,
-    input_ids,
-    attention_mask=None,
-    max_new_tokens=None,
-    min_length=None,
-    do_sample=None,
-    early_stopping=None,
-    num_beams=None,
-    temperature=1.1,
-    streamer=None,
-    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,
-    merged_lm_and_talk_heads=False,
-    merged_lm_and_think_heads=True,
-    use_concat_talk_head=True,
-    use_shallow_think=True,
-    use_shallow_talk=False,
-    use_complex_think_head=False,
-    use_complex_talk_head=True,
-    use_weighted_talk_head=True,
-    trust_remote_code=True,
-    torch_dtype=torch.bfloat16,
-    **model_kwargs,
-):
-
-    if max_new_tokens is None:
-        max_new_tokens = 128 
-    
-    # Set model attributes
-    self.max_thoughts = n_ahead + n_ahead_talk + 1
-    self.merged_talk_heads = merged_talk_heads
-    self.merged_lm_and_talk_heads = merged_lm_and_talk_heads
-    self.merged_lm_and_think_heads = merged_lm_and_think_heads
-    self.use_concat_talk_head = use_concat_talk_head
-    self.use_shallow_think = use_shallow_think
-    self.use_shallow_talk = use_shallow_talk
-    self.use_complex_think_head = use_complex_think_head
-    self.use_complex_talk_head = use_complex_talk_head
-    self.use_weighted_talk_head = use_weighted_talk_head
-
-    # Set model properties
-    self.use_end_thought_token = True
-    self.use_start_thought_token = True
-    self.n_ahead = n_ahead
-    self.n_passes = 1
-    self.eval_mode = True
-    self.first_run = False
-    self.rm_initialized = True
-    self.original_mode = False
-
-    # Check if the input is a string (for compatibility with text-generation-webui)
-    if isinstance(input_ids, str):
-        input_ids = self.tokenizer.encode(input_ids, return_tensors='pt')
-
-    # Move input_ids and attention_mask to the same device as the model
-    input_ids = input_ids.to(self.device)
-    if attention_mask is not None:
-        attention_mask = attention_mask.to(self.device)
-
-    generated_token_ids = custom_generate(
-        self,
-        input_ids=input_ids, 
-        attention_mask=attention_mask,
-        max_new_tokens=max_new_tokens,
-        min_length=min_length,
-        do_sample=do_sample,
-        early_stopping=early_stopping,
-        num_beams=num_beams,
-        temperature=temperature,
-        top_k=top_k,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        bad_words_ids=bad_words_ids,
-        bos_token_id=bos_token_id,
-        pad_token_id=pad_token_id,
-        eos_token_id=eos_token_id,
-        length_penalty=length_penalty,
-        no_repeat_ngram_size=no_repeat_ngram_size,
-        num_return_sequences=num_return_sequences,
-        decoder_start_token_id=decoder_start_token_id,
-        use_cache=use_cache,
-        num_beam_groups=num_beam_groups,
-        diversity_penalty=diversity_penalty,
-        prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
-        output_attentions=output_attentions,
-        output_hidden_states=output_hidden_states,
-        output_scores=output_scores,
-        return_dict_in_generate=return_dict_in_generate,
-        forced_bos_token_id=forced_bos_token_id,
-        forced_eos_token_id=forced_eos_token_id,
-        remove_invalid_values=remove_invalid_values,
-        synced_gpus=synced_gpus,
-        streamer=streamer,
-        **model_kwargs,
-    )
+        print("Generated Token IDs shape:", generated_token_ids.shape)
+        print("Generated Token IDs:", generated_token_ids)
 
-    return generated_token_ids
+        return generated_token_ids