Update generate.py

generate.py

@@ -1,10 +1,4 @@
 import torch
-from transformers.generation.utils import (
-    GenerationMixin,
-    validate_stopping_criteria,
-    StoppingCriteriaList,
-)
-from transformers import TextStreamer
 
 
 def custom_generate(
@@ -50,7 +44,6 @@ def custom_generate(
         if max_new_tokens is None:
             max_new_tokens = 50  # Default value if not specified
         for cur_token_idx in range(max_new_tokens):
-            # Sample the next token
             new_ids = self(
                 input_ids[~finished_generating],
                 attention_mask=attention_mask[~finished_generating] if attention_mask is not None else None,
@@ -95,14 +88,12 @@ def custom_generate(
             if streamer is not None:
                 streamer.put(new_ids_sampled)
 
-    # Create a named tuple to match the expected output format
     from collections import namedtuple
     GenerateOutput = namedtuple("GenerateOutput", ["sequences", "scores", "attentions", "hidden_states"])
     
     # Convert the generated token IDs to a tensor
     generated_token_ids_tensor = input_ids
     
-    # Create the GenerateOutput named tuple
     output = GenerateOutput(
         sequences=generated_token_ids_tensor,
         scores=None,