feat/tf_support #37

Merged
Fabel merged 13 commits from feat/tf_support into develop 2025-04-16 20:59:48 +00:00
1 changed files with 7 additions and 5 deletions
Showing only changes of commit ba6da9ef02 - Show all commits

View File

@ -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)