fix attention mask

Haskely · Aug 10, 2023 · dd8e714 · dd8e714
1 parent 2d12248
commit dd8e714
Show file tree

Hide file tree

Showing 4 changed files with 17,156 additions and 4 deletions.
diff --git a/llama_gen_and_eval.py b/llama_gen_and_eval.py
@@ -97,21 +97,22 @@ def gsm8k_batch_gen(
 def get_batch_llama(model: LlamaForCausalLM, tokenizer: LlamaTokenizer):
     @torch.inference_mode()
     def batch_llama(input_strs: list[str]) -> list[str]:
-        input_ids = tokenizer(
+        input_ids_w_attnmask = tokenizer(
             input_strs,
             padding=True,
             return_tensors="pt",
-        ).input_ids.to(model.device)
+        ).to(model.device)
         output_ids = model.generate(
-            inputs=input_ids,
+            input_ids=input_ids_w_attnmask.input_ids,
+            attention_mask=input_ids_w_attnmask.attention_mask,
             generation_config=GenerationConfig(
                 max_length=512,
                 do_sample=False,
                 temperature=0.0,  # t=0.0 raise error if do_sample=True
             ),
         ).tolist()
         real_output_ids = [
-            output_id[len(input_ids[i]) :] for i, output_id in enumerate(output_ids)
+            output_id[len(input_ids_w_attnmask.input_ids[i]) :] for i, output_id in enumerate(output_ids)
         ]
         output_strs = tokenizer.batch_decode(real_output_ids, skip_special_tokens=True)
         return output_strs