mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-24 17:31:18 +08:00
ab91ac1075
* chore: add load model with adapter support and fix bug in sample * chore: ignore temp during calculating prob in sample
24 lines
719 B
Python
24 lines
719 B
Python
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
from mlx.utils import tree_unflatten
|
|
|
|
from .lora import LoRALinear
|
|
|
|
|
|
def apply_lora_layers(model: nn.Module, adapter_file: str) -> nn.Module:
|
|
adapters = list(mx.load(adapter_file).items())
|
|
linear_replacements = {}
|
|
lora_layers = set(
|
|
[name.replace(".lora_a", "").replace(".lora_b", "") for name, _ in adapters]
|
|
)
|
|
|
|
for name, module in model.named_modules():
|
|
if name in lora_layers:
|
|
replacement_module = LoRALinear.from_linear(module)
|
|
linear_replacements[name] = replacement_module
|
|
|
|
model.update_modules(tree_unflatten(list(linear_replacements.items())))
|
|
|
|
model.update(tree_unflatten(adapters))
|
|
return model
|