Update generate.py

generate.py

@@ -52,7 +52,6 @@ def custom_generate(
             new_ids = self(
                 input_ids[~finished_generating],
                 attention_mask=attention_mask[~finished_generating] if attention_mask is not None else None,
-                streamer=streamer,
                 **kwargs
             )['logits']
 
@@ -88,6 +87,9 @@ def custom_generate(
                 if new_ids_sampled == self.tokenizer.convert_tokens_to_ids("<|/assistant|>"):
                     finished_generating[answer_idx] = 1
 
+            if streamer is not None:
+                streamer.put(new_ids_sampled)
+
         generated_token_ids = input_ids.tolist()
     
         return generated_token_ids
@@ -165,12 +167,6 @@ def generate(
     self.rm_initialized = True
     self.original_mode = False
 
-    # # Validate stopping criteria
-    # stopping_criteria = self._get_stopping_criteria(generation_config=self.config, stopping_criteria=StoppingCriteriaList())
-    # if stopping_criteria is None:
-    #     stopping_criteria = StoppingCriteriaList()
-    # if max_length is not None:
-    #     stopping_criteria = validate_stopping_criteria(max_length, stopping_criteria=stopping_criteria)
     # Generate using the custom generate function
     generated_token_ids = custom_generate(
         self,
@@ -209,4 +205,4 @@ def generate(
         **model_kwargs,
     )
 
-    return generated_token_ids
+    return generated_token_ids