Update generate.py

generate.py

@@ -1,71 +1,214 @@
 import torch
-from transformers import LogitsProcessorList, StoppingCriteriaList
 
-def generate(
+def custom_generate(
     self,
     input_ids,
     attention_mask=None,
     max_new_tokens=None,
-    temperature=1.0,
-    do_sample=True,
+    min_length=None,
+    do_sample=None,
+    early_stopping=None,
+    num_beams=None,
+    temperature=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,
-    **kwargs
+    streamer=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,
+    **kwargs,
 ):
     device = input_ids.device
-    logits_processor = LogitsProcessorList()
-    stopping_criteria = StoppingCriteriaList()
-    
-    if attention_mask is None:
-        attention_mask = torch.ones_like(input_ids)
-    
-    # Initialize unfinished sentences to manage the loop for early stopping
-    unfinished_sents = input_ids.new(input_ids.shape[0]).fill_(1)
-    
-    cur_len = input_ids.shape[1]
+    with torch.no_grad():
+        finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=device)
+
+        if max_new_tokens is None:
+            max_new_tokens = 50  # Default value if not specified
+        for cur_token_idx in range(max_new_tokens):
+            new_ids = self(
+                input_ids[~finished_generating],
+                attention_mask=attention_mask[~finished_generating] if attention_mask is not None else None,
+                **kwargs
+            )['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] / temperature, 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=device)
+                    input_ids = torch.cat([input_ids, new_padding], dim=-1)
+                    if attention_mask is not None:
+                        attention_mask = torch.cat([attention_mask, torch.zeros_like(new_padding)], dim=-1)
+
+                if attention_mask is not None:
+                    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
+
+                # Check if the end token is generated
+                if new_ids_sampled == self.tokenizer.convert_tokens_to_ids("</s>"):
+                    finished_generating[answer_idx] = 1
+            
+            if finished_generating.all():
+                break
+
+            if streamer is not None:
+                streamer.put(new_ids_sampled)
+
+    from collections import namedtuple
+    GenerateOutput = namedtuple("GenerateOutput", ["sequences", "scores", "attentions", "hidden_states"])
     
-    while cur_len < max_new_tokens:
-        model_outputs = self(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            use_cache=True,
-            return_dict=True
-        )
-        
-        next_token_logits = model_outputs.logits[:, -1, :]
-        
-        # Processing logits to avoid generating undesired tokens
-        next_token_logits[:, pad_token_id] = -float('inf')  # Never select pad
-        next_token_logits[:, eos_token_id] = -float('inf')  # Avoid generating end token prematurely
-
-        # Apply temperature scaling and softmax to generate probabilities
-        if do_sample:
-            probabilities = torch.nn.functional.softmax(next_token_logits / temperature, dim=-1)
-            next_token = torch.multinomial(probabilities, num_samples=1).squeeze(1)
-        else:
-            next_token = next_token_logits.argmax(dim=-1)
-        
-        # Update input_ids and attention_mask
-        input_ids = torch.cat([input_ids, next_token.unsqueeze(-1)], dim=1)
-        new_attention = torch.ones_like(input_ids[:, 0:1])
-        attention_mask = torch.cat([attention_mask, new_attention], dim=1)
-        
-        # Check unfinished sentences
-        unfinished_sents.mul_(next_token.ne(eos_token_id).long())
-        if unfinished_sents.max() == 0:
-            break
-        
-        cur_len += 1
+    # Convert the generated token IDs to a tensor
+    generated_token_ids_tensor = input_ids
     
-    # Optionally return additional information
-    if kwargs.get('return_dict_in_generate', False):
-        output = {
-            "sequences": input_ids,
-            "scores": None,  # Placeholder for when score calculations are implemented
-            "attentions": model_outputs.attentions,
-            "hidden_states": model_outputs.hidden_states
-        }
-    else:
-        output = input_ids
+    output = GenerateOutput(
+        sequences=generated_token_ids_tensor,
+        scores=None,
+        attentions=None,
+        hidden_states=None
+    )
     
     return output
+
+
+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,
+):
+    # 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
+
+    output = 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 output