File size: 2,571 Bytes
a164e13 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
# Copyright (C) 2024 Charles O. Goddard
#
# This software is free software: you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This software is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see http://www.gnu.org/licenses/.
from typing import List, Optional
import click
import yaml
from pydantic import BaseModel
from mergekit.common import MergeOptions
from mergekit.config import (
ConditionalParameter,
InputSliceDefinition,
MergeConfiguration,
)
from mergekit.merge import run_merge
class LayerSlice(BaseModel):
model: str
start: int
end: int
scale: Optional[float] = None
class BakllamaConfig(BaseModel):
layer_slices: List[LayerSlice]
embedding_source: Optional[str] = None
lm_head_source: Optional[str] = None
@click.command("bakllama")
@click.argument("config_path", type=click.Path(exists=True, dir_okay=False))
@click.argument("out_path", type=str)
@click.option(
"--clone-tensors/--no-clone-tensors",
type=bool,
is_flag=True,
help="Clone tensors before saving, to allow multiple occurrences of the same layer",
default=False,
)
@click.option("--fp16/--no-fp16", type=bool, default=False)
def main(
config_path: str,
out_path: str,
clone_tensors: bool,
fp16: bool,
):
"""Wrapper for using legacy bakllama configuration files."""
with open(config_path, "r", encoding="utf-8") as file:
config = BakllamaConfig.model_validate(yaml.safe_load(file))
slices = []
for s in config.layer_slices:
parameters = {}
if s.scale is not None:
parameters["scale"] = ConditionalParameter(
value=s.scale, filter="down_proj"
)
slices.append(
InputSliceDefinition(
model=s.model, layer_range=(s.start, s.end), parameters=parameters
)
)
merge_config = MergeConfiguration(
merge_method="passthrough", slices=slices, dtype="float16" if fp16 else None
)
run_merge(merge_config, out_path, MergeOptions(clone_tensors=clone_tensors))
if __name__ == "__main__":
main()
|