forked from luopingyi/xc_agent_platform_demo
37 lines
1.7 KiB
Python
37 lines
1.7 KiB
Python
|
import unittest
|
||
|
|
|
||
|
|
from src.banned import banned_rule
|
||
|
|
|
||
|
|
|
||
|
|
class TestBanned(unittest.TestCase):
|
||
|
|
def test_default_engine_and_openchat(self):
|
||
|
|
for model, rule in [('org/Model-GGUF', 'gguf'), ('org/model.gguf', 'gguf'),
|
||
|
|
('org/GGML-model', 'gguf'), ('OpenChat/model', 'openchat')]:
|
||
|
|
with self.subTest(model=model):
|
||
|
|
self.assertEqual(banned_rule(model), rule)
|
||
|
|
|
||
|
|
def test_unmerged_adapter_hints(self):
|
||
|
|
for model in ('org/model-lora', 'org/LORA', 'org/adapter_model',
|
||
|
|
'org/model-adapter-config', 'org/Model-PEFT'):
|
||
|
|
with self.subTest(model=model):
|
||
|
|
self.assertEqual(banned_rule(model), 'peft-lora')
|
||
|
|
|
||
|
|
def test_full_weights_and_unrelated_names(self):
|
||
|
|
for model in ('org/lora-merged', 'org/merge-lora', 'org/PEFT-FULLFT',
|
||
|
|
'org/full-ft-adapter_model', 'org/adapter-config-MERGED',
|
||
|
|
'org/flora', 'org/lorax', 'org/model_lora',
|
||
|
|
'org/model-GRPO', 'abaryan/CyberXP_Agent_Llama_3.2_1B'):
|
||
|
|
with self.subTest(model=model):
|
||
|
|
self.assertIsNone(banned_rule(model))
|
||
|
|
|
||
|
|
def test_full_weight_hint_does_not_override_other_bans(self):
|
||
|
|
self.assertEqual(banned_rule('org/lora-merged-GGUF'), 'gguf')
|
||
|
|
self.assertEqual(banned_rule('org/openchat-fullft'), 'openchat')
|
||
|
|
|
||
|
|
def test_existing_rules(self):
|
||
|
|
for model in ('gpt2', 'org/Phi-3-mini', 'nvidia/Nemotron-H-4B',
|
||
|
|
'allenai/OLMo-2-7B', 'org/LFM2', 'org/Spark-X2.5',
|
||
|
|
'org/Breeze', 'org/gemma2', 'org/mistral-v0.1', 'org/tinyllama'):
|
||
|
|
with self.subTest(model=model):
|
||
|
|
self.assertIsNotNone(banned_rule(model))
|