automodel and autotok

README.md

@@ -33,11 +33,10 @@ Train loss: ~2.18 · Dev loss: ~2.20
 ```python
 import torch
 import torch.nn.functional as F
-from transformers import AutoTokenizer
-from modeling_makemore import MakemoreMLP
+from transformers import AutoModel, AutoTokenizer
 
-model = MakemoreMLP.from_pretrained("iand666/makemore-mlp", trust_remote_code=True)
-tok = AutoTokenizer.from_pretrained("iand666/makemore-mlp", trust_remote_code=True)
+model = AutoModel.from_pretrained("iand666/makemore-mlp", trust_remote_code=True)
+tok   = AutoTokenizer.from_pretrained("iand666/makemore-mlp", trust_remote_code=True)
 model.eval()
 
 def gen_name(model, tok, block_size=3):

config.json

@@ -4,5 +4,8 @@
   "emb_dim": 10,
   "hidden_dim": 200,
   "vocab_size": 27,
-  "tokenizer_class": "MakemoreTokenizer"
+  "tokenizer_class": "MakemoreTokenizer",
+  "auto_map": {
+    "AutoModel": "modeling_makemore.MakemoreMLP"
+  }
 }

tokenization_makemore.py

@@ -1,5 +1,9 @@
 from transformers import PreTrainedTokenizer
 
+AUTO_MAP = {
+    "AutoTokenizer": ["tokenization_makemore.MakemoreTokenizer", None]
+}
+
 
 class MakemoreTokenizer(PreTrainedTokenizer):
     def __init__(self, **kwargs):