HF integration

config.json

@@ -8,6 +8,10 @@
   "attention_dropout": 0.0,
   "attention_layer_norm": true,
   "attention_layer_norm_with_affine": true,
+  "auto_map": {
+    "AutoConfig": "configuration_olmo.OLMoConfig",
+    "AutoModelForCausalLM": "modeling_fan.OLMoForCausalLM"
+  },
   "bias_for_layer_norm": false,
   "block_group_size": 1,
   "block_type": "sequential",

configuration_olmo.py

@@ -1 +1,4 @@
+from transformers import AutoConfig
 from hf_olmo import OLMoConfig
+
+AutoConfig.register("hf_olmo", OLMoConfig)

modeling_olmo.py

@@ -1 +1,4 @@
-from hf_olmo import OLMoForCausalLM
+from hf_olmo import OLMoForCausalLM, OLMoConfig
+from transformers import AutoModel
+
+AutoModel.register(OLMoConfig, OLMoForCausalLM)

tokenization_olmo_fast.py

@@ -1 +1,4 @@
-from hf_olmo.tokenization_olmo_fast import OLMoTokenizerFast
+from hf_olmo.tokenization_olmo_fast import OLMoTokenizerFast, OLMoConfig
+from transformers import AutoTokenizer
+
+AutoTokenizer.register(OLMoConfig, fast_tokenizer_class=OLMoTokenizerFast)