diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py index c59c7421..c68c6d5e 100644 --- a/tests/base/model_test_registry.py +++ b/tests/base/model_test_registry.py @@ -2356,6 +2356,90 @@ "input_factory_kwargs": {"seq_len": 16}, "expected_output_shape": {"start_logits": (2, 16), "end_logits": (2, 16)}, }, + "BartModel": { + "module": "zeromodels.models.bart", + "model_cls": "BartModel", + "model_type": "llm", + "init_kwargs": { + "vocab_size": 128, + "hidden_dim": 32, + "encoder_num_layers": 2, + "decoder_num_layers": 2, + "encoder_attention_heads": 4, + "decoder_attention_heads": 4, + "encoder_ffn_dim": 64, + "decoder_ffn_dim": 64, + "max_position_embeddings": 32, + }, + "input_factory": "t5_input", + "input_factory_kwargs": {"seq_len": 16, "decoder_seq_len": 8}, + "expected_output_shape": { + "last_hidden_state": (2, 8, 32), + "encoder_last_hidden_state": (2, 16, 32), + }, + }, + "BartConditionalGenerate": { + "module": "zeromodels.models.bart", + "model_cls": "BartConditionalGenerate", + "model_type": "llm", + "init_kwargs": { + "vocab_size": 128, + "hidden_dim": 32, + "encoder_num_layers": 2, + "decoder_num_layers": 2, + "encoder_attention_heads": 4, + "decoder_attention_heads": 4, + "encoder_ffn_dim": 64, + "decoder_ffn_dim": 64, + "max_position_embeddings": 32, + }, + "input_factory": "t5_input", + "input_factory_kwargs": {"seq_len": 16, "decoder_seq_len": 8}, + "expected_output_shape": { + "last_hidden_state": (2, 8, 32), + "encoder_last_hidden_state": (2, 16, 32), + "logits": (2, 8, 128), + }, + }, + "BartSequenceClassify": { + "module": "zeromodels.models.bart", + "model_cls": "BartSequenceClassify", + "model_type": "llm", + "init_kwargs": { + "vocab_size": 128, + "hidden_dim": 32, + "encoder_num_layers": 2, + "decoder_num_layers": 2, + "encoder_attention_heads": 4, + "decoder_attention_heads": 4, + "encoder_ffn_dim": 64, + "decoder_ffn_dim": 64, + "max_position_embeddings": 32, + "num_labels": 3, + }, + "input_factory": "t5_encoder_input", + "input_factory_kwargs": {"seq_len": 16}, + "expected_output_shape": (2, 3), + }, + "BartQnA": { + "module": "zeromodels.models.bart", + "model_cls": "BartQnA", + "model_type": "llm", + "init_kwargs": { + "vocab_size": 128, + "hidden_dim": 32, + "encoder_num_layers": 2, + "decoder_num_layers": 2, + "encoder_attention_heads": 4, + "decoder_attention_heads": 4, + "encoder_ffn_dim": 64, + "decoder_ffn_dim": 64, + "max_position_embeddings": 32, + }, + "input_factory": "t5_encoder_input", + "input_factory_kwargs": {"seq_len": 16}, + "expected_output_shape": {"start_logits": (2, 16), "end_logits": (2, 16)}, + }, "BertModel": { "module": "zeromodels.models.bert", "model_cls": "BertModel", diff --git a/tests/fixtures/cross_backend_parity.json b/tests/fixtures/cross_backend_parity.json index 70e27c96..c9be97b1 100644 --- a/tests/fixtures/cross_backend_parity.json +++ b/tests/fixtures/cross_backend_parity.json @@ -1,4 +1,471 @@ { + "BartConditionalGenerate": [ + { + "shape": [ + 2, + 16, + 32 + ], + "sample": [ + -0.045629, + -0.015995, + -0.04584, + -0.016489, + -0.046488, + 0.007977, + 0.08007, + 0.006455, + 0.077616, + 0.01086, + 0.040044, + -0.002418, + 0.039532, + 8.6e-05, + -0.005188, + -0.001721, + -0.004609, + -0.004943, + -0.014366, + -0.005109, + -0.014451, + 0.004622, + -0.009322, + 0.004263, + -0.012089, + 0.004254, + 0.00167, + -0.02902, + 0.013354, + -0.029042, + -0.018903, + 0.003579, + -0.01923, + 0.003615, + -0.031424, + 0.014926, + -0.031926, + 0.016058, + -0.014396, + -0.003727, + -0.014458, + -0.003636, + 0.003686, + 0.020906, + 0.006667, + 0.021055, + 0.007164, + -0.06098, + 0.002818, + -0.060315, + 0.002815, + 0.014172, + -0.005455, + 0.015238, + -0.005348, + -0.000647, + -0.014952, + -0.000546, + -0.015215, + -0.008536, + 0.016109, + -0.008881, + 0.016623, + -0.008043 + ] + }, + { + "shape": [ + 2, + 8, + 32 + ], + "sample": [ + -0.003057, + -0.032961, + 0.007522, + -0.002054, + -0.003044, + -0.032961, + 0.00752, + -0.002026, + -0.003054, + -0.034907, + 0.026296, + -0.042875, + -0.040833, + -0.035029, + 0.026275, + -0.042822, + -0.040937, + -0.035068, + 0.030709, + -0.034514, + -0.065917, + 0.00501, + 0.030799, + -0.034461, + -0.065893, + 0.00502, + 0.030487, + 0.004971, + -0.036882, + -0.024988, + -0.009892, + 0.004953, + -0.037139, + -0.02531, + -0.009921, + 0.004986, + 0.033659, + -0.006406, + -0.015556, + -0.000915, + 0.032601, + -0.006399, + -0.015502, + -0.000936, + 0.032995, + -0.006932, + -0.020991, + 0.019133, + -0.015691, + -0.006929, + -0.020998, + 0.019234, + -0.015538, + -0.006927, + 0.000436, + -0.070966, + 0.019877, + -0.025448, + 0.00028, + -0.068742, + 0.019893, + -0.025393, + 0.000336, + -0.04234 + ] + }, + { + "shape": [ + 2, + 8, + 128 + ], + "sample": [ + 0.017134, + -0.00781, + 0.000839, + 0.005244, + 0.000137, + -0.005838, + -0.014909, + 0.016994, + 0.005994, + -0.028245, + -0.013149, + 0.010292, + -0.000747, + -0.006892, + 0.024968, + 0.004003, + 0.012783, + 0.007287, + 0.012294, + -0.008029, + 0.037378, + -0.020914, + 0.035682, + 0.019984, + -0.002951, + 0.025353, + 0.013376, + -0.029729, + -0.015842, + -4.9e-05, + 0.032712, + -0.027618, + -0.008007, + 0.011081, + -0.011713, + -0.005795, + 0.009777, + -0.019756, + -0.035538, + -0.011392, + -0.002458, + -0.004896, + 0.006239, + 0.01323, + -0.029386, + -0.006449, + -0.025356, + 0.035827, + 0.01171, + -0.011475, + 0.007274, + -0.016466, + 0.007971, + 0.010488, + -0.00854, + -0.021537, + -0.023269, + -0.003491, + -0.006641, + 0.017053, + 0.041385, + 0.002612, + -0.030184, + 0.021758 + ] + } + ], + "BartModel": [ + { + "shape": [ + 2, + 16, + 32 + ], + "sample": [ + -0.045629, + -0.015995, + -0.04584, + -0.016489, + -0.046488, + 0.007977, + 0.08007, + 0.006455, + 0.077616, + 0.01086, + 0.040044, + -0.002418, + 0.039532, + 8.6e-05, + -0.005188, + -0.001721, + -0.004609, + -0.004943, + -0.014366, + -0.005109, + -0.014451, + 0.004622, + -0.009322, + 0.004263, + -0.012089, + 0.004254, + 0.00167, + -0.02902, + 0.013354, + -0.029042, + -0.018903, + 0.003579, + -0.01923, + 0.003615, + -0.031424, + 0.014926, + -0.031926, + 0.016058, + -0.014396, + -0.003727, + -0.014458, + -0.003636, + 0.003686, + 0.020906, + 0.006667, + 0.021055, + 0.007164, + -0.06098, + 0.002818, + -0.060315, + 0.002815, + 0.014172, + -0.005455, + 0.015238, + -0.005348, + -0.000647, + -0.014952, + -0.000546, + -0.015215, + -0.008536, + 0.016109, + -0.008881, + 0.016623, + -0.008043 + ] + }, + { + "shape": [ + 2, + 8, + 32 + ], + "sample": [ + -0.003057, + -0.032961, + 0.007522, + -0.002054, + -0.003044, + -0.032961, + 0.00752, + -0.002026, + -0.003054, + -0.034907, + 0.026296, + -0.042875, + -0.040833, + -0.035029, + 0.026275, + -0.042822, + -0.040937, + -0.035068, + 0.030709, + -0.034514, + -0.065917, + 0.00501, + 0.030799, + -0.034461, + -0.065893, + 0.00502, + 0.030487, + 0.004971, + -0.036882, + -0.024988, + -0.009892, + 0.004953, + -0.037139, + -0.02531, + -0.009921, + 0.004986, + 0.033659, + -0.006406, + -0.015556, + -0.000915, + 0.032601, + -0.006399, + -0.015502, + -0.000936, + 0.032995, + -0.006932, + -0.020991, + 0.019133, + -0.015691, + -0.006929, + -0.020998, + 0.019234, + -0.015538, + -0.006927, + 0.000436, + -0.070966, + 0.019877, + -0.025448, + 0.00028, + -0.068742, + 0.019893, + -0.025393, + 0.000336, + -0.04234 + ] + } + ], + "BartQnA": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + -0.026378, + -0.026354, + -0.026368, + -0.026356, + -0.026373, + -0.026342, + -0.02633, + -0.026312, + -0.026355, + -0.026359, + -0.026343, + -0.026371, + -0.026366, + -0.026391, + -0.02638, + -0.026337, + -0.026379, + -0.026345, + -0.026348, + -0.02637, + -0.026363, + -0.026382, + -0.026368, + -0.026316, + -0.026337, + -0.026352, + -0.026355, + -0.02636, + -0.026342, + -0.026362, + -0.026365, + -0.026347 + ] + }, + { + "shape": [ + 2, + 16 + ], + "sample": [ + 0.019905, + 0.019923, + 0.019907, + 0.019921, + 0.019927, + 0.019932, + 0.01992, + 0.019927, + 0.019935, + 0.019911, + 0.019958, + 0.019945, + 0.019925, + 0.019943, + 0.019943, + 0.019925, + 0.019905, + 0.019927, + 0.01993, + 0.019928, + 0.019941, + 0.019939, + 0.019929, + 0.019922, + 0.019937, + 0.019915, + 0.019958, + 0.019922, + 0.019904, + 0.019924, + 0.019948, + 0.019911 + ] + } + ], + "BartSequenceClassify": [ + { + "shape": [ + 2, + 3 + ], + "sample": [ + 0.003965, + -0.002365, + 0.03728, + 0.00396, + -0.002363, + 0.037283 + ] + } + ], "BeitImageClassify": [ { "shape": [ diff --git a/zeromodels/auto/auto_mapping_names.py b/zeromodels/auto/auto_mapping_names.py index 514fd155..c3ec9f08 100644 --- a/zeromodels/auto/auto_mapping_names.py +++ b/zeromodels/auto/auto_mapping_names.py @@ -21,6 +21,7 @@ "granite_speech5_ctc": "GraniteSpeech5CTC", }, "ConditionalGenerate": { + "bart": "BartConditionalGenerate", "deepseek_vl": "DeepseekVLConditionalGenerate", "deepseek_vl_hybrid": "DeepseekVLHybridConditionalGenerate", "gemma3": "Gemma3ConditionalGenerate", @@ -145,6 +146,7 @@ "xlm_roberta": "XLMRobertaMaskedLM", }, "Model": { + "bart": "BartModel", "beit": "BeitModel", "bert": "BertModel", "cait": "CaiTModel", @@ -307,6 +309,7 @@ "detr": "DETRPanopticSegment", }, "QnA": { + "bart": "BartQnA", "bert": "BertQnA", "deberta": "DebertaQnA", "deberta_v2": "DebertaV2QnA", @@ -332,6 +335,7 @@ "tipsv2_dpt": "Tipsv2DptSemanticSegment", }, "SequenceClassify": { + "bart": "BartSequenceClassify", "bert": "BertSequenceClassify", "deberta": "DebertaSequenceClassify", "deberta_v2": "DebertaV2SequenceClassify", @@ -490,6 +494,7 @@ CONFIG_MAPPING_NAMES = { "autoencoder_kl": "AutoencoderKLConfig", + "bart": "BartConfig", "beit": "BeitConfig", "bert": "BertConfig", "cait": "CaiTConfig", @@ -694,6 +699,7 @@ } TOKENIZER_MAPPING_NAMES = { + "bart": "BartTokenizer", "bert": "BertTokenizer", "clip": "CLIPTokenizer", "deberta": "DebertaTokenizer", diff --git a/zeromodels/models/__init__.py b/zeromodels/models/__init__.py index 4707d7a9..67346365 100644 --- a/zeromodels/models/__init__.py +++ b/zeromodels/models/__init__.py @@ -1,4 +1,5 @@ from zeromodels.models import ( + bart, beit, bert, cait,