Tweak the GPTNeoConfigCPU config

minimaxir · Apr 18, 2021 · 791e0e1 · 791e0e1
1 parent f2ef5fd
commit 791e0e1
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/aitextgen/utils.py b/aitextgen/utils.py
@@ -149,7 +149,10 @@ def GPTNeoConfigCPU(
         vocab_size=vocab_size,
         max_position_embeddings=64,
         hidden_size=256,
-        attention_types=[[["global", "local"], 4]],
+        window_size=32,
+        intermediate_size=256,
+        attention_types=[[["global", "local"], 2]],
+        num_layers=4,
         num_heads=4,
         bos_token_id=bos_token_id,
         eos_token_id=eos_token_id,