Run formatter

allenai · Nov 26, 2024 · 8239b05 · 8239b05
1 parent 6d49aff
commit 8239b05
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/scripts/convert_olmo2_to_hf.py b/scripts/convert_olmo2_to_hf.py
@@ -22,11 +22,9 @@
 import torch
 import yaml
 from tokenizers import Tokenizer
-
 from transformers import Olmo2Config, Olmo2ForCausalLM
 from transformers.models.gpt2.tokenization_gpt2_fast import GPT2TokenizerFast
 
-
 """
 Sample usage:
 
@@ -125,9 +123,15 @@ def write_model(
             f"model.layers.{layer_i}.self_attn.q_proj.weight": q_proj_weight,
             f"model.layers.{layer_i}.self_attn.k_proj.weight": k_proj_weight,
             f"model.layers.{layer_i}.self_attn.v_proj.weight": v_proj_weight,
-            f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"transformer.blocks.{layer_i}.attn_out.weight"],
-            f"model.layers.{layer_i}.self_attn.q_norm.weight": loaded[f"transformer.blocks.{layer_i}.q_norm.weight"],
-            f"model.layers.{layer_i}.self_attn.k_norm.weight": loaded[f"transformer.blocks.{layer_i}.k_norm.weight"],
+            f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[
+                f"transformer.blocks.{layer_i}.attn_out.weight"
+            ],
+            f"model.layers.{layer_i}.self_attn.q_norm.weight": loaded[
+                f"transformer.blocks.{layer_i}.q_norm.weight"
+            ],
+            f"model.layers.{layer_i}.self_attn.k_norm.weight": loaded[
+                f"transformer.blocks.{layer_i}.k_norm.weight"
+            ],
             f"model.layers.{layer_i}.mlp.gate_proj.weight": gate_proj_weight,
             f"model.layers.{layer_i}.mlp.down_proj.weight": loaded[f"transformer.blocks.{layer_i}.ff_out.weight"],
             f"model.layers.{layer_i}.mlp.up_proj.weight": up_proj_weight,