Update generate.py

generate.py

@@ -104,4 +104,127 @@ def custom_generate(
         print("Generated Token IDs shape:", generated_token_ids.shape)
         print("Generated Token IDs:", generated_token_ids)
 
-        return generated_token_ids
+        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,
+    )
+
+    return generated_token_ids