Update generate.py

generate.py

@@ -103,10 +103,7 @@ def custom_generate(
             if streamer is not None:
                 streamer.put(new_ids_sampled)
 
-        # Convert generated token IDs to text
-        generated_text = self.tokenizer.decode(generated_token_ids[0], skip_special_tokens=False)
-
-        return generated_token_ids, generated_text
+        return generated_token_ids
 
 def generate(
     self,
@@ -193,7 +190,7 @@ def generate(
     if attention_mask is not None:
         attention_mask = attention_mask.to(self.device)
 
-    generated_token_ids, generated_text = custom_generate(
+    generated_token_ids = custom_generate(
         self,
         input_ids=input_ids,
         attention_mask=attention_mask,
@@ -230,4 +227,7 @@ def generate(
         **model_kwargs,
     )
 
-    return generated_token_ids, generated_text
+    # Convert the generated token IDs tensor to text
+    generated_text = self.tokenizer.decode(generated_token_ids[0], skip_special_tokens=False)
+
+    return generated_token_ids, generated_text