added new models

Author: noanabeshima

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tinymodel"
-version = "0.1.1-9"
+version = "0.1.2-1"
 description = "A small TinyStories LM with SAEs and transcoders"
 authors = ["Noa Nabeshima <noanabeshima@gmail.com>"]
 readme = "README.md"

tiny_model/lm.py

@@ -249,13 +249,11 @@ def __getitem__(self, mlp_tag):
 
 
 def get_state_dict(model_fname="tiny_model"):
-    """
-    There are two models available: `tiny_model` and `tiny_model_1_epoch`.
-    """
     assert model_fname in [
         "tiny_model",
-        "tiny_model_1_epoch",
-    ], "There are two models available: `tiny_model` and `tiny_model_1_epoch`."
+        "tiny_model_2L_1E",
+        "tiny_model_2L_3E"
+    ], "There are 3 models available: `tiny_model`, `tiny_model_2L_1E`, and `tiny_model_2L_3E`."
     state_dict = torch.load(
         hf_hub_download(repo_id="noanabeshima/tiny_model", filename=f"{model_fname}.pt"),
         map_location=torch.device('cpu')