fix import

myolmoe/__init__.py

@@ -1 +1,13 @@
-from .modeling_myolmoe import MyOlmoeForCausalLM
+from .modeling_myolmoe import (
+    MyOlmoeForCausalLM,
+    OlmoeModel,
+    OlmoePreTrainedModel,
+    OlmoeConfig  # Add this line
+)
+
+__all__ = [
+    "MyOlmoeForCausalLM",
+    "OlmoeModel", 
+    "OlmoePreTrainedModel",
+    "OlmoeConfig"  # Add this line
+]

scripts/train.py

@@ -12,8 +12,14 @@ from myolmoe import MyOlmoeForCausalLM, OlmoeConfig
 import os
 
 def main():
-    # Load config and model
-    config = OlmoeConfig.from_pretrained("myolmoe/config.json")
+    # Load config - first try from local file, then from pretrained
+    config_path = os.path.join("myolmoe", "config.json")
+    if os.path.exists(config_path):
+        config = OlmoeConfig.from_json_file(config_path)
+    else:
+        config = OlmoeConfig.from_pretrained("myolmoe")
+    
+    # Load model
     model = MyOlmoeForCausalLM.from_pretrained(
         "myolmoe",
         config=config,