Machine-Learning
/
AIIA
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

develop #41

Merged
Fabel merged 27 commits from develop into main 2025-04-17 17:08:57 +00:00
Conversation 0 Commits 27 Files Changed 13 +15 -15
1 changed files with 15 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 9ec01e2dd0 - Show all commits

30
tests/model/test_aiia.py
View File

@ -12,7 +12,7 @@ def test_aiiabase_save_pretrained_from_pretrained():
model = AIIABase(config)
save_pretrained_path = "test_aiiabase_save_pretrained_load"
# save_pretrained the model
# Save the model
model.save_pretrained(save_pretrained_path)
assert os.path.exists(os.path.join(save_pretrained_path, "model.safetensors"))
assert os.path.exists(os.path.join(save_pretrained_path, "config.json"))
@ -38,7 +38,7 @@ def test_aiiabase_shared_save_pretrained_from_pretrained():
model = AIIABaseShared(config)
save_pretrained_path = "test_aiiabase_shared_save_pretrained_load"
# save_pretrained the model
# Save the model
model.save_pretrained(save_pretrained_path)
assert os.path.exists(os.path.join(save_pretrained_path, "model.safetensors"))
assert os.path.exists(os.path.join(save_pretrained_path, "config.json"))
@ -64,7 +64,7 @@ def test_aiiaexpert_save_pretrained_from_pretrained():
model = AIIAExpert(config)
save_pretrained_path = "test_aiiaexpert_save_pretrained_load"
# save_pretrained the model
# Save the model
model.save_pretrained(save_pretrained_path)
assert os.path.exists(os.path.join(save_pretrained_path, "model.safetensors"))
assert os.path.exists(os.path.join(save_pretrained_path, "config.json"))
@ -81,16 +81,16 @@ def test_aiiaexpert_save_pretrained_from_pretrained():
os.rmdir(save_pretrained_path)
def test_aiiamoe_creation():
config = AIIAConfig()
model = AIIAmoe(config, num_experts=5)
config = AIIAConfig(num_experts=3)
model = AIIAmoe(config)
assert isinstance(model, AIIAmoe)
def test_aiiamoe_save_pretrained_from_pretrained():
config = AIIAConfig()
model = AIIAmoe(config, num_experts=5)
config = AIIAConfig(num_experts=3)
model = AIIAmoe(config)
save_pretrained_path = "test_aiiamoe_save_pretrained_load"
# save_pretrained the model
# Save the model
model.save_pretrained(save_pretrained_path)
assert os.path.exists(os.path.join(save_pretrained_path, "model.safetensors"))
assert os.path.exists(os.path.join(save_pretrained_path, "config.json"))
@ -107,16 +107,16 @@ def test_aiiamoe_save_pretrained_from_pretrained():
os.rmdir(save_pretrained_path)
def test_aiiasparsemoe_creation():
config = AIIAConfig()
model = AIIASparseMoe(config, num_experts=5, top_k=2)
config = AIIAConfig(num_experts=5, top_k=2)
model = AIIASparseMoe(config, base_class=AIIABaseShared)
assert isinstance(model, AIIASparseMoe)
def test_aiiasparsemoe_save_pretrained_from_pretrained():
config = AIIAConfig()
model = AIIASparseMoe(config, num_experts=3, top_k=1)
config = AIIAConfig(num_experts=3, top_k=1)
model = AIIASparseMoe(config)
save_pretrained_path = "test_aiiasparsemoe_save_pretrained_load"
# save_pretrained the model
# Save the model
model.save_pretrained(save_pretrained_path)
assert os.path.exists(os.path.join(save_pretrained_path, "model.safetensors"))
assert os.path.exists(os.path.join(save_pretrained_path, "config.json"))
@ -130,4 +130,4 @@ def test_aiiasparsemoe_save_pretrained_from_pretrained():
# Clean up
os.remove(os.path.join(save_pretrained_path, "model.safetensors"))
os.remove(os.path.join(save_pretrained_path, "config.json"))
os.rmdir(save_pretrained_path)
os.rmdir(save_pretrained_path)