Update generate.py

generate.py

@@ -56,9 +56,12 @@ def custom_generate(
     with torch.no_grad():
         batch_size = input_ids.shape[0]
         finished_generating = torch.zeros(batch_size, dtype=torch.bool, device=device)
+
+        max_length = input_ids.shape[1] + max_new_tokens
         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):
+        cur_token_idx = 0
+        while cur_token_idx < max_length:
             # Sample the next token
             new_ids = self(
                 input_ids[~finished_generating],
@@ -96,13 +99,11 @@ def custom_generate(
                 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:
+                if new_ids_sampled == self.tokenizer.eos_token_id or cur_token_idx + 1 == max_length:
                     finished_generating[answer_idx] = 1
 
-                # Check if the end token is generated
-                if new_ids_sampled == self.tokenizer.convert_tokens_to_ids("</s>"):
-                    finished_generating[answer_idx] = 1
-            
+            cur_token_idx += 1
+
             if finished_generating.all():
                 break
 
@@ -112,7 +113,13 @@ def custom_generate(
         print("Generated Token IDs shape:", generated_token_ids.shape)
         print("Generated Token IDs:", generated_token_ids)
 
-        return generated_token_ids
+        # Decode the generated token IDs into text
+        generated_text = self.tokenizer.decode(generated_token_ids[0], skip_special_tokens=True)
+        print("Generated Text:")
+        print(generated_text)
+
+        return generated_token_ids, generated_text
+
 
 def generate(
     self,
@@ -161,6 +168,7 @@ def generate(
     use_weighted_talk_head=True,
     trust_remote_code=True,
     torch_dtype=torch.bfloat16,
+    dynamic_temperature=None,
     **model_kwargs,
 ):
 
@@ -198,7 +206,7 @@ def generate(
     if attention_mask is not None:
         attention_mask = attention_mask.to(self.device)
 
-    generated_token_ids = custom_generate(
+    generated_token_ids, generated_text = custom_generate(
         self,
         input_ids=input_ids, 
         attention_mask=attention_mask,
@@ -235,4 +243,7 @@ def generate(
         **model_kwargs,
     )
 
-    return generated_token_ids
+    if dynamic_temperature is not None:
+        return generated_text
+    else:
+        return generated_token_ids