Update generate.py

generate.py

@@ -86,10 +86,10 @@ class LmGeneration:
         total_len = args.seq_length
 
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        tokens = torch.full((batch, total_len), self.tokenizer.pad_token_id).to(device).long()
+        tokens = torch.full((batch, total_len), self.tokenizer.pad_token).to(device).long()
         for idx, t in enumerate(prompt_tokens):
             tokens[idx, : len(t)] = torch.tensor(t).long()
-        mask = tokens != self.tokenizer.pad_token_id
+        mask = tokens != self.tokenizer.pad_token
         start_pos = min_prompt_len
         prev_pos = 0
         continue_exsample = [i for i in range(batch)]
@@ -118,7 +118,7 @@ class LmGeneration:
                 continue_exsample = []
                 for i, t in enumerate(tokens.tolist()):
                     try:
-                        t.index(self.tokenizer.eos_token_id)
+                        t.index(self.tokenizer.eos_token)
                     except ValueError:
                         if cut_off is not None:
                             if cut_off == self.tokenizer.decode(t[:cur_pos + 1])[-len(cut_off):]:
@@ -134,8 +134,8 @@ class LmGeneration:
         for i, t in enumerate(tokens.tolist()):
             t = t[: args.seq_length]
             try:
-                t = t[: t.index(self.tokenizer.pad_token_id)]
-                t = t[: t.index(self.tokenizer.eos_token_id)]
+                t = t[: t.index(self.tokenizer.pad_token)]
+                t = t[: t.index(self.tokenizer.eos_token)]
             except ValueError:
                 pass
             decoder.append(self.tokenizer.decode(t))