From ba6da9ef020498a84a35de362249dc13c6e7bd84 Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Sun, 13 Apr 2025 22:17:46 +0200
Subject: [PATCH] restored full example

---
 example.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/example.py b/example.py
index 8dbb67e..2ce0b6a 100644
--- a/example.py
+++ b/example.py
@@ -1,10 +1,12 @@
-from aiia.model import AIIABase
-from aiia.model import AIIAConfig
-from aiia.pretrain import Pretrainer
+from src.aiia.model import AIIAmoe
+from src.aiia.model import AIIAConfig
+from src.aiia.pretrain import Pretrainer
 
 # Create your model
-config = AIIAConfig(model_name="AIIA-Base-512x20k")
-model = AIIABase(config)
+config = AIIAConfig(num_experts=5)
+model = AIIAmoe(config)
+model.save_pretrained("test")
+model = AIIAmoe.from_pretrained("test")
 
 # Initialize pretrainer with the model
 pretrainer = Pretrainer(model, learning_rate=1e-4, config=config)