Update modeling_quiet.py

modeling_quiet.py

@@ -1334,6 +1334,17 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        max_length: Optional[int] = None,
+        num_return_sequences: Optional[int] = 1,
+        no_repeat_ngram_size: Optional[int] = 2,
+        early_stopping: Optional[bool] = True,
+        num_beams: Optional[int] = 1,
+        temperature: Optional[float] = 1.0,
+        repetition_penalty: Optional[float] = 1.2,
+        length_penalty: Optional[float] = 1.0,
+        pad_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        streamer: Optional[TextStreamer] = None,
     ):
         batch_size, seq_len = input_ids.shape
 
@@ -2128,50 +2139,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         del start_embedding
         del end_embedding
         torch.cuda.empty_cache()
-       
-        if not self.training:
-            # Inference mode
-            if max_length is None:
-                max_length = self.config.max_length
-
-            finished_generating = torch.zeros(batch_size, dtype=torch.bool, device=input_ids.device)
-            for cur_token_idx in range(max_length):
-                outputs = self.model(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    past_key_values=past_key_values,
-                    inputs_embeds=inputs_embeds,
-                    use_cache=use_cache,
-                    output_attentions=output_attentions,
-                    output_hidden_states=output_hidden_states,
-                    return_dict=return_dict,
-                )
-                hidden_states = outputs[0]
-                logits = self.lm_head(hidden_states)
-
-                # Mask out the start and end thought tokens
-                logits[:, :, self.start_token_id] = -float("inf")
-                logits[:, :, self.end_token_id] = -float("inf")
-
-                for batch_idx in range(batch_size):
-                    if not finished_generating[batch_idx]:
-                        last_token_idx = (input_ids[batch_idx] != self.tokenizer.pad_token_id).nonzero(as_tuple=True)[0].max()
-                        new_id_sampled = torch.multinomial(
-                            torch.nn.functional.softmax(logits[batch_idx, last_token_idx] / temperature, dim=-1), 1
-                        )
-                        if last_token_idx + 1 >= input_ids.shape[1]:
-                            # Add padding
-                            new_padding = torch.full((batch_size, 1), self.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[batch_idx, last_token_idx + 1] = 1
-                        input_ids[batch_idx, last_token_idx + 1] = new_id_sampled
-                        if new_id_sampled == self.tokenizer.eos_token_id or new_id_sampled == self.tokenizer.bos_token_id or new_id_sampled == self.tokenizer.pad_token_id:
-                            finished_generating[batch_idx] = True
-
-                if finished_generating.all():
-                    break
+        
 
         return CausalLMOutputWithPast(
             loss=loss if loss is not None else None,