tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix: config to allow optional input

Browse Source
This commit is contained in:
NanoCode012
2025-02-05 15:52:30 +07:00
parent 2bc7833a4e
commit 4cc60df876
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/axolotl/integrations/lolcats/linear_llama/configuration_linear_llama.py
View File

@@ -14,6 +14,8 @@
# limitations under the License. # limitations under the License.
"""Linear LLaMA model configuration""" """Linear LLaMA model configuration"""
from typing import Optional
from transformers import LlamaConfig from transformers import LlamaConfig
@@ -26,6 +28,8 @@ class LinearLlamaConfig(LlamaConfig):
attention_config (`dict`): attention_config (`dict`):
Dictionary containing the configuration for linear attention mechanism. Dictionary containing the configuration for linear attention mechanism.
Expected contents: Expected contents:
`attention_type` (str):
The type of attention to convert to.
`feature_map` (`str`): `feature_map` (`str`):
The type of feature map to use for linear attention. The type of feature map to use for linear attention.
`feature_map_kwargs` (`dict`): `feature_map_kwargs` (`dict`):
@@ -57,11 +61,11 @@ class LinearLlamaConfig(LlamaConfig):
model_type = "linear_llama" model_type = "linear_llama"
def __init__(self, attention_config: dict, **kwargs): def __init__(self, attention_config: Optional[dict] = None, **kwargs):
super().__init__(**kwargs) super().__init__(**kwargs)
# Set default attention config if none provided # Set default attention config if none provided
self.attention_config = attention_config self.attention_config = attention_config or {"attention_type": "softmax"}
@classmethod @classmethod
def from_llama(cls, llama_config: LlamaConfig, attention_config: dict): def from_llama(cls, llama_config: LlamaConfig, attention_config: dict):