Update generate.py

generate.py

@@ -95,8 +95,22 @@ def custom_generate(
             if streamer is not None:
                 streamer.put(new_ids_sampled)
 
-        generated_token_ids = input_ids.tolist()
-        return generated_token_ids, attention_mask
+    # 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,
+        attentions=None,
+        hidden_states=None
+    )
+    
+    return output
 
 
 def generate(