fusion-bench 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (78) hide show
  1. fusion_bench/__init__.py +4 -0
  2. fusion_bench/compat/method/__init__.py +5 -2
  3. fusion_bench/compat/method/base_algorithm.py +3 -2
  4. fusion_bench/compat/modelpool/base_pool.py +3 -3
  5. fusion_bench/compat/taskpool/clip_image_classification.py +1 -1
  6. fusion_bench/dataset/gpt2_glue.py +1 -1
  7. fusion_bench/method/__init__.py +12 -2
  8. fusion_bench/method/analysis/task_vector_cos_similarity.py +95 -12
  9. fusion_bench/method/analysis/task_vector_violin_plot.py +160 -52
  10. fusion_bench/method/bitdelta/bitdelta.py +7 -23
  11. fusion_bench/method/ensemble.py +17 -2
  12. fusion_bench/method/expert_sparsity/mixtral/dynamic_skipping.py +2 -0
  13. fusion_bench/method/expert_sparsity/mixtral/layer_wise_pruning.py +2 -0
  14. fusion_bench/method/expert_sparsity/mixtral/progressive_pruning.py +2 -0
  15. fusion_bench/method/linear/__init__.py +6 -2
  16. fusion_bench/method/linear/{simple_average_for_llama.py → simple_average_for_causallm.py} +8 -4
  17. fusion_bench/method/linear/{task_arithmetic_for_llama.py → task_arithmetic_for_causallm.py} +22 -12
  18. fusion_bench/method/linear/ties_merging_for_causallm.py +70 -0
  19. fusion_bench/method/model_stock/__init__.py +1 -0
  20. fusion_bench/method/model_stock/model_stock.py +309 -0
  21. fusion_bench/method/regmean/clip_regmean.py +3 -6
  22. fusion_bench/method/regmean/regmean.py +27 -56
  23. fusion_bench/method/regmean/utils.py +56 -0
  24. fusion_bench/method/regmean_plusplus/regmean_plusplus.py +21 -60
  25. fusion_bench/method/simple_average.py +2 -2
  26. fusion_bench/method/slerp/__init__.py +1 -1
  27. fusion_bench/method/slerp/slerp.py +110 -14
  28. fusion_bench/method/task_arithmetic/task_arithmetic.py +35 -10
  29. fusion_bench/method/ties_merging/ties_merging.py +22 -6
  30. fusion_bench/method/we_moe/flan_t5_we_moe.py +9 -20
  31. fusion_bench/method/wudi/__init__.py +1 -0
  32. fusion_bench/method/wudi/wudi.py +105 -0
  33. fusion_bench/mixins/clip_classification.py +26 -6
  34. fusion_bench/mixins/lightning_fabric.py +4 -0
  35. fusion_bench/mixins/serialization.py +40 -83
  36. fusion_bench/modelpool/base_pool.py +1 -1
  37. fusion_bench/modelpool/causal_lm/causal_lm.py +285 -44
  38. fusion_bench/modelpool/seq2seq_lm/modelpool.py +146 -0
  39. fusion_bench/models/hf_clip.py +4 -0
  40. fusion_bench/models/hf_utils.py +10 -4
  41. fusion_bench/models/linearized/vision_model.py +6 -6
  42. fusion_bench/models/model_card_templates/default.md +8 -1
  43. fusion_bench/models/modeling_smile_mistral/__init__.py +1 -0
  44. fusion_bench/models/we_moe.py +8 -8
  45. fusion_bench/models/wrappers/ensemble.py +136 -7
  46. fusion_bench/scripts/cli.py +2 -2
  47. fusion_bench/taskpool/base_pool.py +99 -17
  48. fusion_bench/taskpool/clip_vision/taskpool.py +12 -5
  49. fusion_bench/taskpool/dummy.py +101 -13
  50. fusion_bench/taskpool/lm_eval_harness/taskpool.py +80 -0
  51. fusion_bench/taskpool/nyuv2_taskpool.py +28 -0
  52. fusion_bench/utils/__init__.py +1 -0
  53. fusion_bench/utils/data.py +6 -4
  54. fusion_bench/utils/devices.py +36 -11
  55. fusion_bench/utils/dtype.py +3 -2
  56. fusion_bench/utils/lazy_state_dict.py +85 -19
  57. fusion_bench/utils/packages.py +3 -3
  58. fusion_bench/utils/parameters.py +0 -2
  59. fusion_bench/utils/rich_utils.py +7 -3
  60. fusion_bench/utils/timer.py +92 -10
  61. {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/METADATA +10 -3
  62. {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/RECORD +77 -64
  63. fusion_bench_config/_get_started/llm_slerp.yaml +12 -0
  64. fusion_bench_config/method/ensemble/simple_ensemble.yaml +1 -0
  65. fusion_bench_config/method/linear/{simple_average_for_llama.yaml → simple_average_for_causallm.yaml} +1 -1
  66. fusion_bench_config/method/linear/task_arithmetic_for_causallm.yaml +4 -0
  67. fusion_bench_config/method/linear/ties_merging_for_causallm.yaml +13 -0
  68. fusion_bench_config/method/model_stock/model_stock.yaml +12 -0
  69. fusion_bench_config/method/slerp/slerp_lm.yaml +4 -0
  70. fusion_bench_config/method/wudi/wudi.yaml +4 -0
  71. fusion_bench_config/modelpool/CausalLMPool/{Qwen2.5-1.5B_math_and_coder.yaml → Qwen2.5-1.5B_math_and_code.yaml} +1 -2
  72. fusion_bench_config/modelpool/CausalLMPool/Qwen2.5-1.5B_three_models.yaml +11 -0
  73. fusion_bench_config/modelpool/CausalLMPool/llama-7b_3-models_v1.yaml +11 -0
  74. fusion_bench_config/method/linear/task_arithmetic_for_llama.yaml +0 -4
  75. {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/WHEEL +0 -0
  76. {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/entry_points.txt +0 -0
  77. {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/licenses/LICENSE +0 -0
  78. {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/top_level.txt +0 -0
@@ -1,5 +1,5 @@
1
1
  """
2
- Online documentation for this module: https://tanganke.github.io/fusion_bench/modelpool/causal_lm
2
+ Online documentation for this module: https://tanganke.github.io/fusion_bench/modelpool/llm
3
3
  """
4
4
 
5
5
  import logging
@@ -26,6 +26,7 @@ from fusion_bench import (
26
26
  instantiate,
27
27
  parse_dtype,
28
28
  )
29
+ from fusion_bench.models.hf_utils import create_default_model_card
29
30
  from fusion_bench.utils.lazy_state_dict import LazyStateDict
30
31
 
31
32
  log = logging.getLogger(__name__)
@@ -33,6 +34,44 @@ log = logging.getLogger(__name__)
33
34
 
34
35
  @auto_register_config
35
36
  class CausalLMPool(BaseModelPool):
37
+ """A model pool for managing and loading causal language models.
38
+
39
+ This class provides a unified interface for loading and managing multiple
40
+ causal language models, typically used in model fusion and ensemble scenarios.
41
+ It supports both eager and lazy loading strategies, and handles model
42
+ configuration through YAML configs or direct instantiation.
43
+
44
+ The pool can manage models from Hugging Face Hub, local paths, or custom
45
+ configurations. It also provides tokenizer management and model saving
46
+ capabilities with optional Hugging Face Hub integration.
47
+
48
+ Args:
49
+ models: Dictionary or configuration specifying the models to be managed.
50
+ Can contain model names mapped to paths or detailed configurations.
51
+ tokenizer: Tokenizer configuration, either a string path/name or
52
+ a DictConfig with detailed tokenizer settings.
53
+ model_kwargs: Additional keyword arguments passed to model loading.
54
+ Common options include torch_dtype, device_map, etc.
55
+ enable_lazy_loading: Whether to use lazy loading for models. When True,
56
+ models are loaded as LazyStateDict objects instead of actual models,
57
+ which can save memory for large model collections.
58
+ **kwargs: Additional arguments passed to the parent BaseModelPool.
59
+
60
+ Example:
61
+ ```python
62
+ >>> pool = CausalLMPool(
63
+ ... models={
64
+ ... "model_a": "microsoft/DialoGPT-medium",
65
+ ... "model_b": "/path/to/local/model"
66
+ ... },
67
+ ... tokenizer="microsoft/DialoGPT-medium",
68
+ ... model_kwargs={"torch_dtype": "bfloat16"}
69
+ ... )
70
+ >>> model = pool.load_model("model_a")
71
+ >>> tokenizer = pool.load_tokenizer()
72
+ ```
73
+ """
74
+
36
75
  def __init__(
37
76
  self,
38
77
  models,
@@ -47,6 +86,20 @@ class CausalLMPool(BaseModelPool):
47
86
  self.model_kwargs = DictConfig({})
48
87
 
49
88
  def get_model_path(self, model_name: str):
89
+ """Extract the model path from the model configuration.
90
+
91
+ Args:
92
+ model_name: The name of the model as defined in the models configuration.
93
+
94
+ Returns:
95
+ str: The path or identifier for the model. For string configurations,
96
+ returns the string directly. For dict configurations, extracts
97
+ the 'pretrained_model_name_or_path' field.
98
+
99
+ Raises:
100
+ RuntimeError: If the model configuration is invalid or the model
101
+ name is not found in the configuration.
102
+ """
50
103
  model_name_or_config = self._models[model_name]
51
104
  if isinstance(model_name_or_config, str):
52
105
  return model_name_or_config
@@ -56,6 +109,16 @@ class CausalLMPool(BaseModelPool):
56
109
  raise RuntimeError("Invalid model configuration")
57
110
 
58
111
  def get_model_kwargs(self):
112
+ """Get processed model keyword arguments for model loading.
113
+
114
+ Converts the stored `model_kwargs` from DictConfig to a regular dictionary
115
+ and processes special arguments like torch_dtype for proper model loading.
116
+
117
+ Returns:
118
+ dict: Processed keyword arguments ready to be passed to model
119
+ loading functions. The torch_dtype field, if present, is
120
+ converted from string to the appropriate torch dtype object.
121
+ """
59
122
  model_kwargs = (
60
123
  OmegaConf.to_container(self.model_kwargs, resolve=True)
61
124
  if isinstance(self.model_kwargs, DictConfig)
@@ -71,31 +134,52 @@ class CausalLMPool(BaseModelPool):
71
134
  model_name_or_config: str | DictConfig,
72
135
  *args,
73
136
  **kwargs,
74
- ) -> PreTrainedModel:
75
- """
76
- Example of YAML config:
137
+ ) -> Union[PreTrainedModel, LazyStateDict]:
138
+ """Load a causal language model from the model pool.
77
139
 
78
- ```yaml
79
- models:
80
- _pretrained_: path_to_pretrained_model # if a plain string, it will be passed to AutoModelForCausalLM.from_pretrained
81
- model_a: path_to_model_a
82
- model_b: path_to_model_b
83
- ```
140
+ This method supports multiple loading strategies:
141
+ 1. Loading by model name from the configured model pool
142
+ 2. Loading from a direct configuration dictionary
143
+ 3. Lazy loading using LazyStateDict for memory efficiency
84
144
 
85
- or equivalently,
86
-
87
- ```yaml
88
- models:
89
- _pretrained_:
90
- _target_: transformers.AutoModelForCausalLM # any callable that returns a model
91
- pretrained_model_name_or_path: path_to_pretrained_model
92
- model_a:
93
- _target_: transformers.AutoModelForCausalLM
94
- pretrained_model_name_or_path: path_to_model_a
95
- model_b:
96
- _target_: transformers.AutoModelForCausalLM
97
- pretrained_model_name_or_path: path_to_model_b
98
- ```
145
+ The method automatically handles different model configuration formats
146
+ and applies the appropriate loading strategy based on the enable_lazy_loading flag.
147
+
148
+ Args:
149
+ model_name_or_config: Either a string model name that exists in the
150
+ model pool configuration, or a DictConfig/dict containing the
151
+ model configuration directly.
152
+ *args: Additional positional arguments passed to the model constructor.
153
+ **kwargs: Additional keyword arguments passed to the model constructor.
154
+ These will be merged with the pool's model_kwargs.
155
+
156
+ Returns:
157
+ Union[PreTrainedModel, LazyStateDict]: The loaded model. Returns a
158
+ PreTrainedModel for normal loading or a LazyStateDict for lazy loading.
159
+
160
+ Raises:
161
+ RuntimeError: If the model configuration is invalid.
162
+ KeyError: If the model name is not found in the model pool.
163
+
164
+ Example YAML configurations:
165
+ Simple string configuration:
166
+ ```yaml
167
+ models:
168
+ _pretrained_: path_to_pretrained_model
169
+ model_a: path_to_model_a
170
+ model_b: path_to_model_b
171
+ ```
172
+
173
+ Detailed configuration:
174
+ ```yaml
175
+ models:
176
+ _pretrained_:
177
+ _target_: transformers.AutoModelForCausalLM
178
+ pretrained_model_name_or_path: path_to_pretrained_model
179
+ model_a:
180
+ _target_: transformers.AutoModelForCausalLM
181
+ pretrained_model_name_or_path: path_to_model_a
182
+ ```
99
183
  """
100
184
  model_kwargs = self.get_model_kwargs()
101
185
  model_kwargs.update(kwargs)
@@ -139,23 +223,36 @@ class CausalLMPool(BaseModelPool):
139
223
  return model
140
224
 
141
225
  def load_tokenizer(self, *args, **kwargs) -> PreTrainedTokenizer:
142
- """
143
- Example of YAML config:
144
-
145
- ```yaml
146
- tokenizer: google/gemma-2-2b-it # if a plain string, it will be passed to AutoTokenizer.from_pretrained
147
- ```
226
+ """Load the tokenizer associated with this model pool.
148
227
 
149
- or equivalently,
228
+ Loads a tokenizer based on the tokenizer configuration provided during
229
+ pool initialization. Supports both simple string paths and detailed
230
+ configuration dictionaries.
150
231
 
151
- ```yaml
152
- tokenizer:
153
- _target_: transformers.AutoTokenizer # any callable that returns a tokenizer
154
- pretrained_model_name_or_path: google/gemma-2-2b-it
155
- ```
232
+ Args:
233
+ *args: Additional positional arguments passed to the tokenizer constructor.
234
+ **kwargs: Additional keyword arguments passed to the tokenizer constructor.
156
235
 
157
236
  Returns:
158
- PreTrainedTokenizer: The tokenizer.
237
+ PreTrainedTokenizer: The loaded tokenizer instance.
238
+
239
+ Raises:
240
+ AssertionError: If no tokenizer is defined in the configuration.
241
+
242
+ Example YAML configurations:
243
+ Simple string configuration:
244
+ ```yaml
245
+ tokenizer: google/gemma-2-2b-it
246
+ ```
247
+
248
+ Detailed configuration:
249
+ ```yaml
250
+ tokenizer:
251
+ _target_: transformers.AutoTokenizer
252
+ pretrained_model_name_or_path: google/gemma-2-2b-it
253
+ use_fast: true
254
+ padding_side: left
255
+ ```
159
256
  """
160
257
  assert self.tokenizer is not None, "Tokenizer is not defined in the config"
161
258
  log.info("Loading tokenizer.", stacklevel=2)
@@ -175,17 +272,54 @@ class CausalLMPool(BaseModelPool):
175
272
  save_tokenizer: bool = False,
176
273
  tokenizer_kwargs=None,
177
274
  tokenizer: Optional[PreTrainedTokenizer] = None,
275
+ algorithm_config: Optional[DictConfig] = None,
276
+ description: Optional[str] = None,
277
+ base_model_in_modelcard: bool = True,
178
278
  **kwargs,
179
279
  ):
180
- """
181
- Save the model to the specified path.
280
+ """Save a model to the specified path with optional tokenizer and Hub upload.
281
+
282
+ This method provides comprehensive model saving capabilities including
283
+ optional tokenizer saving, dtype conversion, model card creation, and
284
+ Hugging Face Hub upload. The model is saved in the standard Hugging Face format.
182
285
 
183
286
  Args:
184
- model (PreTrainedModel): The model to be saved.
185
- path (str): The path where the model will be saved.
186
- push_to_hub (bool, optional): Whether to push the model to the Hugging Face Hub. Defaults to False.
187
- save_tokenizer (bool, optional): Whether to save the tokenizer along with the model. Defaults to False.
188
- **kwargs: Additional keyword arguments passed to the `save_pretrained` method.
287
+ model: The PreTrainedModel instance to be saved.
288
+ path: The local path where the model will be saved. Supports tilde
289
+ expansion for home directory paths.
290
+ push_to_hub: Whether to push the saved model to the Hugging Face Hub.
291
+ Requires proper authentication and repository permissions.
292
+ model_dtype: Optional string specifying the target dtype for the model
293
+ before saving (e.g., "float16", "bfloat16"). The model will be
294
+ converted to this dtype before saving.
295
+ save_tokenizer: Whether to save the tokenizer alongside the model.
296
+ If True, the tokenizer will be loaded using the pool's tokenizer
297
+ configuration and saved to the same path.
298
+ tokenizer_kwargs: Additional keyword arguments for tokenizer loading
299
+ when save_tokenizer is True.
300
+ tokenizer: Optional pre-loaded tokenizer instance. If provided, this
301
+ tokenizer will be saved regardless of the save_tokenizer flag.
302
+ algorithm_config: Optional DictConfig containing algorithm configuration.
303
+ If provided, a model card will be created with algorithm details.
304
+ description: Optional description for the model card. If not provided
305
+ and algorithm_config is given, a default description will be generated.
306
+ **kwargs: Additional keyword arguments passed to the model's
307
+ save_pretrained method.
308
+
309
+ Example:
310
+ ```python
311
+ >>> pool = CausalLMPool(models=..., tokenizer=...)
312
+ >>> model = pool.load_model("my_model")
313
+ >>> pool.save_model(
314
+ ... model,
315
+ ... "/path/to/save",
316
+ ... save_tokenizer=True,
317
+ ... model_dtype="float16",
318
+ ... push_to_hub=True,
319
+ ... algorithm_config=algorithm_config,
320
+ ... description="Custom merged model"
321
+ ... )
322
+ ```
189
323
  """
190
324
  path = os.path.expanduser(path)
191
325
  # NOTE: if tokenizer is provided, it will be saved regardless of `save_tokenizer`
@@ -207,11 +341,75 @@ class CausalLMPool(BaseModelPool):
207
341
  **kwargs,
208
342
  )
209
343
 
344
+ # Create and save model card if algorithm_config is provided
345
+ if algorithm_config is not None:
346
+ if description is None:
347
+ description = "Model created using FusionBench."
348
+ model_card_str = create_default_model_card(
349
+ base_model=(
350
+ self.get_model_path("_pretrained_")
351
+ if base_model_in_modelcard and self.has_pretrained
352
+ else None
353
+ ),
354
+ models=[self.get_model_path(m) for m in self.model_names],
355
+ description=description,
356
+ algorithm_config=algorithm_config,
357
+ modelpool_config=self.config,
358
+ )
359
+ with open(os.path.join(path, "README.md"), "w") as f:
360
+ f.write(model_card_str)
361
+
210
362
 
211
363
  class CausalLMBackbonePool(CausalLMPool):
364
+ """A specialized model pool that loads only the transformer backbone layers.
365
+
366
+ This class extends CausalLMPool to provide access to just the transformer
367
+ layers (backbone) of causal language models, excluding the language modeling
368
+ head and embeddings. This is useful for model fusion scenarios where only
369
+ the core transformer layers are needed.
370
+
371
+ The class automatically extracts the `model.layers` component from loaded
372
+ AutoModelForCausalLM instances, providing direct access to the transformer
373
+ blocks. Lazy loading is not supported for this pool type.
374
+
375
+ Note:
376
+ This pool automatically disables lazy loading as it needs to access
377
+ the internal structure of the model to extract the backbone layers.
378
+
379
+ Example:
380
+ ```python
381
+ >>> backbone_pool = CausalLMBackbonePool(
382
+ ... models={"model_a": "microsoft/DialoGPT-medium"},
383
+ ... tokenizer="microsoft/DialoGPT-medium"
384
+ ... )
385
+ >>> layers = backbone_pool.load_model("model_a") # Returns nn.ModuleList of transformer layers
386
+ ```
387
+ """
388
+
212
389
  def load_model(
213
390
  self, model_name_or_config: str | DictConfig, *args, **kwargs
214
391
  ) -> Module:
392
+ """Load only the transformer backbone layers from a causal language model.
393
+
394
+ This method loads a complete causal language model and then extracts
395
+ only the transformer layers (backbone), discarding the embedding layers
396
+ and language modeling head. This is useful for model fusion scenarios
397
+ where only the core transformer computation is needed.
398
+
399
+ Args:
400
+ model_name_or_config: Either a string model name from the pool
401
+ configuration or a DictConfig with model loading parameters.
402
+ *args: Additional positional arguments passed to the parent load_model method.
403
+ **kwargs: Additional keyword arguments passed to the parent load_model method.
404
+
405
+ Returns:
406
+ Module: The transformer layers (typically a nn.ModuleList) containing
407
+ the core transformer blocks without embeddings or output heads.
408
+
409
+ Note:
410
+ Lazy loading is automatically disabled for this method as it needs
411
+ to access the internal model structure to extract the layers.
412
+ """
215
413
  if self.enable_lazy_loading:
216
414
  log.warning(
217
415
  "CausalLMBackbonePool does not support lazy loading. "
@@ -231,6 +429,49 @@ def load_peft_causal_lm(
231
429
  is_trainable: bool = True,
232
430
  merge_and_unload: bool = False,
233
431
  ):
432
+ """Load a causal language model with PEFT (Parameter-Efficient Fine-Tuning) adapters.
433
+
434
+ This function loads a base causal language model and applies PEFT adapters
435
+ (such as LoRA, AdaLoRA, or other parameter-efficient fine-tuning methods)
436
+ to create a fine-tuned model. It supports both keeping the adapters separate
437
+ or merging them into the base model.
438
+
439
+ Args:
440
+ base_model_path: Path or identifier for the base causal language model.
441
+ Can be a Hugging Face model name or local path.
442
+ peft_model_path: Path to the PEFT adapter configuration and weights.
443
+ This should contain the adapter_config.json and adapter weights.
444
+ torch_dtype: The torch data type to use for the model. Common options
445
+ include "float16", "bfloat16", "float32". Defaults to "bfloat16".
446
+ is_trainable: Whether the loaded PEFT model should be trainable.
447
+ Set to False for inference-only usage to save memory.
448
+ merge_and_unload: Whether to merge the PEFT adapters into the base model
449
+ and unload the adapter weights. When True, returns a standard
450
+ PreTrainedModel instead of a PeftModel.
451
+
452
+ Returns:
453
+ Union[PeftModel, PreTrainedModel]: The loaded model with PEFT adapters.
454
+ Returns a PeftModel if merge_and_unload is False, or a PreTrainedModel
455
+ if the adapters are merged and unloaded.
456
+
457
+ Example:
458
+ ```python
459
+ >>> # Load model with adapters for training
460
+ >>> model = load_peft_causal_lm(
461
+ ... "microsoft/DialoGPT-medium",
462
+ ... "/path/to/lora/adapters",
463
+ ... is_trainable=True
464
+ ... )
465
+
466
+ >>> # Load and merge adapters for inference
467
+ >>> merged_model = load_peft_causal_lm(
468
+ ... "microsoft/DialoGPT-medium",
469
+ ... "/path/to/lora/adapters",
470
+ ... merge_and_unload=True,
471
+ ... is_trainable=False
472
+ ... )
473
+ ```
474
+ """
234
475
  base_model = AutoModelForCausalLM.from_pretrained(
235
476
  base_model_path, torch_dtype=torch_dtype
236
477
  )
@@ -18,6 +18,48 @@ def load_lora_model(
18
18
  is_trainable: bool = True,
19
19
  merge_and_unload: bool = True,
20
20
  ):
21
+ """Load a sequence-to-sequence model with LoRA (Low-Rank Adaptation) fine-tuning.
22
+
23
+ This function loads a base sequence-to-sequence language model and applies
24
+ LoRA adapters for parameter-efficient fine-tuning. LoRA allows for efficient
25
+ adaptation of large models by adding trainable low-rank matrices to the
26
+ existing weights without modifying the original parameters.
27
+
28
+ Args:
29
+ base_model_path: Path or identifier for the base sequence-to-sequence model.
30
+ Can be a Hugging Face model name (e.g., "t5-base") or local path.
31
+ peft_model_path: Path to the directory containing LoRA adapter weights
32
+ and configuration. Should include adapter_config.json and adapter weights.
33
+ is_trainable: Whether the loaded model should be trainable. Set to False
34
+ for inference-only usage to save memory and computation.
35
+ merge_and_unload: Whether to merge the LoRA weights into the base model
36
+ and unload the adapter. When True, returns a standard model instead
37
+ of a PeftModel, which can be more efficient for inference.
38
+
39
+ Returns:
40
+ Union[PeftModel, AutoModelForSeq2SeqLM]: The loaded model with LoRA
41
+ adapters. Returns a PeftModel if merge_and_unload is False, or
42
+ a standard AutoModelForSeq2SeqLM if adapters are merged.
43
+
44
+ Example:
45
+ ```python
46
+ >>> # Load model with separate adapters for training
47
+ >>> model = load_lora_model(
48
+ ... "t5-base",
49
+ ... "/path/to/lora/adapters",
50
+ ... is_trainable=True,
51
+ ... merge_and_unload=False
52
+ ... )
53
+
54
+ >>> # Load and merge adapters for efficient inference
55
+ >>> merged_model = load_lora_model(
56
+ ... "t5-base",
57
+ ... "/path/to/lora/adapters",
58
+ ... is_trainable=False,
59
+ ... merge_and_unload=True
60
+ ... )
61
+ ```
62
+ """
21
63
  base_model = AutoModelForSeq2SeqLM.from_pretrained(base_model_path)
22
64
  model = PeftModel.from_pretrained(
23
65
  base_model,
@@ -30,6 +72,46 @@ def load_lora_model(
30
72
 
31
73
 
32
74
  class Seq2SeqLMPool(BaseModelPool):
75
+ """A model pool specialized for sequence-to-sequence language models.
76
+
77
+ This model pool provides management and loading capabilities for sequence-to-sequence
78
+ (seq2seq) language models such as T5, BART, and mT5. It extends the base model pool
79
+ functionality with seq2seq-specific features including tokenizer management and
80
+ model configuration handling.
81
+
82
+ Seq2seq models are particularly useful for tasks that require generating output
83
+ sequences from input sequences, such as translation, summarization, question
84
+ answering, and text generation. This pool streamlines the process of loading
85
+ and configuring multiple seq2seq models for fusion and ensemble scenarios.
86
+
87
+ Key Features:
88
+ - Specialized loading for AutoModelForSeq2SeqLM models
89
+ - Integrated tokenizer management
90
+ - Support for model-specific keyword arguments
91
+ - Automatic dtype parsing and configuration
92
+ - Compatible with PEFT (Parameter-Efficient Fine-Tuning) adapters
93
+
94
+ Attributes:
95
+ _tokenizer: Configuration for the tokenizer associated with the models
96
+ _model_kwargs: Default keyword arguments applied to all model loading operations
97
+
98
+ Example:
99
+ ```python
100
+ pool = Seq2SeqLMPool(
101
+ models={
102
+ "t5_base": "t5-base",
103
+ "t5_large": "t5-large",
104
+ "custom_model": "/path/to/local/model"
105
+ },
106
+ tokenizer={"_target_": "transformers.T5Tokenizer",
107
+ "pretrained_model_name_or_path": "t5-base"},
108
+ model_kwargs={"torch_dtype": "float16", "device_map": "auto"}
109
+ )
110
+ model = pool.load_model("t5_base")
111
+ tokenizer = pool.load_tokenizer()
112
+ ```
113
+ """
114
+
33
115
  _config_mapping = BaseModelPool._config_mapping | {
34
116
  "_tokenizer": "tokenizer",
35
117
  "_model_kwargs": "model_kwargs",
@@ -43,6 +125,35 @@ class Seq2SeqLMPool(BaseModelPool):
43
125
  model_kwargs: Optional[DictConfig] = None,
44
126
  **kwargs,
45
127
  ):
128
+ """Initialize the sequence-to-sequence language model pool.
129
+
130
+ Sets up the model pool with configurations for models, tokenizer, and
131
+ default model loading parameters. Automatically processes model kwargs
132
+ to handle special configurations like torch_dtype parsing.
133
+
134
+ Args:
135
+ models: Configuration dictionary specifying the seq2seq models to manage.
136
+ Keys are model names, values can be model paths/names or detailed configs.
137
+ tokenizer: Configuration for the tokenizer to use with the models.
138
+ Can be a simple path/name or detailed configuration with _target_.
139
+ model_kwargs: Default keyword arguments applied to all model loading
140
+ operations. Common options include torch_dtype, device_map, etc.
141
+ The torch_dtype field is automatically parsed from string to dtype.
142
+ **kwargs: Additional arguments passed to the parent BaseModelPool.
143
+
144
+ Example:
145
+ ```python
146
+ pool = Seq2SeqLMPool(
147
+ models={
148
+ "base": "t5-base",
149
+ "large": {"_target_": "transformers.AutoModelForSeq2SeqLM",
150
+ "pretrained_model_name_or_path": "t5-large"}
151
+ },
152
+ tokenizer="t5-base",
153
+ model_kwargs={"torch_dtype": "bfloat16"}
154
+ )
155
+ ```
156
+ """
46
157
  super().__init__(models, **kwargs)
47
158
  self._tokenizer = tokenizer
48
159
  self._model_kwargs = model_kwargs
@@ -55,11 +166,46 @@ class Seq2SeqLMPool(BaseModelPool):
55
166
  )
56
167
 
57
168
  def load_model(self, model_name_or_config: str | DictConfig, *args, **kwargs):
169
+ """Load a sequence-to-sequence language model from the pool.
170
+
171
+ Loads a seq2seq model using the parent class loading mechanism while
172
+ automatically applying the pool's default model kwargs. The method
173
+ merges the pool's model_kwargs with any additional kwargs provided,
174
+ giving priority to the explicitly provided kwargs.
175
+
176
+ Args:
177
+ model_name_or_config: Either a string model name from the pool
178
+ configuration or a DictConfig containing model loading parameters.
179
+ *args: Additional positional arguments passed to the parent load_model method.
180
+ **kwargs: Additional keyword arguments that override the pool's default
181
+ model_kwargs. Common options include device, torch_dtype, etc.
182
+
183
+ Returns:
184
+ AutoModelForSeq2SeqLM: The loaded sequence-to-sequence language model.
185
+ """
58
186
  model_kwargs = deepcopy(self._model_kwargs)
59
187
  model_kwargs.update(kwargs)
60
188
  return super().load_model(model_name_or_config, *args, **model_kwargs)
61
189
 
62
190
  def load_tokenizer(self, *args, **kwargs):
191
+ """Load the tokenizer associated with the sequence-to-sequence models.
192
+
193
+ Loads a tokenizer based on the tokenizer configuration provided during
194
+ pool initialization. The tokenizer should be compatible with the seq2seq
195
+ models in the pool and is typically used for preprocessing input text
196
+ and postprocessing generated output.
197
+
198
+ Args:
199
+ *args: Additional positional arguments passed to the tokenizer constructor.
200
+ **kwargs: Additional keyword arguments passed to the tokenizer constructor.
201
+
202
+ Returns:
203
+ PreTrainedTokenizer: The loaded tokenizer instance compatible with
204
+ the seq2seq models in this pool.
205
+
206
+ Raises:
207
+ AssertionError: If no tokenizer configuration is provided.
208
+ """
63
209
  assert self._tokenizer is not None, "Tokenizer is not defined in the config"
64
210
  tokenizer = isinstance(self._tokenizer, *args, **kwargs)
65
211
  return tokenizer
@@ -195,5 +195,9 @@ class HFCLIPClassifier(nn.Module):
195
195
  pass
196
196
  elif isinstance(image_embeds, BaseModelOutputWithPooling):
197
197
  image_embeds = image_embeds[1]
198
+ elif isinstance(image_embeds, dict) and "pooler_output" in image_embeds:
199
+ image_embeds = image_embeds["pooler_output"]
200
+ else:
201
+ raise ValueError("Unsupported output type from vision model outputs")
198
202
  image_embeds = self.clip_model.visual_projection(image_embeds)
199
203
  return image_embeds
@@ -42,12 +42,14 @@ def load_model_card_template(basename: str) -> str:
42
42
  FileNotFoundError: If the template file is not found in any of the search locations.
43
43
  """
44
44
  if os.path.exists(basename):
45
- return open(basename).read()
45
+ with open(basename, "r") as f:
46
+ return f.read()
46
47
 
47
48
  for template_dir in MODEL_CARD_TEMPLATE_DIRS:
48
49
  template_path = os.path.join(template_dir, basename)
49
50
  if os.path.exists(template_path):
50
- return open(template_path).read()
51
+ with open(template_path, "r") as f:
52
+ return f.read()
51
53
 
52
54
  raise FileNotFoundError(f"Model card template '{basename}' not found.")
53
55
 
@@ -141,6 +143,9 @@ def save_pretrained_with_remote_code(
141
143
 
142
144
  def create_default_model_card(
143
145
  models: list[str],
146
+ base_model: Optional[str] = None,
147
+ title: str = "Deep Model Fusion",
148
+ tags: list[str] = ["fusion-bench", "merge"],
144
149
  description=None,
145
150
  algorithm_config: DictConfig = None,
146
151
  modelpool_config: DictConfig = None,
@@ -149,10 +154,11 @@ def create_default_model_card(
149
154
 
150
155
  template: Template = Template(load_model_card_template("default.md"))
151
156
  card = template.render(
157
+ base_model=base_model,
152
158
  models=models,
153
159
  library_name="transformers",
154
- tags=["fusion-bench", "merge"],
155
- title="Deep Model Fusion",
160
+ title=title,
161
+ tags=tags,
156
162
  description=description,
157
163
  algorithm_config_str=try_to_yaml(algorithm_config),
158
164
  modelpool_config_str=try_to_yaml(modelpool_config),