Update generate.py

generate.py

@@ -45,7 +45,9 @@ def custom_generate(
 ):
     device = input_ids.device
     with torch.no_grad():
-        finished_generating = torch.zeros(len(input_ids), dtype=torch.bool, device=device)
+        batch_size = input_ids.shape[0]
+        finished_generating = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        generated_token_ids = torch.full((batch_size, max_new_tokens), self.tokenizer.pad_token_id, dtype=torch.long, device=device)
 
         for cur_token_idx in range(max_new_tokens):
             # Sample the next token
@@ -70,7 +72,7 @@ def custom_generate(
                 # Assign the new id to the last token
                 if last_token_idx + 1 >= len(base_answer_ids):
                     # Add padding everywhere
-                    new_padding = torch.full((len(input_ids), 1), self.tokenizer.pad_token_id, dtype=torch.long,
+                    new_padding = torch.full((batch_size, 1), self.tokenizer.pad_token_id, dtype=torch.long,
                                             device=device)
                     input_ids = torch.cat([input_ids, new_padding], dim=-1)
                     if attention_mask is not None:
@@ -79,6 +81,7 @@ def custom_generate(
                 if attention_mask is not None:
                     attention_mask[answer_idx, last_token_idx + 1] = 1
                 input_ids[answer_idx, last_token_idx + 1] = new_ids_sampled
+                generated_token_ids[answer_idx, cur_token_idx] = new_ids_sampled
 
                 if new_ids_sampled == self.tokenizer.eos_token_id or new_ids_sampled == self.tokenizer.bos_token_id or new_ids_sampled == self.tokenizer.pad_token_id:
                     finished_generating[answer_idx] = 1
@@ -93,7 +96,7 @@ def custom_generate(
             if streamer is not None:
                 streamer.put(new_ids_sampled)
 
-        return input_ids
+        return generated_token_ids
 
 
 def generate(