Update ts_generation_mixin.py

ts_generation_mixin.py

@@ -6,8 +6,38 @@ from transformers.generation import validate_stopping_criteria, EosTokenCriteria
 from transformers.generation.utils import GenerateNonBeamOutput, GenerateEncoderDecoderOutput, GenerateDecoderOnlyOutput, GenerationConfig, GenerateOutput
 from transformers.utils import ModelOutput
 
+
 class TSGenerationMixin(GenerationMixin):
 
+    @torch.no_grad()
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        logits_processor: Optional[LogitsProcessorList] = None,
+        stopping_criteria: Optional[StoppingCriteriaList] = None,
+        prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
+        synced_gpus: Optional[bool] = None,
+        assistant_model: Optional["PreTrainedModel"] = None,
+        streamer: Optional["BaseStreamer"] = None,
+        negative_prompt_ids: Optional[torch.Tensor] = None,
+        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[GenerateOutput, torch.LongTensor]:
+        if len(inputs.shape) == 2:
+            batch_size, cur_len = inputs.shape
+            if cur_len < self.config.input_token_len:
+                raise ValueError(
+                    f"Input length must be at least {self.config.input_token_len}")
+            elif cur_len % self.config.input_token_len != 0:
+                new_len = (cur_len // self.config.input_token_len) * \
+                    self.config.input_token_len
+                inputs = inputs[:, -new_len:]
+        else:
+            raise ValueError('Input shape must be: [batch_size, seq_len]')
+        return super().generate(inputs=inputs, generation_config=generation_config, logits_processor=logits_processor, stopping_criteria=stopping_criteria, prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, synced_gpus=synced_gpus, assistant_model=assistant_model, streamer=streamer, negative_prompt_ids=negative_prompt_ids, negative_prompt_attention_mask=negative_prompt_attention_mask, **kwargs)
+        
+        
     def _greedy_search(
             self,
             input_ids: torch.Tensor,
@@ -26,19 +56,7 @@ class TSGenerationMixin(GenerationMixin):
             **model_kwargs,
     ) -> Union[GenerateNonBeamOutput, torch.Tensor]:
         input_ids = input_ids.to(self.device)
-        initial_input_length = input_ids.shape[1]
-        if len(input_ids.shape) == 2:
-            batch_size, cur_len = input_ids.shape
-            if cur_len < self.config.input_token_len:
-                raise ValueError(
-                    f"Input length must be at least {self.config.input_token_len}")
-            elif cur_len % self.config.input_token_len != 0:
-                new_len = (cur_len // self.config.input_token_len) * \
-                    self.config.input_token_len
-                input_ids = input_ids[:, -new_len:]
-        else:
-            raise ValueError('Input shape must be: [batch_size, seq_len]')
-
+        batch_size, cur_len = input_ids.shape
         # init values
         logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
         stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
@@ -106,9 +124,8 @@ class TSGenerationMixin(GenerationMixin):
             batch_size, dtype=torch.long, device=input_ids.device)
         model_kwargs["cache_position"] = torch.arange(
             cur_len, device=input_ids.device)
-        true_seq_len = input_ids.shape[1] // self.config.input_token_len
+        true_seq_len = cur_len // self.config.input_token_len
         model_kwargs["attention_mask"] = model_kwargs["attention_mask"][:, -true_seq_len:]
-
         max_length = stopping_criteria.max_length
         while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):
             # prepare model inputs
@@ -129,7 +146,7 @@ class TSGenerationMixin(GenerationMixin):
             if synced_gpus and this_peer_finished:
                 continue  # don't waste resources running the code we don't need
 
-            next_token_logits = outputs.logits[:, -1, :]
+            next_token_logits = outputs.logits
 
             # pre-process distribution
             next_tokens_scores = logits_processor(input_ids, next_token_logits)
@@ -212,7 +229,7 @@ class TSGenerationMixin(GenerationMixin):
                     past_key_values=model_kwargs.get("past_key_values"),
                 )
         else:
-            return input_ids[:, -(max_length - initial_input_length):]
+            return input_ids[:, -(max_length - cur_len):]
 
     def _update_model_kwargs_for_generation(
             self,