🦖 X—LLM: Simple & Cutting Edge LLM Finetuning
from typing import Dict, List, Optional, Tuple

import datasets
from tqdm import tqdm

from src.xllm import enums
from src.xllm.core.config import HuggingFaceConfig
from src.xllm.datasets.base import BaseDataset
from src.xllm.types import RawSample


class AntropicDataset(BaseDataset):
    _HF_DATASET_ID = "Anthropic/hh-rlhf"

    @classmethod
    def download(cls, config: HuggingFaceConfig) -> Tuple[List[RawSample], Optional[List[RawSample]]]:
        rlhf_dataset = datasets.load_dataset(cls._HF_DATASET_ID)

        parsed_data: Dict[str, List[RawSample]] = dict()

        for split in ["train", "test"]:
            parsed_data[split] = list()

            for sample in tqdm(rlhf_dataset[split], desc=f"Parsing {split}"):
                text_parts = sample["chosen"].split("\n\n")[1:]

                parsed_data[split].append(text_parts)

        train = parsed_data["train"]
        evaluation = parsed_data["test"]

        return train, evaluation

    def get_sample(self, index: int) -> RawSample:
        sample = {enums.General.text_parts: self.data[index]}
        return sample

examples/minimal/dataset.py
KompleteAI/xllm
from examples.minimal.dataset import AntropicDataset
from src.xllm.cli.train import cli_run_train
from src.xllm.core.config import HuggingFaceConfig
from src.xllm.datasets.registry import datasets_registry

if __name__ == "__main__":
    datasets_registry.add(key="antropic", value=AntropicDataset)
    cli_run_train(config_cls=HuggingFaceConfig)

examples/minimal/train.py
KompleteAI/xllm
from setuptools import find_packages, setup

# Extras
extras = dict()

extras["quality"] = [
    "black",
    "mypy",
    "mypy-extensions",
    "pre-commit",
    "ruff",
]

extras["tests"] = [
    "pytest",
    "coverage",
    "pytest-html",
    "pytest-cov",
]

extras["build"] = [
    "wheel",
    "twine",
]

extras["dev"] = extras["quality"] + extras["tests"] + extras["build"]

extras["deepspeed"] = ["deepspeed"]

extras["auto-gptq"] = ["auto-gptq"]

extras["flash-attn"] = ["flash-attn>=2.2.1"]

extras["train"] = extras["deepspeed"] + extras["flash-attn"] + extras["auto-gptq"]

extras["all"] = extras["train"] + extras["dev"]

# Install deps
install_requires = [
    "numpy>=1.17",
    "packaging>=20.0",
    "psutil",
    "torch>=2.0.1",
    "loguru",
    "peft>=0.5.0",
    "wandb",
    "python-dotenv",
    "requests",
    "optimum>=1.12.0",
    "bitsandbytes>=0.41.1",
    "scipy",
    "transformers",
    "tqdm",
    "safetensors",
]

# Setup
setup(
    name="xllm",
    version="0.3.24",
    description="Simple & Cutting Edge LLM Finetuning",
    license_files=["LICENSE"],
    long_description=open("README.md", "r", encoding="utf-8").read(),
    long_description_content_type="text/markdown",
    keywords="ai nlp llm text deep-learning",
    license="Apache",
    author="BobaZooba",
    author_email="bobazooba@gmail.com",
    url="https://github.com/kompleteai/xllm",
    package_dir={"": "src"},
    packages=find_packages("src"),
    package_data={"xllm": ["py.typed"]},
    entry_points={},
    python_requires=">=3.8",
    install_requires=install_requires,
    extras_require=extras,
    classifiers=[
        "Development Status :: 5 - Production/Stable",
        "Intended Audience :: Developers",
        "Intended Audience :: Education",
        "Intended Audience :: Science/Research",
        "License :: OSI Approved :: Apache Software License",
        "Operating System :: OS Independent",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.8",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
)

# Release checklist
# 1. Change the version in __init__.py and setup.py.
# 2. Commit these changes with the message: "Release: VERSION"
# 3. Add a tag in git to mark the release: "git tag VERSION -m 'Adds tag VERSION for pypi' "
#    Push the tag to git: git push --tags origin main
# 4. Run the following commands in the top-level directory:
#      [Optional] rm -rf dist/
#      python setup.py bdist_wheel
#      python setup.py sdist
# 5. Upload the package to the pypi test server first:
#      twine upload dist/* -r testpypi
# 6. Check that you can install it in a virtualenv by running:
#      pip install --upgrade -i https://testpypi.python.org/pypi xllm
# 7. Upload the final version to actual pypi:
#      twine upload dist/* -r pypi
# 8. Add release notes to the tag in github once everything is looking hunky-dory.
# 9. Update the version in __init__.py, setup.py to the new version "-dev" and push to master

setup.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: F401

__version__ = "0.3.24"

from . import enums, types
from .cli.download import cli_run_download
from .cli.fuse import cli_run_fuse
from .cli.train import cli_run_train
from .core.config import HuggingFaceConfig
from .run.download import download
from .run.fuse import fuse
from .run.train import train
from .utils.cli import setup_cli
from .utils.logger import dist_logger
from .utils.post_training import fuse_lora

src/xllm/__init__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .cli.train import cli_run_train
from .core.config import HuggingFaceConfig

if __name__ == "__main__":
    cli_run_train(config_cls=HuggingFaceConfig)

src/xllm/__main__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: F401

from .download import cli_run_download
from .fuse import cli_run_fuse
from .quantize import cli_run_quantize
from .train import cli_run_train

src/xllm/cli/__init__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Tuple, Type

from transformers import HfArgumentParser, PreTrainedModel, PreTrainedTokenizer

from ..core.config import HuggingFaceConfig
from ..run.download import download
from ..utils.cli import setup_cli


def cli_run_download(
    config_cls: Type[HuggingFaceConfig] = HuggingFaceConfig,
) -> Tuple[PreTrainedTokenizer, PreTrainedModel]:
    parser = HfArgumentParser(config_cls)
    config = parser.parse_args_into_dataclasses()[0]
    setup_cli(config=config, logger_path="./xllm_download.log")
    tokenizer, model = download(config=config)
    return tokenizer, model


if __name__ == "__main__":
    cli_run_download(config_cls=HuggingFaceConfig)

src/xllm/cli/download.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Tuple, Type

from transformers import HfArgumentParser, PreTrainedModel, PreTrainedTokenizer

from ..core.config import HuggingFaceConfig
from ..run.fuse import fuse
from ..utils.cli import setup_cli


def cli_run_fuse(
    config_cls: Type[HuggingFaceConfig] = HuggingFaceConfig,
) -> Tuple[PreTrainedTokenizer, PreTrainedModel]:
    parser = HfArgumentParser(config_cls)
    config = parser.parse_args_into_dataclasses()[0]
    setup_cli(config=config, logger_path="./xllm_fuse.log")
    tokenizer, model = fuse(config=config)
    return tokenizer, model


if __name__ == "__main__":
    cli_run_fuse(config_cls=HuggingFaceConfig)

src/xllm/cli/fuse.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Type

from transformers import HfArgumentParser

from ..core.config import HuggingFaceConfig
from ..quantization.quantizer import Quantizer
from ..run.quantize import quantize
from ..utils.cli import setup_cli


def cli_run_quantize(
    config_cls: Type[HuggingFaceConfig] = HuggingFaceConfig,
) -> Quantizer:
    parser = HfArgumentParser(config_cls)
    config = parser.parse_args_into_dataclasses()[0]
    setup_cli(config=config, logger_path="./xllm_gptq_quantize.log")
    quantizer = quantize(config=config)
    return quantizer


if __name__ == "__main__":
    cli_run_quantize(config_cls=HuggingFaceConfig)

src/xllm/cli/quantize.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Type

from transformers import (
    HfArgumentParser,
)

from ..core.config import HuggingFaceConfig
from ..experiments.base import Experiment
from ..run.train import train
from ..utils.cli import setup_cli


def cli_run_train(
    config_cls: Type[HuggingFaceConfig] = HuggingFaceConfig,
) -> Experiment:
    parser = HfArgumentParser(config_cls)
    config = parser.parse_args_into_dataclasses()[0]
    setup_cli(config=config, logger_path="./xllm_train.log")
    experiment = train(config=config)
    return experiment


if __name__ == "__main__":
    cli_run_train(config_cls=HuggingFaceConfig)

src/xllm/cli/train.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: F401

from .base import BaseCollator
from .completion import CompletionCollator
from .lm import LMCollator
from .registry import collators_registry

src/xllm/collators/__init__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from abc import ABC, abstractmethod
from typing import List

from transformers import PreTrainedTokenizer

from ..core.constants import BATCH_MUST_HAVE_KEYS
from ..types import Batch, RawSample
from ..utils.miscellaneous import have_missing_keys


class BaseCollator(ABC):
    _BATCH_MUST_HAVE_KEYS = BATCH_MUST_HAVE_KEYS

    def __init__(self, tokenizer: PreTrainedTokenizer, max_length: int, separator: str = "\n"):
        super().__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.separator = separator

    def __call__(self, raw_batch: List[RawSample]) -> Batch:
        batch = self.parse_batch(raw_batch=raw_batch)
        flag, difference = have_missing_keys(data=batch, must_have_keys=self._BATCH_MUST_HAVE_KEYS)
        if flag:
            raise ValueError(f"Batch from {self.__class__.__name__} must have {difference} keys")
        return batch

    @abstractmethod
    def parse_batch(self, raw_batch: List[RawSample]) -> Batch:
        raise NotImplementedError

src/xllm/collators/base.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import List, Optional, Tuple

import torch
from transformers import PreTrainedTokenizer

from .. import enums
from ..collators.base import BaseCollator
from ..types import Batch, RawSample


class CompletionCollator(BaseCollator):
    def __init__(
        self,
        tokenizer: PreTrainedTokenizer,
        max_length: int,
        prefix_end: Optional[str] = None,
        separator: str = "\n",
    ):
        super().__init__(tokenizer=tokenizer, max_length=max_length, separator=separator)

        self.prefix_end = prefix_end

    def parse_sample(self, sample: List[str]) -> Tuple[List[int], List[int]]:
        text_parts = [text + self.separator for text in sample]
        tokenized = self.tokenizer(text_parts)[enums.Transformers.input_ids]

        token_indices = list()
        mask = list()

        if self.prefix_end and self.prefix_end in text_parts[-1]:
            prefix = text_parts[-1][: text_parts[-1].index(self.prefix_end) + len(self.prefix_end)]
            n_prefix = len(self.tokenizer(prefix)[enums.Transformers.input_ids]) - 1
        else:
            n_prefix = 0

        for n_sample, text_indices in enumerate(tokenized):
            if n_sample > 0:
                text_indices = text_indices[1:]
            if n_sample == len(tokenized) - 1:
                sample_mask = [0] * n_prefix + [1] * (len(text_indices) - n_prefix)
            else:
                sample_mask = [0] * len(text_indices)

            token_indices.extend(text_indices)
            mask.extend(sample_mask)

        input_ids = token_indices[:-1]

        targets = list()

        for token_index, flag in zip(token_indices[1:], mask[1:]):
            if flag:
                targets.append(token_index)
            else:
                targets.append(self.tokenizer.pad_token_id)

        return input_ids, targets

    def parse_batch(self, raw_batch: List[RawSample]) -> Batch:
        input_ids = list()
        targets = list()
        attention_masks = list()

        batch_max_length = 0

        for sample in raw_batch:
            text_parts = sample[enums.General.text_parts]

            if isinstance(text_parts, list):
                sample_input_ids, sample_targets = self.parse_sample(sample=[str(item) for item in text_parts])
                sample_length = len(sample_input_ids)

                input_ids.append(sample_input_ids)
                targets.append(sample_targets)

                if sample_length > batch_max_length:
                    batch_max_length = sample_length

                attention_masks.append([1] * sample_length)

        for n_sample in range(len(input_ids)):
            pad_sequence = [self.tokenizer.pad_token_id] * (batch_max_length - len(input_ids[n_sample]))
            if pad_sequence:
                additional_attention_mask = [0] * len(pad_sequence)
                if self.tokenizer.padding_side == "left":
                    input_ids[n_sample] = pad_sequence + input_ids[n_sample]
                    targets[n_sample] = pad_sequence + targets[n_sample]
                    attention_masks[n_sample] = additional_attention_mask + attention_masks[n_sample]
                else:
                    input_ids[n_sample] += pad_sequence
                    targets[n_sample] += pad_sequence
                    attention_masks[n_sample] += additional_attention_mask

        batch = {
            enums.Transformers.input_ids: torch.tensor(input_ids),
            enums.Transformers.attention_mask: torch.tensor(attention_masks),
            enums.Transformers.labels: torch.tensor(targets),
        }

        return batch

src/xllm/collators/completion.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import List

from .. import enums
from ..collators.base import BaseCollator
from ..types import Batch, RawSample


class LMCollator(BaseCollator):
    def parse_batch(self, raw_batch: List[RawSample]) -> Batch:
        texts = list()

        for sample in raw_batch:
            item = sample[enums.General.text_parts]
            if isinstance(item, (str, int, float)):
                texts.append(self.separator.join(str(item)))
            else:
                texts.append(self.separator.join(str(i) for i in item))

        tokenized = self.tokenizer(
            texts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.max_length,
        )

        batch = {
            enums.Transformers.input_ids: tokenized.input_ids[:, :-1],
            enums.Transformers.attention_mask: tokenized.attention_mask[:, :-1],
            enums.Transformers.labels: tokenized.input_ids[:, 1:],
        }

        return batch

src/xllm/collators/lm.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .. import enums
from ..collators.completion import CompletionCollator
from ..collators.lm import LMCollator
from ..utils.registry import Registry

collators_registry = Registry(name=enums.Registry.collators, default_value=LMCollator)

collators_registry.add(key=enums.Collators.lm, value=LMCollator)
collators_registry.add(key=enums.Collators.completion, value=CompletionCollator)

src/xllm/collators/registry.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: F401

from ..utils.logger import dist_logger
from ..utils.post_training import fuse_lora
from .config import HuggingFaceConfig

src/xllm/core/__init__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import os
from dataclasses import dataclass, field
from importlib.util import find_spec
from typing import Any, Dict, List, Optional, Union

import torch
from transformers.trainer_utils import FSDPOption

from .. import enums
from ..core.deepspeed_configs import DS_CONFIG_MAPPER
from ..utils.logger import dist_logger


@dataclass
class HuggingFaceConfig:
    # general
    experiment_key: str = field(
        default=enums.Experiments.base,
        metadata={"help": "Experiment class key"},
    )
    save_safetensors: bool = field(
        default=True,
        metadata={
            "help": "Safe serialization",
        },
    )
    max_shard_size: str = field(
        default="10GB", metadata={"help": "max_shard_size for the model pushing to the HuggingFace Hub"}
    )
    local_rank: int = field(
        default=0,
        metadata={
            "help": "Local rank for logging and saving. Works only in distributed training",
        },
    )
    use_gradient_checkpointing: bool = field(
        default=False,
        metadata={
            "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass",
        },
    )
    trainer_key: str = field(
        default=enums.Trainers.lm,
        metadata={
            "help": "Key of the trainer for loading from trainers_registry",
        },
    )
    force_fp32: bool = field(
        default=False,
        metadata={
            "help": "Force using fp32",
        },
    )
    force_fp16: bool = field(
        default=False,
        metadata={
            "help": "Force using fp16",
        },
    )
    from_gptq: bool = field(
        default=False,
        metadata={
            "help": "If you loadining GPTQ quantized model",
        },
    )
    huggingface_hub_token: Optional[str] = field(
        default=None,
        metadata={
            "help": "HuggingFace Hub token. You can also set this key using .env file",
        },
    )
    deepspeed_stage: Union[int, str, None] = field(
        default=0,
        metadata={
            "help": "DeepSpeed stage",
        },
    )
    deepspeed_config_path: Optional[int] = field(
        default=None,
        metadata={
            "help": "Path to DeepSpeed config",
        },
    )
    fsdp_strategy: str = field(
        default="",
        metadata={
            "help": "FSDP strategy",
        },
    )
    fsdp_offload: bool = field(
        default=True,
        metadata={
            "help": "Offload weights when using FSDP",
        },
    )
    seed: int = field(
        default=42,
        metadata={
            "help": "Seed value for random operations",
        },
    )
    stabilize: bool = field(
        default=False,
        metadata={
            "help": "Stabilize the model. Convert some weights to fp32, some to fp16/bf16",
        },
    )
    path_to_env_file: Optional[str] = field(
        default=None,
        metadata={"help": "Custom path to .env file"},
    )

    # download
    prepare_dataset: bool = field(
        default=True,
        metadata={
            "help": "Prepare the dataset. Works only when you use download",
        },
    )

    # fuse
    lora_hub_model_id: Optional[str] = field(
        default=None,
        metadata={
            "help": "Fusing LoRA. The name of the LoRA model at the hub for fusing. Example: KompleteAI/Shurale",
        },
    )
    lora_model_local_path: Optional[str] = field(
        default=None,
        metadata={
            "help": "Fusing LoRA. Local path to the LoRA model",
        },
    )
    fused_model_local_path: Optional[str] = field(
        default=None,
        metadata={
            "help": "Local path to fused model. Useful if you want to quantize model after fusing on the same machine",
        },
    )

    # gptq quantization
    quantization_dataset_id: Optional[str] = field(
        default=None,
        metadata={
            "help": "Dataset id for GPTQ quantization. You can install either the idi dataset, or use any dataset",
        },
    )
    quantization_max_samples: int = field(
        default=1024,
        metadata={
            "help": "Max samples for GPTQ quantization if you use custom dataset",
        },
    )
    quantized_model_path: str = field(
        default="./quantized_model/",
        metadata={
            "help": "Path to GPTQ quantized model",
        },
    )
    quantized_hub_model_id: Optional[str] = field(
        default=None,
        metadata={
            "help": "The name of the model at the hub for GPTQ quantization. Example: KompleteAI/Shurale-GPTQ",
        },
    )
    quantized_hub_private_repo: Optional[str] = field(
        default=None,
        metadata={
            "help": "Private repository for GPTQ quantization model or not",
        },
    )

    # dataset
    dataset_key: str = field(
        default=enums.Datasets.soda,
        metadata={
            "help": "Key of the dataset for loading from datasets_registry",
        },
    )
    train_local_path_to_data: str = field(
        default="./train.jsonl",
        metadata={
            "help": "The path to the local training data file",
        },
    )
    eval_local_path_to_data: Optional[str] = field(
        default=None,
        metadata={
            "help": "The path to the local eval data file",
        },
    )
    shuffle: bool = field(
        default=True,
        metadata={
            "help": "Shuffle training data",
        },
    )
    max_eval_samples: int = field(
        default=1_000,
        metadata={
            "help": "Maximum number of examples for evaluation",
        },
    )
    add_eval_to_train_if_no_path: bool = field(
        default=False,
        metadata={
            "help": "Add evaluation data to training data if their number is greater than max_eval_samples",
        },
    )

    # tokenizer
    tokenizer_name_or_path: Optional[str] = field(
        default=None,
        metadata={
            "help": "Tokenizer name or path. If the value is not set, "
            "then it will be taken from the model_name_or_path",
        },
    )
    tokenizer_use_fast: Optional[bool] = field(
        default=None,
        metadata={
            "help": "Use fast flag for the tokenizer",
        },
    )
    tokenizer_padding_side: Optional[str] = field(
        default=None,
        metadata={
            "help": "Padding side of the collator: right or left",
        },
    )

    # collator
    collator_key: str = field(
        default=enums.Collators.lm,
        metadata={
            "help": "Key of the collator for loading from collators_registry",
        },
    )
    max_length: int = field(
        default=2048,
        metadata={
            "help": "Max length of the model",
        },
    )

    # model
    model_name_or_path: str = field(
        default="meta-llama/Llama-2-7b-hf",
        metadata={
            "help": "Model name or path. It could be from HuggingFace or locally",
        },
    )
    push_to_hub_bos_add_bos_token: bool = field(
        default=False,
        metadata={
            "help": "Upload to the hub tokenization config with add_bos_token equals to True. Might be helpful for TGI"
        },
    )
    use_flash_attention_2: bool = field(
        default=False,
        metadata={
            "help": "Use or not flash attention 2. Requires 1) CUDA >= 11.6; 2) install flash-attn 3) compatible model",
        },
    )
    low_cpu_mem_usage: Optional[bool] = field(
        default=None,
        metadata={"help": "low_cpu_mem_usage when loading model"},
    )
    trust_remote_code: bool = field(
        default=True,
        metadata={
            "help": "Trust remote code from HuggingFace",
        },
    )
    device_map: Optional[str] = field(
        default=None,
        metadata={
            "help": "Device map for loading the model",
        },
    )
    prepare_model_for_kbit_training: bool = field(
        default=True,
        metadata={
            "help": "Prepare or not for kbit training",
        },
    )

    # bitsandbytes
    load_in_8bit: bool = field(
        default=False,
        metadata={
            "help": "Load the model in 8 bit using bitsandbytes",
        },
    )
    load_in_4bit: bool = field(
        default=False,
        metadata={
            "help": "Load the model in 4 bit using bitsandbytes",
        },
    )
    llm_int8_threshold: float = field(
        default=6.0,
        metadata={
            "help": "Threshold for outlier detection",
        },
    )
    llm_int8_has_fp16_weight: bool = field(
        default=True,
        metadata={
            "help": "LLM has weights in fp16",
        },
    )
    bnb_4bit_use_double_quant: bool = field(
        default=True,
        metadata={
            "help": "Double quantization. "
            "This will enable a second quantization after the first "
            "one to save an additional 0.4 bits per parameter",
        },
    )
    bnb_4bit_quant_type: str = field(
        default="nf4",
        metadata={
            "help": "Quantization type for 4 bit",
        },
    )
    bnb_quantize_after_model_init: bool = field(
        default=False, metadata={"help": "If False, quantization will be at model init"}
    )

    # gptq
    gptq_bits: int = field(
        default=4,
        metadata={
            "help": "Bits for GPTQ quantization",
        },
    )
    gptq_group_size: int = field(
        default=128,
        metadata={
            "help": "Group size for GPTQ quantization",
        },
    )
    gptq_disable_exllama: bool = field(
        default=True,
        metadata={
            "help": "Disable ExLlama kernels for GPTQ quantization",
        },
    )

    # lora
    apply_lora: bool = field(
        default=False,
        metadata={
            "help": "Apply LoRA to the model or not",
        },
    )
    lora_rank: int = field(
        default=64,
        metadata={
            "help": "LoRA rank value. LoRA matrices W_A x R and R x W_B, where R is LoRA rank",
        },
    )
    lora_alpha: int = field(
        default=32,
        metadata={
            "help": "LoRA alpha value. The resulting LoRA matrix will be multiplied by this value",
        },
    )
    lora_dropout: float = field(
        default=0.1,
        metadata={
            "help": "LoRA dropout value",
        },
    )
    raw_lora_target_modules: str = field(
        default="all",
        metadata={
            "help": 'Names of modules to apply LoRA. A comma-separated string, for example: "k,q,v". '
            'When setting the value "all", LoRA will be applied to all linear layers, except for the '
            "input embeddings and the lm_head.",
        },
    )

    # training arguments
    output_dir: str = field(
        default="./outputs/",
        metadata={
            "help": "The path to the directory where the artifacts will be saved",
        },
    )
    per_device_train_batch_size: int = field(
        default=2,
        metadata={
            "help": "Batch size on each GPU",
        },
    )
    do_eval: bool = field(
        default=False,
        metadata={
            "help": "Run eval or not",
        },
    )
    per_device_eval_batch_size: Optional[int] = field(
        default=None,
        metadata={
            "help": "Batch size on each GPU for evaluation. "
            "If None per_device_eval_batch_size equals to per_device_train_batch_size",
        },
    )
    gradient_accumulation_steps: int = field(
        default=1,
        metadata={
            "help": "Number of steps to accumulate gradients",
        },
    )
    eval_accumulation_steps: Optional[int] = field(
        default=None,
        metadata={
            "help": "Number of steps to accumulate gradients at evaluation."
            "If None eval_accumulation_steps equals to gradient_accumulation_steps",
        },
    )
    eval_delay: float = field(
        default=0,
        metadata={
            "helps": "Number of epochs or steps to wait for before the first "
            "evaluation can be performed, depending on the evaluation_strategy"
        },
    )
    eval_steps: Union[int, float, None] = field(
        default=1000, metadata={"helps": "Number of update steps between two evaluations"}
    )
    warmup_steps: int = field(
        default=1000,
        metadata={
            "help": "Number of steps to warm up",
        },
    )
    max_steps: Optional[int] = field(
        default=None,
        metadata={
            "help": "Maximum number of training steps",
        },
    )
    num_train_epochs: int = field(
        default=1,
        metadata={
            "help": "Number of training epochs",
        },
    )
    learning_rate: float = field(
        default=2e-4,
        metadata={
            "help": "Learning rate value",
        },
    )
    max_grad_norm: float = field(
        default=1.0,
        metadata={
            "help": "Clip grad value",
        },
    )
    weight_decay: float = field(
        default=0.001,
        metadata={
            "help": "Weight decay value",
        },
    )
    label_smoothing_factor: float = field(
        default=0.0,
        metadata={
            "help": "Label smoothing value",
        },
    )
    logging_steps: int = field(
        default=10,
        metadata={
            "help": "Number of steps between logging",
        },
    )
    save_steps: int = field(
        default=100,
        metadata={
            "help": "The number of training steps between saving the checkpoint and uploading to the hub",
        },
    )
    save_total_limit: int = field(
        default=1,
        metadata={
            "help": "The number of checkpoints that are saved locally",
        },
    )
    optim: Optional[str] = field(
        default="paged_adamw_8bit",
        metadata={
            "help": "Optimizer name. It will be overwritten if you use deepspeed",
        },
    )
    push_to_hub: bool = field(
        default=False,
        metadata={
            "help": "Upload the model to the hub. "
            "The model will be uploaded to the hub every save_steps. "
            "If LoRA is used, then LoRA's weights will be loaded onto the hub.",
        },
    )
    hub_model_id: Optional[str] = field(
        default=None,
        metadata={
            "help": "The name of the model at the hub. Example: KompleteAI/Shurale",
        },
    )
    hub_private_repo: bool = field(
        default=True,
        metadata={
            "help": "Private repository or not",
        },
    )

    # wandb
    report_to_wandb: bool = field(
        default=False,
        metadata={
            "help": "Report or not to Weight & Biases",
        },
    )
    wandb_api_key: Optional[str] = field(
        default=None,
        metadata={
            "help": "Weight & Biases API key. You can also set this key using .env file",
        },
    )
    wandb_project: Optional[str] = field(
        default=None,
        metadata={
            "help": "Weight & Biases project name",
        },
    )
    wandb_entity: Optional[str] = field(
        default=None,
        metadata={
            "help": "Weight & Biases entity name (user or company)",
        },
    )

    def __post_init__(self):
        if self.huggingface_hub_token is not None:
            os.environ[enums.EnvironmentVariables.huggingface_hub_token] = self.huggingface_hub_token
            dist_logger(message=f"Environment variable {enums.EnvironmentVariables.huggingface_hub_token} set")

        if self.report_to_wandb:
            for key, value in zip(
                [
                    enums.EnvironmentVariables.wandb_api_key,
                    enums.EnvironmentVariables.wandb_project,
                    enums.EnvironmentVariables.wandb_entity,
                ],
                [
                    self.wandb_api_key,
                    self.wandb_project,
                    self.wandb_entity,
                ],
            ):
                if value is not None:
                    os.environ[key] = value
                    dist_logger(message=f"Environment variable {key} set")
        else:
            os.environ[enums.EnvironmentVariables.wandb_disabled] = "true"

    def check_hub(self) -> None:
        if self.push_to_hub and self.hub_model_id is None:
            raise ValueError("You want to push to HF hub, but hub_model_id is None")
        elif self.hub_model_id is not None and not self.push_to_hub:
            dist_logger.warning("You set hub_model_id, but push_to_hub is False")

        return None

    def check_deepspeed(self) -> None:
        if self.deepspeed is not None:
            spec = find_spec("deepspeed")

            if spec is None:
                raise ImportError("Deepspeed is not None, but failed to import deepspeed. Please install deepspeed.")

        return None

    def check_flash_attention(self) -> None:
        if self.use_flash_attention_2:
            if not torch.cuda.is_available():
                raise ImportError("You want to use flash_attention_2, but CUDA is not available")

            spec = find_spec("flash_attn")

            if spec is None:
                raise ImportError(
                    "You want to use flash_attention_2, but flash-attn is not installed. Please install flash-attn."
                )

        return None

    def check_auto_gptq(self) -> None:
        spec = find_spec("auto_gptq")

        if spec is None:
            raise ImportError(
                "You want to quantize model using GPTQ, but auto-gptq is not installed. Please install auto-gptq."
            )

        return None

    def check(self) -> None:
        self.check_hub()
        self.check_deepspeed()
        self.check_flash_attention()

        return None

    @property
    def correct_tokenizer_name_or_path(self) -> str:
        if self.tokenizer_name_or_path is not None:
            return self.tokenizer_name_or_path
        else:
            return self.model_name_or_path

    @property
    def lora_target_modules(self) -> Optional[List[str]]:
        if self.raw_lora_target_modules == "all":
            return None
        elif self.raw_lora_target_modules is not None:
            modules_names = [module_name.strip() for module_name in self.raw_lora_target_modules.split(",")]
            return modules_names
        else:
            raise ValueError("raw_lora_target_modules doesn't set")

    @property
    def dtype(self) -> torch.dtype:
        if not torch.cuda.is_available() or self.force_fp32:
            return torch.float32
        elif self.force_fp16:
            return torch.float16
        elif torch.cuda.is_bf16_supported():
            return torch.bfloat16
        else:
            return torch.float16

    @property
    def deepspeed(self) -> Optional[Dict[str, Any]]:
        if self.deepspeed_stage in [0, "0", "stage_0"]:
            return None

        deepspeed_config: Optional[Dict[str, Any]] = None

        if self.deepspeed_stage is not None:
            deepspeed_config = DS_CONFIG_MAPPER.get(self.deepspeed_stage, None)
            if deepspeed_config is None:
                raise ValueError(
                    f'Deepspeed stage "{self.deepspeed_stage}" not found in keys: {list(DS_CONFIG_MAPPER.keys())}'
                )

        if self.deepspeed_config_path is not None:
            if os.path.isfile(self.deepspeed_config_path):
                with open(self.deepspeed_config_path) as file_object:
                    deepspeed_config = json.load(file_object)
            else:
                raise ValueError(f"deepspeed_config_path set to {self.deepspeed_config_path}, but not found")

        return deepspeed_config

    @property
    def fsdp(self) -> Union[str, List[str]]:
        fsdp_options = list()

        if self.fsdp_strategy is not None and self.fsdp_strategy != "":
            fsdp_options.append(self.fsdp_strategy)
        else:
            return ""

        if self.fsdp_offload:
            fsdp_options.append(FSDPOption.OFFLOAD)

        return fsdp_options

    @property
    def lora_model_name_or_path_for_fusing(self) -> str:
        if self.lora_hub_model_id is not None:
            return self.lora_hub_model_id
        elif self.lora_model_local_path is not None:
            return self.lora_model_local_path
        else:
            raise ValueError("Please set lora_hub_model_id or lora_model_local_path for fusing")

src/xllm/core/config.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .. import enums

DATASETS_MUST_HAVE_KEYS = [enums.General.text_parts]
BATCH_MUST_HAVE_KEYS = [
    enums.Transformers.input_ids,
    enums.Transformers.attention_mask,
    enums.Transformers.labels,
]

src/xllm/core/constants.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Dict, Union

STAGE_1 = {
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1,
    },
    "bf16": {"enabled": "auto"},
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto",
        },
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
        },
    },
    "zero_optimization": {
        "stage": 1,
        "offload_optimizer": {"device": "cpu", "pin_memory": False},
        "allgather_partitions": True,
        "allgather_bucket_size": 200000000.0,
        "overlap_comm": True,
        "reduce_scatter": True,
        "reduce_bucket_size": 200000000.0,
        "contiguous_gradients": True,
        "round_robin_gradients": True,
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": False,
}

STAGE_2 = {
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1,
    },
    "bf16": {"enabled": "auto"},
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto",
        },
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "total_num_steps": "auto",
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
        },
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {"device": "cpu", "pin_memory": False},
        "allgather_partitions": True,
        "allgather_bucket_size": 200000000.0,
        "overlap_comm": True,
        "reduce_scatter": True,
        "reduce_bucket_size": 200000000.0,
        "contiguous_gradients": True,
        "round_robin_gradients": True,
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": False,
}

STAGE_3 = {
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1,
    },
    "bf16": {"enabled": "auto"},
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto",
        },
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "total_num_steps": "auto",
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
        },
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "cpu", "pin_memory": False},
        "offload_param": {"device": "cpu", "pin_memory": False},
        "overlap_comm": True,
        "contiguous_gradients": True,
        "sub_group_size": 1000000000.0,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1000000000.0,
        "stage3_max_reuse_distance": 1000000000.0,
        "gather_16bit_weights_on_model_save": True,
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": False,
}

DS_CONFIG_MAPPER: Dict[Union[str, int], Dict[str, Any]] = {
    1: STAGE_1,
    "1": STAGE_1,
    "stage_1": STAGE_1,
    2: STAGE_2,
    "2": STAGE_2,
    "stage_2": STAGE_2,
    3: STAGE_3,
    "3": STAGE_3,
    "stage_3": STAGE_3,
}

src/xllm/core/deepspeed_configs.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from typing import Any, Dict, Optional, Type, Union

import torch
from peft import (  # type: ignore
    PeftModel,
    prepare_model_for_kbit_training,
)
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    GPTQConfig,
    IntervalStrategy,
    PreTrainedModel,
    PreTrainedTokenizer,
    Trainer,
    TrainingArguments,
)

from .. import enums
from ..collators.base import BaseCollator
from ..collators.registry import collators_registry
from ..core.config import HuggingFaceConfig
from ..datasets.base import BaseDataset
from ..datasets.registry import datasets_registry
from ..trainers.lm import LMTrainer
from ..trainers.registry import trainers_registry
from ..utils.logger import dist_logger


def build_training_arguments(config: HuggingFaceConfig) -> TrainingArguments:
    if torch.cuda.is_available():
        if torch.cuda.is_bf16_supported() and not config.force_fp16:
            fp16 = False
            bf16 = True
        else:
            fp16 = True
            bf16 = False
    else:
        fp16 = False
        bf16 = False

    training_arguments = TrainingArguments(
        output_dir=config.output_dir,
        per_device_train_batch_size=config.per_device_train_batch_size,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        warmup_steps=config.warmup_steps,
        learning_rate=config.learning_rate,
        max_steps=config.max_steps if config.max_steps is not None else -1,
        num_train_epochs=config.num_train_epochs,
        weight_decay=config.weight_decay,
        max_grad_norm=config.max_grad_norm,
        label_smoothing_factor=config.label_smoothing_factor,
        fp16=fp16,
        bf16=bf16,
        logging_steps=config.logging_steps,
        report_to=["wandb"] if config.report_to_wandb else None,
        save_strategy="steps",
        save_steps=config.save_steps,
        save_total_limit=config.save_total_limit,
        hub_model_id=config.hub_model_id,
        hub_strategy="checkpoint",
        hub_token=os.environ.get(enums.EnvironmentVariables.huggingface_hub_token, None),
        push_to_hub=config.push_to_hub,
        hub_private_repo=config.hub_private_repo,
        save_safetensors=config.save_safetensors,
        fsdp=config.fsdp,
        deepspeed=config.deepspeed,
        remove_unused_columns=False,
        log_level=enums.LogLevel.info,
        disable_tqdm=False,
        logging_first_step=True,
        optim=config.optim,  # will be overwriten by deepspeed config if exist
        do_eval=config.do_eval,
        evaluation_strategy="steps" if config.do_eval else IntervalStrategy.NO,
        per_device_eval_batch_size=config.per_device_eval_batch_size or config.per_device_train_batch_size,
        eval_accumulation_steps=config.eval_accumulation_steps or config.gradient_accumulation_steps,
        eval_delay=config.eval_delay,
        eval_steps=config.eval_steps,
        seed=config.seed,
        data_seed=config.seed,
        metric_for_best_model="eval_loss" if config.do_eval else "loss",
    )
    return training_arguments


def build_dataset(config: HuggingFaceConfig, is_train: bool = True, **kwargs: Any) -> Optional[BaseDataset]:
    if is_train:
        path_to_data = config.train_local_path_to_data
    elif config.eval_local_path_to_data is not None:
        path_to_data = config.eval_local_path_to_data
    else:
        return None

    dataset_cls: Type[BaseDataset] = datasets_registry.get(key=config.dataset_key)

    if issubclass(dataset_cls, BaseDataset):
        dataset = dataset_cls.load(path_to_data=path_to_data, **kwargs)
    else:
        raise ValueError(f"Unknown type of dataset: {dataset_cls.__name__}")

    return dataset


def build_tokenizer(config: HuggingFaceConfig, use_fast: Optional[bool] = None) -> PreTrainedTokenizer:
    kwargs = dict()

    if use_fast is not None:
        kwargs["use_fast"] = use_fast

    tokenizer = AutoTokenizer.from_pretrained(
        pretrained_model_name_or_path=config.correct_tokenizer_name_or_path, **kwargs
    )

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        dist_logger.info(message="Tokenizer pad token set to eos token", local_rank=config.local_rank)

    if config.tokenizer_padding_side is not None:
        tokenizer.padding_side = config.tokenizer_padding_side
        dist_logger.info(
            message=f"Tokenizer padding side set to {config.tokenizer_padding_side}", local_rank=config.local_rank
        )

    return tokenizer


def build_collator(config: HuggingFaceConfig, tokenizer: PreTrainedTokenizer, **kwargs: Any) -> BaseCollator:
    collator_cls: Type[BaseCollator] = collators_registry.get(key=config.collator_key)

    if not issubclass(collator_cls, BaseCollator):
        raise ValueError(f"Unknown type of collator: {collator_cls.__name__}")

    collator = collator_cls(tokenizer=tokenizer, max_length=config.max_length, **kwargs)

    return collator


def build_quantization_config(
    config: HuggingFaceConfig,
) -> Union[BitsAndBytesConfig, GPTQConfig, None]:
    if config.from_gptq:
        quantization_config = GPTQConfig(
            bits=config.gptq_bits,
            group_size=config.gptq_group_size,
            disable_exllama=config.gptq_disable_exllama,
        )
    elif config.load_in_8bit or config.load_in_4bit:
        quantization_config = BitsAndBytesConfig(
            load_in_8bit=config.load_in_8bit,
            load_in_4bit=config.load_in_4bit,
            llm_int8_threshold=config.llm_int8_threshold,
            llm_int8_has_fp16_weight=config.llm_int8_has_fp16_weight,
            bnb_4bit_compute_dtype=config.dtype,
            bnb_4bit_use_double_quant=config.bnb_4bit_use_double_quant,
            bnb_4bit_quant_type=config.bnb_4bit_quant_type,
        )
    else:
        quantization_config = None

    return quantization_config


def build_model(
    config: HuggingFaceConfig,
    quantization_config: Union[BitsAndBytesConfig, GPTQConfig, None],
) -> PreTrainedModel:
    if config.bnb_quantize_after_model_init:
        quantization_config = None
        dist_logger("bnb quantization is expected later")

    if config.use_gradient_checkpointing:
        use_cache = False
    else:
        use_cache = True

    kwargs: Dict[str, Any] = dict()

    if config.use_flash_attention_2:
        kwargs["use_flash_attention_2"] = True

    if config.low_cpu_mem_usage is not None:
        kwargs["low_cpu_mem_usage"] = config.low_cpu_mem_usage

    model = AutoModelForCausalLM.from_pretrained(
        pretrained_model_name_or_path=config.model_name_or_path,
        quantization_config=quantization_config,
        torch_dtype=config.dtype,
        trust_remote_code=config.trust_remote_code,
        device_map=config.device_map,
        use_cache=use_cache,
        **kwargs,
    )
    model.config.pretraining_tp = 1

    if quantization_config is not None and config.prepare_model_for_kbit_training:
        model = prepare_model_for_kbit_training(
            model=model, use_gradient_checkpointing=config.use_gradient_checkpointing
        )

    return model


def build_trainer(
    config: HuggingFaceConfig,
    pad_token_id: int,
    training_arguments: TrainingArguments,
    model: Union[PreTrainedModel, PeftModel],
    train_dataset: BaseDataset,
    collator: BaseCollator,
    eval_dataset: Optional[BaseDataset] = None,
    **kwargs: Any,
) -> LMTrainer:
    trainer_cls = trainers_registry.get(key=config.trainer_key)

    if not issubclass(trainer_cls, Trainer):
        raise ValueError(f"Unknown type of trainer: {trainer_cls.__name__}")

    trainer: LMTrainer = trainer_cls(
        config=config,
        model=model,
        args=training_arguments,
        data_collator=collator,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        ignore_index=pad_token_id,
        **kwargs,
    )

    try:
        model.config.use_cache = False  # type: ignore
    except Exception as exception:
        dist_logger.warning(
            message=f"Can't set use cache to false, because: {exception}",
            local_rank=config.local_rank,
        )

    return trainer

src/xllm/core/dependencies.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: F401

from .base import BaseDataset
from .general import GeneralDataset
from .registry import datasets_registry
from .soda import SodaDataset

src/xllm/datasets/__init__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import os
import random
from abc import ABC, abstractmethod
from typing import Any, List, Optional, Tuple

from loguru import logger
from torch.utils.data import Dataset

from ..core.config import HuggingFaceConfig
from ..core.constants import DATASETS_MUST_HAVE_KEYS
from ..types import RawSample
from ..utils.miscellaneous import have_missing_keys


class BaseDataset(Dataset[RawSample], ABC):
    _DATASET_MUST_HAVE_KEYS = DATASETS_MUST_HAVE_KEYS

    def __init__(self, data: List[RawSample]):
        super().__init__()
        self.data = data

    @classmethod
    def prepare(cls, config: HuggingFaceConfig) -> None:
        download_result = cls.download(config=config)

        if download_result is None:
            raise ValueError("download method returned None")
        else:
            train_data, eval_data = download_result

        if config.eval_local_path_to_data is None and eval_data is not None:
            logger.warning("eval_local_path_to_data is None, but eval_data is not None")
            if config.add_eval_to_train_if_no_path:
                train_data += eval_data
                logger.info("Add eval data to train")

        if eval_data is not None and config.eval_local_path_to_data is not None:
            if len(eval_data) > config.max_eval_samples and config.max_eval_samples > 0:
                train_data += eval_data[config.max_eval_samples :]
                eval_data = eval_data[: config.max_eval_samples]
                logger.info(f"Eval data size truncated to {config.max_eval_samples}")
            else:
                logger.info(f"Eval data size: {len(eval_data)}")

            with open(config.eval_local_path_to_data, mode="w") as file_object:
                for raw_sample in eval_data:
                    file_object.write(json.dumps(raw_sample) + "\n")
        else:
            logger.warning("eval data or eval_local_path_to_data is None")

        if config.shuffle:
            random.shuffle(train_data)
            logger.info("Train data shuffled")

        with open(config.train_local_path_to_data, mode="w") as file_object:
            for raw_sample in train_data:
                file_object.write(json.dumps(raw_sample) + "\n")

        logger.info(f"Train data size: {len(train_data)}")

        return None

    @classmethod
    def load(cls, path_to_data: str, **kwargs: Any) -> "BaseDataset":
        data = list()

        if not os.path.isfile(path_to_data):
            raise FileNotFoundError(
                f"File {path_to_data} not found. Probably you should run download_and_prepare before"
            )

        with open(path_to_data) as file_object:
            for line in file_object:
                sample = json.loads(line)
                data.append(sample)

        dataset = cls(data=data)

        return dataset

    def __len__(self) -> int:
        return len(self.data)

    def __getitem__(self, index: int) -> RawSample:
        sample = self.get_sample(index=index)
        flag, difference = have_missing_keys(data=sample, must_have_keys=self._DATASET_MUST_HAVE_KEYS)
        if flag:
            raise ValueError(f"RawSample from {self.__class__.__name__} must have {difference} keys")
        return sample

    @classmethod
    @abstractmethod
    def download(cls, config: HuggingFaceConfig) -> Optional[Tuple[List[RawSample], Optional[List[RawSample]]]]:
        raise NotImplementedError

    @abstractmethod
    def get_sample(self, index: int) -> RawSample:
        raise NotImplementedError

src/xllm/datasets/base.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import List, Optional, Tuple

from .. import enums
from ..core.config import HuggingFaceConfig
from ..datasets.base import BaseDataset
from ..types import RawSample
from ..utils.logger import dist_logger


class GeneralDataset(BaseDataset):
    def __init__(
        self,
        data: List[RawSample],
        sample_field: str = enums.General.default_sample_field,
        separator: Optional[str] = None,
    ):
        super().__init__(data=data)

        self.sample_field = sample_field
        self.separator = separator

    @classmethod
    def download(cls, config: HuggingFaceConfig) -> Optional[Tuple[List[RawSample], Optional[List[RawSample]]]]:
        dist_logger.warning(
            "This is a special type of dataset in which it is not supposed to download anything. "
            "You must pass the data here through __init__, "
            "or through the path in config.train_local_path_to_data and config.eval_local_path_to_data (optional)"
        )
        return None

    def get_sample(self, index: int) -> RawSample:
        text = self.data[index][self.sample_field]

        assert isinstance(text, str)

        if self.separator is not None:
            text_parts: List[str] = text.split(self.separator)
        else:
            text_parts = [text]

        sample: RawSample = {enums.General.text_parts: text_parts}

        return sample

src/xllm/datasets/general.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .. import enums
from ..datasets.general import GeneralDataset
from ..datasets.soda import SodaDataset
from ..utils.registry import Registry

datasets_registry = Registry(name=enums.Registry.datasets)

datasets_registry.add(key=enums.Datasets.general, value=GeneralDataset)
datasets_registry.add(key=enums.Datasets.soda, value=SodaDataset)

src/xllm/datasets/registry.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Dict, List, Optional, Tuple

import datasets
import numpy as np
from tqdm import tqdm

from .. import enums
from ..core.config import HuggingFaceConfig
from ..datasets.base import BaseDataset
from ..types import RawSample


class SodaDataset(BaseDataset):
    HEADER_KEY = "header"
    DIALOG_KEY = "dialog"

    _HF_DATASET_ID = "allenai/soda"

    def __init__(self, data: List[RawSample], header_drop_probability: float = 0.05):
        super().__init__(data=data)
        self.header_drop_probability = header_drop_probability

    @classmethod
    def download(cls, config: HuggingFaceConfig) -> Optional[Tuple[List[RawSample], Optional[List[RawSample]]]]:
        soda_dataset = datasets.load_dataset(cls._HF_DATASET_ID)

        parsed_data: Dict[str, List[RawSample]] = dict()

        known_indices = set()

        for split in ["train", "test"]:
            parsed_data[split] = list()

            for sample in tqdm(soda_dataset[split], desc=f"Parsing SODA {split}"):
                index = sample.get("original_index")

                if index in known_indices:
                    continue

                parsed_sample = {
                    cls.HEADER_KEY: sample.get("narrative"),
                    cls.DIALOG_KEY: [
                        f"{speaker}: {phrase}"
                        for speaker, phrase in zip(sample.get("speakers"), sample.get("dialogue"))
                    ],
                }

                parsed_data[split].append(parsed_sample)
                known_indices.add(index)

        train = parsed_data["train"]
        valid = parsed_data["test"]

        return train, valid

    def get_sample(self, index: int) -> RawSample:
        sample = self.data[index]

        dialog = sample[self.DIALOG_KEY]

        phrases = list()

        if not isinstance(dialog, list):
            raise ValueError(f"{self.DIALOG_KEY} of sample is not a list: {type(dialog)}")

        for phrase in dialog:
            if isinstance(phrase, str):
                phrases.append(phrase)

        if self.HEADER_KEY in sample:
            header = sample[self.HEADER_KEY]

            is_drop_header = np.random.rand() <= self.header_drop_probability

            if not is_drop_header and isinstance(header, str):
                phrases.insert(0, header)

        sample = {enums.General.text_parts: [phrase.replace("\n", " ").replace("\r", " ") for phrase in phrases]}

        return sample

src/xllm/datasets/soda.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass


@dataclass
class General:
    text_parts: str = "text_parts"
    default_sample_field: str = "text"


@dataclass
class Transformers:
    input_ids: str = "input_ids"
    attention_mask: str = "attention_mask"
    labels: str = "labels"
    logits: str = "logits"


@dataclass
class Registry:
    datasets: str = "datasets"
    collators: str = "collators"
    trainers: str = "trainers"
    experiments: str = "experiments"


@dataclass
class Datasets:
    general: str = "general"
    soda: str = "soda"


@dataclass
class Collators:
    lm: str = "lm"
    completion: str = "completion"


@dataclass
class Trainers:
    lm: str = "lm"


@dataclass
class Experiments:
    base: str = "base"


@dataclass
class EnvironmentVariables:
    huggingface_hub_token: str = "HUGGING_FACE_HUB_TOKEN"
    wandb_api_key: str = "WANDB_API_KEY"
    wandb_entity: str = "WANDB_ENTITY"
    wandb_project: str = "WANDB_PROJECT"
    wandb_disabled: str = "WANDB_DISABLED"
    tokenizers_parallelism: str = "TOKENIZERS_PARALLELISM"


@dataclass
class LogLevel:
    info: str = "info"
    warning: str = "warning"
    error: str = "error"
    critical: str = "critical"

src/xllm/enums.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: F401

from .base import Experiment
from .registry import experiments_registry

src/xllm/experiments/__init__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import os
from time import sleep
from typing import Any, Dict, Optional, Union

import torch
import torch.distributed as distributed
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training  # type: ignore
from transformers import (
    BitsAndBytesConfig,
    GPTQConfig,
    PreTrainedModel,
    PreTrainedTokenizer,
    TrainingArguments,
)
from transformers.integrations.bitsandbytes import replace_with_bnb_linear

from ..collators.base import BaseCollator
from ..core.config import HuggingFaceConfig
from ..core.dependencies import (
    build_collator,
    build_dataset,
    build_model,
    build_quantization_config,
    build_tokenizer,
    build_trainer,
    build_training_arguments,
)
from ..datasets.base import BaseDataset
from ..trainers.lm import LMTrainer
from ..utils.logger import dist_logger
from ..utils.miscellaneous import is_distributed_training
from ..utils.nn import apply_lora, stabilize_training
from ..utils.post_training import post_training, push_to_hub_bos_add_bos_token


class Experiment:
    def __init__(
        self,
        config: HuggingFaceConfig,
        training_arguments: Optional[TrainingArguments] = None,
        train_dataset: Optional[BaseDataset] = None,
        eval_dataset: Optional[BaseDataset] = None,
        tokenizer: Optional[PreTrainedTokenizer] = None,
        collator: Optional[BaseCollator] = None,
        quantization_config: Union[BitsAndBytesConfig, GPTQConfig, None] = None,
        model: Union[PreTrainedModel, PeftModel, None] = None,
        lora_config: Optional[LoraConfig] = None,
        trainer: Optional[LMTrainer] = None,
    ):
        self.config = config

        self.training_arguments = training_arguments
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.tokenizer = tokenizer
        self.collator = collator
        self.quantization_config = quantization_config
        self.model = model
        self.lora_config = lora_config
        self.trainer = trainer

        self.internal_checks()

    def internal_checks(self) -> None:
        if self.tokenizer is not None and self.collator is not None and self.tokenizer != self.collator.tokenizer:
            dist_logger.warning("Tokenizer not equals to tokenizer in collator")

        return None

    def internal_checks_before_train(self) -> None:
        if self.training_arguments is not None and self.training_arguments.do_eval and self.eval_dataset is None:
            raise ValueError(
                f"You set do_eval at config to {self.config.do_eval}, "
                "but experiment can't run eval, because eval_dataset is None. "
                f"config.eval_local_path_to_data: {self.config.eval_local_path_to_data}"
            )

    def build(self):
        dist_logger("Experiment building has started")
        self.at_beggining()
        self.save_config()
        dist_logger.info("Config saved")

        self.before_checks()
        self.checks()
        dist_logger("Checks passed successfully")
        self.after_checks()

        if self.training_arguments is None:
            self.before_training_arguments_build()
            self.training_arguments = self.build_training_arguments()
            dist_logger(f"Training arguments was built:\n{self.training_arguments.to_json_string()}")
            self.after_training_arguments_build()

        if self.train_dataset is None:
            self.before_train_dataset_build()
            self.train_dataset = self.build_train_dataset()
            dist_logger(
                f"Train dataset {self.train_dataset.__class__.__name__} was built. Size: {len(self.train_dataset)}"
            )
            self.after_train_dataset_build()

        if self.eval_dataset is None:
            self.before_eval_dataset_build()
            self.eval_dataset = self.build_eval_dataset()
            if self.eval_dataset is not None:
                dist_logger(
                    f"Valid dataset {self.eval_dataset.__class__.__name__} was built. Size: {len(self.eval_dataset)}"
                )
            else:
                dist_logger("Valid dataset is None")
            self.after_eval_dataset_build()

        if self.tokenizer is None:
            self.before_tokenizer_build()
            self.tokenizer = self.build_tokenizer()
            dist_logger(f"Tokenizer {self.config.correct_tokenizer_name_or_path} was built")
            self.after_tokenizer_build()

        if self.collator is None:
            self.before_collator_build()
            self.collator = self.build_collator()
            dist_logger(f"Collator {self.collator.__class__.__name__} was built")
            self.after_collator_build()

        if self.quantization_config is None:
            self.before_quantization_config_build()
            self.quantization_config = self.build_quantization_config()
            if self.quantization_config is not None:
                dist_logger(f"Quantization config was built:\n{self.quantization_config.to_json_string()}")
            else:
                dist_logger(f"Quantization config is None. Model will be loaded using {self.config.dtype}")
            self.after_quantization_config_build()

        if self.model is None:
            self.before_model_build()
            self.model = self.build_model()
            dist_logger(f"Model {self.config.model_name_or_path} was built")
            self.after_model_build()
        elif self.quantization_config is not None:
            dist_logger.warning("quantization_config is not None, but the model was built outside of the experiment")

        if self.is_bnb_quantization and self.config.bnb_quantize_after_model_init:
            self.before_bnb_quantization()
            self.bnb_quantization()
            bnb_quantization_type = "int4" if self.quantization_config.load_in_4bit else "int8"
            dist_logger(f"Bnb quantization applyed. Type: {bnb_quantization_type}")
            self.after_bnb_quantization()

        if self.config.apply_lora:
            self.before_lora_apply()
            self.lora_config = self.apply_lora()
            dist_logger(f"LoRA applied to the model {self.config.model_name_or_path}")
            self.after_lora_apply()

        if self.config.stabilize:
            self.before_stabilize_training()
            self.stabilize_training()
            dist_logger(f"Model {self.config.model_name_or_path} is stabilized for training")
            self.after_stabilize_training()

        if self.trainer is None:
            self.before_trainer_build()
            self.trainer = self.build_trainer()
            dist_logger(f"Trainer {self.trainer.__class__.__name__} was built")
            self.after_trainer_build()
        else:
            trainer_components = self.trainer_components
            trainer_components_is_not_none = [key for key, value in trainer_components.items() if value is not None]
            dist_logger.warning(
                "Trainer was built outside of the experiment, "
                f"but this components is not None: {', '.join(trainer_components_is_not_none)}."
            )

        self.internal_checks()

        dist_logger("Experiment built successfully")

    @property
    def is_bnb_quantization(self) -> bool:
        return isinstance(self.quantization_config, BitsAndBytesConfig) and self.config.from_gptq

    @property
    def trainer_components(self) -> Dict[str, Any]:
        components = {
            "training_arguments": self.training_arguments,
            "train_dataset": self.train_dataset,
            "eval_dataset": self.eval_dataset,
            "collator": self.collator,
            "model": self.model,
        }
        return components

    # checks
    def before_checks(self) -> None:
        return None

    def checks(self) -> None:
        if not torch.cuda.is_available():
            dist_logger.warning("CUDA is not available")

        self.config.check()

        return None

    def after_checks(self) -> None:
        return None

    # training arguments
    def before_training_arguments_build(self) -> None:
        return None

    def build_training_arguments(self) -> TrainingArguments:
        training_arguments = build_training_arguments(config=self.config)
        return training_arguments

    def after_training_arguments_build(self) -> None:
        return None

    # train_dataset
    def before_train_dataset_build(self) -> None:
        return None

    def build_train_dataset_additional_kwargs(self) -> Dict[str, Any]:
        return dict()

    def build_train_dataset(self) -> BaseDataset:
        train_dataset_additional_kwargs = self.build_train_dataset_additional_kwargs()
        dataset = build_dataset(config=self.config, is_train=True, **train_dataset_additional_kwargs)
        if dataset is None:
            raise ValueError("Train dataset can't be loaded")
        return dataset

    def after_train_dataset_build(self) -> None:
        return None

    # eval_dataset
    def before_eval_dataset_build(self) -> None:
        return None

    def build_eval_dataset_additional_kwargs(self) -> Dict[str, Any]:
        return self.build_train_dataset_additional_kwargs()

    def build_eval_dataset(self) -> Optional[BaseDataset]:
        eval_dataset_additional_kwargs = self.build_eval_dataset_additional_kwargs()
        dataset = build_dataset(config=self.config, is_train=False, **eval_dataset_additional_kwargs)
        return dataset

    def after_eval_dataset_build(self) -> None:
        return None

    # tokenizer
    def before_tokenizer_build(self) -> None:
        return None

    def build_tokenizer(self) -> PreTrainedTokenizer:
        tokenizer = build_tokenizer(config=self.config, use_fast=self.config.tokenizer_use_fast)
        return tokenizer

    def after_tokenizer_build(self) -> None:
        return None

    # collator
    def before_collator_build(self) -> None:
        return None

    def build_collator_additional_kwargs(self) -> Dict[str, Any]:
        return dict()

    def build_collator(self) -> BaseCollator:
        collator_additional_kwargs = self.build_collator_additional_kwargs()
        collator = build_collator(config=self.config, tokenizer=self.tokenizer, **collator_additional_kwargs)
        return collator

    def after_collator_build(self) -> None:
        return None

    # quantization_config
    def before_quantization_config_build(self) -> None:
        return None

    def build_quantization_config(self) -> Union[BitsAndBytesConfig, GPTQConfig, None]:
        quantization_config = build_quantization_config(config=self.config)
        return quantization_config

    def after_quantization_config_build(self) -> None:
        return None

    # model
    def before_model_build(self) -> None:
        return None

    def build_model(self) -> PreTrainedModel:
        quantization_config = (
            None if self.is_bnb_quantization and self.config.bnb_quantize_after_model_init else self.quantization_config
        )
        model = build_model(config=self.config, quantization_config=quantization_config)
        return model

    def after_model_build(self) -> None:
        return None

    # bnb_quantization
    def before_bnb_quantization(self) -> None:
        return None

    def bnb_quantization(self) -> None:
        self.model = replace_with_bnb_linear(
            model=self.model,
            quantization_config=self.quantization_config,
        )
        self.model.is_loaded_in_4bit = self.config.load_in_4bit
        self.model.is_loaded_in_8bit = self.config.load_in_8bit
        if self.config.prepare_model_for_kbit_training:
            self.model = prepare_model_for_kbit_training(
                model=self.model, use_gradient_checkpointing=self.config.use_gradient_checkpointing
            )

    def after_bnb_quantization(self) -> None:
        return None

    # lora
    def before_lora_apply(self) -> None:
        return None

    def apply_lora(self) -> LoraConfig:
        self.model, lora_config = apply_lora(config=self.config, model=self.model, lora_config=self.lora_config)
        return lora_config

    def after_lora_apply(self) -> None:
        return None

    # stabilize_training
    def before_stabilize_training(self) -> None:
        return None

    def stabilize_training(self) -> None:
        self.model = stabilize_training(model=self.model)

    def after_stabilize_training(self) -> None:
        return None

    # trainer
    def before_trainer_build(self) -> None:
        return None

    def build_trainer_additional_kwargs(self) -> Dict[str, Any]:
        return dict()

    def build_trainer(self) -> LMTrainer:
        additional_trainer_kwargs = self.build_trainer_additional_kwargs()

        if self.tokenizer is None:
            raise ValueError("tokenizer is None")

        if self.train_dataset is None:
            raise ValueError("train_dataset is None")

        if self.collator is None:
            raise ValueError("collator is None")

        trainer = build_trainer(
            config=self.config,
            pad_token_id=self.tokenizer.pad_token_id,
            training_arguments=self.training_arguments,
            model=self.model,
            train_dataset=self.train_dataset,
            collator=self.collator,
            eval_dataset=self.eval_dataset,
            **additional_trainer_kwargs,
        )

        return trainer

    def after_trainer_build(self) -> None:
        return None

    def save_config(self) -> None:
        json_config = json.dumps(self.config.__dict__, indent=2)
        dist_logger(f"Config:\n{json_config}")

        if is_distributed_training():
            if distributed.get_rank() == self.config.local_rank:
                os.makedirs(self.config.output_dir, exist_ok=True)
                with open(os.path.join(self.config.output_dir, "training_config.json"), "w") as file_object:
                    file_object.write(json_config)
        else:
            os.makedirs(self.config.output_dir, exist_ok=True)
            with open(os.path.join(self.config.output_dir, "training_config.json"), "w") as file_object:
                file_object.write(json_config)

        return None

    def before_train(self) -> None:
        return None

    def run(self):
        self.before_train()

        if self.trainer is None:
            raise ValueError("trainer is None")

        if self.training_arguments is None:
            raise ValueError("training_arguments is None")

        self.internal_checks_before_train()
        dist_logger("Training will start soon")
        self.trainer.train()
        dist_logger("Training end")

        self.after_train()

        if is_distributed_training():
            if distributed.get_rank() == self.config.local_rank:
                post_training(config=self.config, tokenizer=self.tokenizer)
        else:
            post_training(config=self.config, tokenizer=self.tokenizer)

        dist_logger(f"Model saved to {self.training_arguments.output_dir}")

        self.at_end()

    def push_to_hub(
        self,
        repo_id: Optional[str] = None,
        private: Optional[bool] = None,
        safe_serialization: Optional[bool] = None,
        need_push_to_hub_bos_add_bos_token: Optional[bool] = None,
    ) -> None:
        repo_id = repo_id or self.config.hub_model_id

        private = private if private is not None else self.config.hub_private_repo
        safe_serialization = safe_serialization if safe_serialization is not None else self.config.save_safetensors
        need_push_to_hub_bos_add_bos_token = (
            need_push_to_hub_bos_add_bos_token
            if need_push_to_hub_bos_add_bos_token is not None
            else self.config.push_to_hub_bos_add_bos_token
        )

        if repo_id is None:
            raise ValueError("repo_id and hub_model_id is None, but you want to push to HF hub")

        if self.tokenizer is None:
            dist_logger.warning("Tokenizer is None. Can't push to the hub")
        else:
            self.tokenizer.push_to_hub(repo_id=repo_id, private=private)
            sleep(10.0)
            if need_push_to_hub_bos_add_bos_token:
                push_to_hub_bos_add_bos_token(repo_id=repo_id)

        if self.model is None:
            raise ValueError("Model is None. Can't push to the hub")
        else:
            self.model.push_to_hub(
                repo_id=repo_id,
                private=private,
                safe_serialization=safe_serialization,
                max_shard_size=self.config.max_shard_size,
            )

    def fuse_lora(self) -> PreTrainedModel:
        if not self.config.apply_lora:
            dist_logger.warning("Apply LoRA set to False at config")

        if isinstance(self.model, PeftModel):
            self.model = self.model.merge_and_unload()
        else:
            raise TypeError(f"Can't fuse model, this the model is not the PeftModel. Model type: {type(self.model)}")

        dist_logger("LoRA fused")

        return self.model

    def after_train(self) -> None:
        return None

    def at_beggining(self) -> None:
        return None

    def at_end(self) -> None:
        return None

src/xllm/experiments/base.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .. import enums
from ..experiments.base import Experiment
from ..utils.registry import Registry

experiments_registry = Registry(name=enums.Registry.experiments)

experiments_registry.add(key=enums.Experiments.base, value=Experiment)

src/xllm/experiments/registry.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: F401

from .quantizer import Quantizer

src/xllm/quantization/__init__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import os
from typing import Any, Dict, List, Optional, Union

import torch
from optimum.gptq import GPTQQuantizer
from tqdm import tqdm
from transformers import PreTrainedModel, PreTrainedTokenizer

from .. import enums
from ..core.config import HuggingFaceConfig
from ..core.dependencies import build_dataset, build_model, build_tokenizer
from ..datasets.registry import datasets_registry
from ..utils.logger import dist_logger
from ..utils.post_training import push_to_hub_bos_add_bos_token


class Quantizer:
    # Thanks to Phil Schmid's post: https://www.philschmid.de/gptq-llama

    def __init__(
        self,
        config: HuggingFaceConfig,
        tokenizer: Optional[PreTrainedTokenizer] = None,
        model: Optional[PreTrainedModel] = None,
        dataset: Union[str, List[str], None] = None,
    ):
        self.config = config

        self.tokenizer = tokenizer
        self.model = model
        self.dataset = dataset

        self.quantizer: Optional[GPTQQuantizer] = None
        self.quantized_model: Optional[PreTrainedModel] = None

    def internal_checks(self) -> None:
        if not torch.cuda.is_available():
            dist_logger.warning("CUDA is not available")

        self.config.check_auto_gptq()

    def build(self) -> None:
        self.internal_checks()

        if self.tokenizer is None:
            self.tokenizer = build_tokenizer(config=self.config, use_fast=False)
            dist_logger.info(f"Tokenizer {self.config.correct_tokenizer_name_or_path} was built")

        if self.model is None:
            if not self.config.low_cpu_mem_usage:
                dist_logger.warning("low_cpu_mem_usage at config is False. Recommended to set to True for quantization")
            self.model = build_model(
                config=self.config,
                quantization_config=None,
            )
            dist_logger.info(f"Model {self.config.model_name_or_path} was built")

        dataset = self.build_dataset() if self.dataset is None else self.dataset

        self.quantizer = GPTQQuantizer(
            bits=self.config.gptq_bits,
            group_size=self.config.gptq_group_size,
            dataset=dataset,
            model_seqlen=self.config.max_length,
        )
        dist_logger.info("Quantizer loaded")

        return None

    def build_dataset(self) -> Union[str, List[str]]:
        dataset_id = None
        samples: List[str] = list()

        if self.config.quantization_dataset_id is not None:
            dataset_id = self.config.quantization_dataset_id
        else:
            if self.config.prepare_dataset:
                dataset_cls = datasets_registry.get(self.config.dataset_key)

                if dataset_cls is None:
                    raise ValueError(f"Dataset with key {self.config.dataset_key} not found")

                dataset_cls.prepare(config=self.config)

            raw_dataset = build_dataset(config=self.config, is_train=True)

            if raw_dataset is None:
                raise ValueError("Quantization dataset can't be loaded")

            samples = list()

            for sample_index in tqdm(
                range(len(raw_dataset)), desc="Loading quantization dataset", total=self.config.quantization_max_samples
            ):
                sample: Dict[str, Any] = raw_dataset[sample_index]
                text_parts = sample[enums.General.text_parts]
                text = "\n".join(text_parts)
                if isinstance(text, str):
                    samples.append(text)

                if 0 < self.config.quantization_max_samples == len(samples):
                    break

        return dataset_id or samples

    def quantize(self) -> None:
        if self.tokenizer is None:
            raise ValueError("tokenizer is None. It is impossible to quantize. Please run build")

        if self.model is None:
            raise ValueError("model is None. It is impossible to quantize. Please run build")

        if self.quantizer is None:
            raise ValueError("quantizer is None. It is impossible to quantize. Please run build")

        dist_logger.info("Start quantization")
        self.quantized_model = self.quantizer.quantize_model(self.model, self.tokenizer)
        dist_logger.info("Quantization complete")

        return None

    def save(self) -> None:
        if self.quantized_model is None:
            raise ValueError("quantized_model is None. Nothing to save")

        dist_logger.info(f"Saving quantized model to {self.config.quantized_model_path}")
        self.quantized_model.save_pretrained(
            save_directory=self.config.quantized_model_path,
            safe_serialization=self.config.save_safetensors,
        )

        tokenizer = build_tokenizer(config=self.config)

        tokenizer.save_pretrained(save_directory=self.config.quantized_model_path)

        path_to_config = os.path.join(self.config.quantized_model_path, "config.json")
        path_to_quantize_config = os.path.join(self.config.quantized_model_path, "quantize_config.json")

        if self.quantizer is not None:
            with open(
                path_to_quantize_config,
                "w",
                encoding="utf-8",
            ) as file_object:
                self.quantizer.disable_exllama = False
                json.dump(self.quantizer.to_dict(), file_object, indent=2)
        else:
            dist_logger.error("quantizer is None. saved quantized model can be broken")

        with open(path_to_config, "r", encoding="utf-8") as file_object:
            model_config = json.load(file_object)
            model_config["quantization_config"]["disable_exllama"] = False

        with open(path_to_config, "w", encoding="utf-8") as file_object:
            json.dump(model_config, file_object, indent=2)

        if self.config.quantized_hub_model_id is not None:
            dist_logger.info(f"Push quantized model to the hub {self.config.quantized_hub_model_id}")
            self.quantized_model.push_to_hub(
                repo_id=self.config.quantized_hub_model_id,
                private=self.config.quantized_hub_private_repo,
                safe_serialization=self.config.save_safetensors,
                max_shard_size=self.config.max_shard_size,
            )
            if self.tokenizer is not None:
                self.tokenizer.push_to_hub(
                    repo_id=self.config.quantized_hub_model_id,
                    private=self.config.quantized_hub_private_repo,
                    safe_serialization=self.config.save_safetensors,
                )
                if self.config.push_to_hub_bos_add_bos_token:
                    push_to_hub_bos_add_bos_token(repo_id=self.config.quantized_hub_model_id)
        else:
            dist_logger.warning("quantized_hub_model_id is None. Model will stay locally")

src/xllm/quantization/quantizer.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: F401

from .download import download
from .fuse import fuse
from .quantize import quantize
from .train import train

src/xllm/run/__init__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
from typing import Tuple

from loguru import logger
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    PreTrainedModel,
    PreTrainedTokenizer,
)

from ..core.config import HuggingFaceConfig
from ..datasets.registry import datasets_registry


def download(config: HuggingFaceConfig) -> Tuple[PreTrainedTokenizer, PreTrainedModel]:
    json_config = json.dumps(config.__dict__, indent=2)
    logger.info(f"Config:\n{json_config}")

    if config.prepare_dataset:
        dataset_cls = datasets_registry.get(config.dataset_key)

        if dataset_cls is None:
            raise ValueError(f"Dataset with key {config.dataset_key} not found")

        dataset_cls.prepare(config=config)
        logger.info(f"Dataset {dataset_cls.__name__} loaded")
    else:
        logger.warning("Dataset is not prepared because this set in config")

    tokenizer = AutoTokenizer.from_pretrained(config.correct_tokenizer_name_or_path)
    logger.info(f"Tokenizer {config.correct_tokenizer_name_or_path} loaded")

    model = AutoModelForCausalLM.from_pretrained(config.model_name_or_path)
    logger.info(f"Model {config.model_name_or_path} loaded")

    return tokenizer, model

src/xllm/run/download.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Tuple

from loguru import logger
from transformers import PreTrainedModel, PreTrainedTokenizer

from ..core.config import HuggingFaceConfig
from ..utils.post_training import fuse_lora


def fuse(config: HuggingFaceConfig) -> Tuple[PreTrainedTokenizer, PreTrainedModel]:
    tokenizer, model = fuse_lora(config=config)
    logger.info("Fusing complete")

    return tokenizer, model

src/xllm/run/fuse.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from ..core.config import HuggingFaceConfig
from ..quantization.quantizer import Quantizer


def quantize(config: HuggingFaceConfig) -> Quantizer:
    quantizer = Quantizer(config=config)
    quantizer.build()
    quantizer.quantize()
    quantizer.save()

    return quantizer

src/xllm/run/quantize.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from ..core.config import HuggingFaceConfig
from ..experiments.base import Experiment
from ..experiments.registry import experiments_registry


def train(
    config: HuggingFaceConfig,
) -> Experiment:
    experiment_cls = experiments_registry.get(config.experiment_key)

    if experiment_cls is None:
        raise ValueError(f"Experiment class {config.experiment_key} not found")

    experiment: Experiment = experiment_cls(config=config)

    experiment.build()

    experiment.run()

    return experiment

src/xllm/run/train.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: F401

from .lm import LMTrainer
from .registry import trainers_registry

src/xllm/trainers/__init__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Dict, Optional, Tuple, Union

from peft import PeftModel  # type: ignore
from torch import Tensor, nn
from transformers import PreTrainedModel, Trainer, TrainingArguments

from .. import enums
from ..collators.base import BaseCollator
from ..core.config import HuggingFaceConfig
from ..datasets.base import BaseDataset


class LMTrainer(Trainer):
    def __init__(
        self,
        config: HuggingFaceConfig,
        model: Union[PreTrainedModel, PeftModel],
        args: TrainingArguments,
        data_collator: BaseCollator,
        train_dataset: BaseDataset,
        ignore_index: int,
        eval_dataset: Optional[BaseDataset] = None,
    ):
        self.config = config

        super().__init__(
            model=model,
            args=args,
            data_collator=data_collator,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
        )

        self.ignore_index = ignore_index

        self.criterion = nn.CrossEntropyLoss(
            label_smoothing=self.args.label_smoothing_factor,
            ignore_index=self.ignore_index,
        )

    def compute_loss(
        self,
        model: Union[PreTrainedModel, PeftModel],
        inputs: Dict[str, Tensor],
        return_outputs: bool = False,
    ) -> Union[Tensor, Tuple[Tensor, Dict[str, Tensor]]]:
        labels = inputs.pop(enums.Transformers.labels)

        outputs = model(**inputs)

        # NOTE from BobaZooba: just copy from hf code. did not go deeper
        #
        # Save past state if it exists
        # TODO: this needs to be fixed and made cleaner later.
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

        loss = self.criterion(outputs.logits.reshape(-1, outputs.logits.size(-1)), labels.reshape(-1))

        return (loss, outputs) if return_outputs else loss

src/xllm/trainers/lm.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .. import enums
from ..trainers.lm import LMTrainer
from ..utils.registry import Registry

trainers_registry = Registry(name=enums.Registry.trainers, default_value=LMTrainer)

trainers_registry.add(key=enums.Trainers.lm, value=LMTrainer)

src/xllm/trainers/registry.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Dict, List, Union

from torch import Tensor

RawSample = Dict[str, Union[str, int, float, List[str]]]
Batch = Dict[str, Tensor]

src/xllm/types.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ruff: noqa: F401

from .cli import setup_cli
from .logger import dist_logger
from .miscellaneous import is_distributed_training
from .nn import apply_lora, stabilize_training
from .registry import Registry

src/xllm/utils/__init__.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

import transformers
from dotenv import load_dotenv
from loguru import logger

from xllm import enums

from ..core.config import HuggingFaceConfig


def setup_cli(config: HuggingFaceConfig, logger_path: str = "xllm.log", rotation: str = "5 MB") -> None:
    logger.add(logger_path, rotation=rotation)
    load_dotenv(dotenv_path=config.path_to_env_file)
    logger.info(".env loaded")

    os.environ[enums.EnvironmentVariables.tokenizers_parallelism] = "false"

    if config.report_to_wandb and enums.EnvironmentVariables.wandb_api_key not in os.environ:
        logger.warning("W&B token not found in env vars")

    if enums.EnvironmentVariables.huggingface_hub_token not in os.environ:
        logger.warning("HuggingFaceHub token not found in env vars")

    transformers.set_seed(seed=config.seed)
    transformers.logging.set_verbosity_error()
    logger.info(f'Logger path "{logger_path}" with rotation "{rotation}"')

    return None

src/xllm/utils/cli.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional

import torch.distributed as distributed
from loguru import logger

from .. import enums
from ..utils.miscellaneous import is_distributed_training


class DistributedLogger:
    @classmethod
    def can_log(cls, local_rank: int = 0) -> bool:
        if is_distributed_training():
            if distributed.get_rank() == local_rank:
                return True
            else:
                return False
        else:
            return True

    @classmethod
    def info(cls, message: str, local_rank: int = 0) -> None:
        if cls.can_log(local_rank=local_rank):
            logger.info(message)
        return None

    @classmethod
    def warning(cls, message: str, local_rank: int = 0) -> None:
        if cls.can_log(local_rank=local_rank):
            logger.warning(message)
        return None

    @classmethod
    def error(cls, message: str, local_rank: int = 0) -> None:
        if cls.can_log(local_rank=local_rank):
            logger.error(message)
        return None

    @classmethod
    def critical(cls, message: str, local_rank: int = 0) -> None:
        if cls.can_log(local_rank=local_rank):
            logger.critical(message)
        return None

    @classmethod
    def log(cls, message: str, level: Optional[str] = None, local_rank: int = 0) -> None:
        if level == enums.LogLevel.warning:
            cls.warning(message, local_rank)
        elif level == enums.LogLevel.error:
            cls.error(message, local_rank)
        elif level == enums.LogLevel.critical:
            cls.critical(message, local_rank)
        else:
            cls.info(message, local_rank)
        return None

    @classmethod
    def __call__(
        cls,
        message: str,
        level: Optional[str] = enums.LogLevel.info,
        local_rank: int = 0,
    ) -> None:
        cls.log(message, level, local_rank)
        return None


dist_logger = DistributedLogger()

src/xllm/utils/logger.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Dict, List, Set, Tuple

import torch.distributed as distributed


def is_distributed_training() -> bool:
    return distributed.is_available() and distributed.is_initialized()


def have_missing_keys(data: Dict[str, Any], must_have_keys: List[str]) -> Tuple[bool, Set[str]]:
    flag = False

    existing_keys = set(data.keys())
    difference = set(must_have_keys).difference(existing_keys)

    if len(difference) > 0:
        flag = True

    return flag, difference

src/xllm/utils/miscellaneous.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional, Tuple

import torch
from peft import LoraConfig, PeftModel, get_peft_model  # type: ignore
from peft.tuners.lora import LoraLayer
from torch import nn
from transformers import (
    PreTrainedModel,
)

from ..core.config import HuggingFaceConfig


def apply_lora(
    config: HuggingFaceConfig, model: PreTrainedModel, lora_config: Optional[LoraConfig] = None
) -> Tuple[PeftModel, LoraConfig]:
    lora_target_modules = config.lora_target_modules

    if lora_target_modules is None:
        target_modules = set()

        for name, module in model.named_modules():
            if isinstance(module, nn.Linear):
                names = name.split(".")
                target_modules.add(names[0] if len(names) == 1 else names[-1])

        if "lm_head" in target_modules:  # stabilize in 16-bit
            target_modules.remove("lm_head")

        lora_target_modules = list(target_modules)

    if lora_config is None:
        lora_config = LoraConfig(
            r=config.lora_rank,
            lora_alpha=config.lora_alpha,
            lora_dropout=config.lora_dropout,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=lora_target_modules,
        )

    model = get_peft_model(model=model, peft_config=lora_config)

    return model, lora_config


def stabilize_training(model: PreTrainedModel) -> PreTrainedModel:
    is_bf16_supported = torch.cuda.is_available() and torch.cuda.is_bf16_supported()

    for name, module in model.named_modules():
        if isinstance(module, LoraLayer) and is_bf16_supported:
            module.lora_A.to(torch.bfloat16)
            module.lora_A.to(torch.bfloat16)
            module.lora_embedding_A.to(torch.bfloat16)
            module.lora_embedding_B.to(torch.bfloat16)
        elif "norm" in name:
            module.to(torch.float32)
        elif (
            ("lm_head" in name or "embed_tokens" in name)
            and hasattr(module, "weight")
            and is_bf16_supported
            and module.weight.dtype == torch.float32
        ):
            module.to(torch.bfloat16)

    return model

src/xllm/utils/nn.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import os
from time import sleep
from typing import Tuple

from huggingface_hub import HfApi, hf_hub_download
from loguru import logger
from peft import PeftModel  # type: ignore
from transformers import AutoModelForCausalLM, PreTrainedModel, PreTrainedTokenizer

from ..core.config import HuggingFaceConfig
from ..core.dependencies import build_tokenizer

TOKENIZER_CONFIG_FILE = "tokenizer_config.json"


def fuse_lora(config: HuggingFaceConfig) -> Tuple[PreTrainedTokenizer, PreTrainedModel]:
    lora_model_name_or_path_for_fusing = config.lora_model_name_or_path_for_fusing

    tokenizer = build_tokenizer(config=config)
    logger.info(f"Tokenizer {config.correct_tokenizer_name_or_path} loaded")
    model = AutoModelForCausalLM.from_pretrained(
        pretrained_model_name_or_path=config.model_name_or_path,
        torch_dtype=config.dtype,
        trust_remote_code=config.trust_remote_code,
    )
    logger.info(f"Model {config.model_name_or_path} loaded")
    model = PeftModel.from_pretrained(model, lora_model_name_or_path_for_fusing)
    logger.info(f"LoRA {lora_model_name_or_path_for_fusing} loaded")
    logger.info("Start fusing")
    model = model.merge_and_unload()
    logger.info("LoRA fused")

    model_dtype = next(iter(model.parameters())).dtype
    if model_dtype != config.dtype:
        model = model.to(config.dtype)
        logger.info(f"Model converted to: {config.dtype}")

    if config.fused_model_local_path is not None:
        logger.info(f"Saving locally to {config.fused_model_local_path}")
        tokenizer.save_pretrained(
            config.fused_model_local_path,
            safe_serialization=config.save_safetensors,
        )
        model.save_pretrained(
            config.fused_model_local_path,
            safe_serialization=config.save_safetensors,
        )
        if config.push_to_hub_bos_add_bos_token:
            path_to_tokenizer_config = os.path.join(config.fused_model_local_path, TOKENIZER_CONFIG_FILE)
            with open(path_to_tokenizer_config) as file_object:
                tokenizer_config = json.load(file_object)

            tokenizer_config["add_bos_token"] = True
            tokenizer_config["add_eos_token"] = False

            with open(path_to_tokenizer_config, "w") as file_object:
                json.dump(tokenizer_config, file_object, indent=2)
        logger.info(f"Model saved locally to {config.fused_model_local_path}")

    if config.push_to_hub:
        logger.info(f"Pushing model to the hub {config.hub_model_id}")
        if config.hub_model_id is not None:
            tokenizer.push_to_hub(
                repo_id=config.hub_model_id,
                private=config.hub_private_repo,
                safe_serialization=config.save_safetensors,
            )
            model.push_to_hub(
                repo_id=config.hub_model_id,
                private=config.hub_private_repo,
                safe_serialization=config.save_safetensors,
                max_shard_size=config.max_shard_size,
            )
            if config.push_to_hub_bos_add_bos_token:
                push_to_hub_bos_add_bos_token(repo_id=config.hub_model_id)
        else:
            raise ValueError("Fused model push to hub failed, because config.hub_model_id if None")

    return tokenizer, model


def push_to_hub_bos_add_bos_token(repo_id: str) -> None:
    local_path = hf_hub_download(repo_id=repo_id, filename=TOKENIZER_CONFIG_FILE)

    with open(local_path) as file_object:
        tokenizer_config = json.load(file_object)

    tokenizer_config["add_bos_token"] = True
    tokenizer_config["add_eos_token"] = False

    local_tokenizer_config_file = os.path.join("./", TOKENIZER_CONFIG_FILE)

    with open(local_tokenizer_config_file, "w") as file_object:
        json.dump(tokenizer_config, file_object, indent=2)

    api = HfApi()
    api.upload_file(
        path_or_fileobj=local_tokenizer_config_file,
        path_in_repo=TOKENIZER_CONFIG_FILE,
        repo_id=repo_id,
        repo_type="model",
    )

    return None


def post_training(config: HuggingFaceConfig, tokenizer: PreTrainedTokenizer) -> None:
    if config.push_to_hub:
        if config.hub_model_id is None:
            raise ValueError("hub_model_id is None, but you want to push to HF hub")
        tokenizer.push_to_hub(repo_id=config.hub_model_id, private=config.hub_private_repo)
        sleep(10.0)
        if config.push_to_hub_bos_add_bos_token:
            push_to_hub_bos_add_bos_token(repo_id=config.hub_model_id)

    return None

src/xllm/utils/post_training.py
KompleteAI/xllm
# Copyright 2023 Komplete AI Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Dict, Optional

from loguru import logger


class Registry:
    DEFAULT_KEY = "default"

    def __init__(self, name: str, default_value: Optional[Any] = None):
        self.name = name

        self.mapper: Dict[str, Any] = dict()

        if default_value is not None:
            self.mapper[self.DEFAULT_KEY] = default_value

    def add(self, key: str, value: Any, override: bool = False) -> None:
        if not override and key in self.mapper:
            raise ValueError(f"Key exist in {self.name} registry")
        self.mapper[key] = value
        return None

    def get(self, key: Optional[str]) -> Any:
        if key is not None:
            value = self.mapper.get(key, None)
        else:
            value = self.mapper.get(self.DEFAULT_KEY, None)

        if value is None:
            value = self.mapper.get(self.DEFAULT_KEY, None)
            if value is not None:
                logger.warning(f"Default item {value.__name__} chosen from {self.name} registry")

        if value is None:
            raise ValueError(f"Item with key {key} not found in {self.name} registry")

        return value

src/xllm/utils/registry.py
KompleteAI/xllm
import json
import os

import pytest
from _pytest.tmpdir import TempPathFactory
from peft import LoraConfig, PeftModel, get_peft_model
from transformers import (
    AutoTokenizer,
    FalconConfig,
    FalconForCausalLM,
    LlamaConfig,
    LlamaForCausalLM,
    PreTrainedTokenizer,
    TrainingArguments,
)

from src.xllm import enums
from src.xllm.collators.lm import LMCollator
from src.xllm.core.config import HuggingFaceConfig
from src.xllm.datasets.soda import SodaDataset
from tests.helpers.constants import (
    FALCON_TOKENIZER_DIR,
    LLAMA_TOKENIZER_DIR,
    LORA_FOR_LLAMA_DEFAULT_TARGET_MODULES,
)
from tests.helpers.dummy_data import DATA, SODA_DATASET


@pytest.fixture(scope="session")
def llama_tokenizer() -> PreTrainedTokenizer:
    tokenizer = AutoTokenizer.from_pretrained(LLAMA_TOKENIZER_DIR)
    tokenizer.pad_token = tokenizer.eos_token
    return tokenizer


@pytest.fixture(scope="session")
def llama_model_config(llama_tokenizer: PreTrainedTokenizer) -> LlamaConfig:
    config = LlamaConfig(
        vocab_size=len(llama_tokenizer),
        hidden_size=8,
        intermediate_size=32,
        num_hidden_layers=2,
        num_attention_heads=2,
        max_position_embeddings=32,
    )
    return config


@pytest.fixture(scope="session")
def llama_model(llama_model_config: LlamaConfig) -> LlamaForCausalLM:
    model = LlamaForCausalLM(config=llama_model_config)
    return model


@pytest.fixture(scope="session")
def llama_lora_config() -> LoraConfig:
    lora_config = LoraConfig(
        r=2,
        target_modules=LORA_FOR_LLAMA_DEFAULT_TARGET_MODULES,
        bias="none",
        task_type="CAUSAL_LM",
        lora_alpha=8,
        lora_dropout=0.1,
    )
    return lora_config


@pytest.fixture(scope="session")
def llama_lora_model(llama_model: LlamaForCausalLM, llama_lora_config: LoraConfig) -> PeftModel:
    llama_model = get_peft_model(model=llama_model, peft_config=llama_lora_config)
    return llama_model


@pytest.fixture(scope="session")
def falcon_tokenizer() -> PreTrainedTokenizer:
    tokenizer = AutoTokenizer.from_pretrained(FALCON_TOKENIZER_DIR)
    tokenizer.pad_token = tokenizer.eos_token
    return tokenizer


@pytest.fixture(scope="session")
def falcon_model_config(falcon_tokenizer: PreTrainedTokenizer) -> FalconConfig:
    config = FalconConfig(
        vocab_size=len(falcon_tokenizer),
        hidden_size=8,
        intermediate_size=32,
        num_hidden_layers=2,
        num_attention_heads=2,
        max_position_embeddings=32,
    )
    return config


@pytest.fixture(scope="session")
def falcon_model(falcon_model_config: FalconConfig) -> FalconForCausalLM:
    model = FalconForCausalLM(config=falcon_model_config)
    return model


@pytest.fixture(scope="session")
def soda_dataset() -> SodaDataset:
    dataset = SodaDataset(data=SODA_DATASET)
    return dataset


@pytest.fixture(scope="session")
def llama_lm_collator(llama_tokenizer: PreTrainedTokenizer) -> LMCollator:
    collator = LMCollator(tokenizer=llama_tokenizer, max_length=32)
    return collator


@pytest.fixture(scope="session")
def training_arguments() -> TrainingArguments:
    arguments = TrainingArguments(
        output_dir="./outputs/",
        per_device_train_batch_size=2,
        gradient_accumulation_steps=2,
        warmup_steps=50,
        learning_rate=2e-4,
        max_steps=500,
        num_train_epochs=1,
        weight_decay=0.001,
        max_grad_norm=1.0,
        label_smoothing_factor=0.1,
        logging_steps=10,
        save_strategy="steps",
        save_steps=100,
        save_total_limit=1,
        hub_strategy="checkpoint",
        push_to_hub=False,
        save_safetensors=True,
        remove_unused_columns=False,
        log_level=enums.LogLevel.info,
        disable_tqdm=False,
    )
    return arguments


@pytest.fixture(scope="session")
def config() -> HuggingFaceConfig:
    hf_config = HuggingFaceConfig(deepspeed_stage=0)
    return hf_config


@pytest.fixture(scope="session")
def path_to_train_dummy_data(tmp_path_factory: TempPathFactory) -> str:
    path = tmp_path_factory.mktemp("tmp") / "train.jsonl"
    with open(path, "w") as file_object:
        for sample in DATA:
            file_object.write(json.dumps(sample) + "\n")
    return os.path.abspath(path)


@pytest.fixture(scope="session")
def path_to_train_prepared_dummy_data(tmp_path_factory: TempPathFactory) -> str:
    path = tmp_path_factory.mktemp("tmp") / "prepared_train.jsonl"
    with open(path, "w") as file_object:
        for sample in SODA_DATASET:
            file_object.write(json.dumps(sample) + "\n")
    return os.path.abspath(path)


@pytest.fixture(scope="session")
def path_to_empty_train_dummy_data(tmp_path_factory: TempPathFactory) -> str:
    path = tmp_path_factory.mktemp("tmp") / "empty_train.jsonl"
    return os.path.abspath(path)


@pytest.fixture(scope="session")
def path_to_eval_dummy_data(tmp_path_factory: TempPathFactory) -> str:
    path = tmp_path_factory.mktemp("tmp") / "eval.jsonl"
    with open(path, "w") as file_object:
        for sample in DATA:
            file_object.write(json.dumps(sample) + "\n")
    return os.path.abspath(path)


@pytest.fixture(scope="session")
def path_to_empty_eval_dummy_data(tmp_path_factory: TempPathFactory) -> str:
    path = tmp_path_factory.mktemp("tmp") / "empty_eval.jsonl"
    return os.path.abspath(path)


@pytest.fixture(scope="session")
def path_to_fused_model_local_path(tmp_path_factory: TempPathFactory) -> str:
    path = tmp_path_factory.mktemp("tmp") / "fused_model/"
    return os.path.abspath(path)


@pytest.fixture(scope="session")
def path_to_download_result(tmp_path_factory: TempPathFactory) -> str:
    path = tmp_path_factory.mktemp("tmp") / "data.jsonl"
    return os.path.abspath(path)

tests/conftest.py
KompleteAI/xllm
import os

from dotenv import load_dotenv

load_dotenv()


PROJECT_DIR: str = os.getcwd()

STATIC_DIR: str = os.path.join(PROJECT_DIR, "tests/static/")
TOKENIZERS_DIR: str = os.path.join(STATIC_DIR, "tokenizers/")

LLAMA_TOKENIZER_DIR: str = os.path.join(TOKENIZERS_DIR, "llama/")
FALCON_TOKENIZER_DIR: str = os.path.join(TOKENIZERS_DIR, "falcon/")

LORA_FOR_LLAMA_DEFAULT_TARGET_MODULES = [
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "gate_proj",
    "up_proj",
    "down_proj",
]

tests/helpers/constants.py
KompleteAI/xllm
from typing import List, Optional, Tuple

from src.xllm import enums
from src.xllm.core.config import HuggingFaceConfig
from src.xllm.datasets.base import BaseDataset
from src.xllm.datasets.soda import SodaDataset
from src.xllm.types import RawSample

DATA = [
    {
        enums.General.text_parts: [
            "Person 1: Hello",
            "Person 2: It's me",
            "Person 1: I was wondering",
        ]
    },
    {
        enums.General.text_parts: [
            "You are a sith lord",
            "Kenobi: Hello there",
            "General Grievous: General Kenobi",
        ]
    },
]

SODA_DATASET = [
    {
        SodaDataset.HEADER_KEY: "This is dialog",
        SodaDataset.DIALOG_KEY: ["Hello", "It's me", "I was wondering"],
    },
    {
        SodaDataset.HEADER_KEY: "This is another dialog",
        SodaDataset.DIALOG_KEY: ["Sup", "Hello", "It's me", "I was wondering", "Please buy some coins"],
    },
]

RAW_SODA_DATASET = {
    "train": [
        {
            "narrative": "You are a sith lord",
            "dialogue": ["Hello there", "General Kenobi"],
            "speakers": ["Kenobi", "General Grievous"],
        }
    ],
    "validation": [
        {
            "narrative": "You are a sith lord",
            "dialogue": ["Hello there", "General Kenobi"],
            "speakers": ["Kenobi", "General Grievous"],
        }
    ],
    "test": [
        {
            "narrative": "You are a sith lord",
            "dialogue": ["Hello there", "General Kenobi"],
            "speakers": ["Kenobi", "General Grievous"],
        }
    ],
}


class DummyDataset(BaseDataset):
    @classmethod
    def download(cls, config: HuggingFaceConfig) -> Tuple[List[RawSample], Optional[List[RawSample]]]:
        return DATA, None

    def get_sample(self, index: int) -> RawSample:
        return self.data[index]

tests/helpers/dummy_data.py
KompleteAI/xllm
from contextlib import contextmanager
from typing import Any, Dict, Union

import datasets
import torch
from peft import PeftModel
from pytest import MonkeyPatch
from torch import dtype
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    GPTQConfig,
    LlamaConfig,
    LlamaForCausalLM,
)

from src.xllm.trainers.lm import LMTrainer
from tests.helpers.dummy_data import RAW_SODA_DATASET


@contextmanager
def patch_from_pretrained_auto_causal_lm(monkeypatch: MonkeyPatch) -> Any:
    def from_pretrained(
        pretrained_model_name_or_path: str,
        quantization_config: Union[BitsAndBytesConfig, GPTQConfig, None] = None,
        torch_dtype: dtype = torch.float16,
        trust_remote_code: bool = True,
        device_map: Union[str, Dict[str, Any], None] = None,
        use_cache: bool = False,
        use_flash_attention_2: bool = True,
    ) -> LlamaForCausalLM:
        config = LlamaConfig(
            vocab_size=32_000,
            hidden_size=8,
            intermediate_size=32,
            num_hidden_layers=2,
            num_attention_heads=2,
            max_position_embeddings=32,
        )
        model = LlamaForCausalLM(config=config)
        return model

    monkeypatch.setattr(AutoModelForCausalLM, "from_pretrained", from_pretrained)
    yield True
    monkeypatch.undo()


@contextmanager
def patch_load_soda_dataset(monkeypatch: MonkeyPatch) -> Any:
    def load_dataset(dataset_name: str):
        return RAW_SODA_DATASET

    monkeypatch.setattr(datasets, "load_dataset", load_dataset)
    yield True
    monkeypatch.undo()


@contextmanager
def patch_trainer_train(monkeypatch: MonkeyPatch) -> Any:
    def train(*args, **kwargs):
        return None

    monkeypatch.setattr(LMTrainer, "train", train)
    yield True
    monkeypatch.undo()


@contextmanager
def patch_tokenizer_from_pretrained(monkeypatch: MonkeyPatch) -> Any:
    def mock_tokenizer_from_pretrained(*args, **kwargs):
        class MockTokenizer:
            def __init__(self):
                self.pad_token = None
                self.eos_token = "EOS"
                self.padding_side = "right"

            def save_pretrained(self, *args, **kwargs):
                return None

            def push_to_hub(self, *args, **kwargs):
                ...

        mock_tokenizer = MockTokenizer()

        return mock_tokenizer

    monkeypatch.setattr(AutoTokenizer, "from_pretrained", mock_tokenizer_from_pretrained)

    yield True
    monkeypatch.undo()


@contextmanager
def patch_peft_model_from_pretrained(monkeypatch: MonkeyPatch) -> Any:
    def mock_peft_model_from_pretrained(*args, **kwargs):
        class MockPeftModel:
            def parameters(self):
                yield torch.rand(1024, 1024)

            def merge_and_unload(self, *args, **kwargs):
                return self

            def save_pretrained(self, *args, **kwargs):
                return None

            def push_to_hub(self, *args, **kwargs):
                ...

        mock_peft_model = MockPeftModel()

        return mock_peft_model

    monkeypatch.setattr(PeftModel, "from_pretrained", mock_peft_model_from_pretrained)

    yield True
    monkeypatch.undo()

tests/helpers/patches.py
KompleteAI/xllm
from typing import Optional

import pytest
from torch import Tensor
from transformers import PreTrainedTokenizer

from src.xllm import enums
from src.xllm.collators.completion import CompletionCollator
from tests.helpers.dummy_data import DATA


@pytest.mark.parametrize("prefix_end", [None, ":"])
def test_completion_collator(llama_tokenizer: PreTrainedTokenizer, prefix_end: Optional[str]):
    collator = CompletionCollator(tokenizer=llama_tokenizer, max_length=128, prefix_end=prefix_end)
    batch = collator(DATA)
    for _key, value in batch.items():
        assert isinstance(value, Tensor)

    condition_result = (batch[enums.Transformers.labels][:, :2] == llama_tokenizer.pad_token_id).unique()

    assert len(condition_result) == 1 and condition_result.item()

tests/unit/collators/test_completion.py
KompleteAI/xllm
from torch import Tensor
from transformers import PreTrainedTokenizer

from src.xllm.collators.lm import LMCollator
from tests.helpers.dummy_data import DATA


def test_lm_collator(llama_tokenizer: PreTrainedTokenizer):
    collator = LMCollator(tokenizer=llama_tokenizer, max_length=128, separator="\n")
    batch = collator(DATA)
    for _key, value in batch.items():
        assert isinstance(value, Tensor)

tests/unit/collators/test_lm.py
KompleteAI/xllm
from transformers import PreTrainedTokenizer

from src.xllm import enums
from src.xllm.collators.lm import LMCollator
from src.xllm.collators.registry import collators_registry


def test_get_lm_collator(llama_tokenizer: PreTrainedTokenizer):
    collator_cls = collators_registry.get(key=enums.Collators.lm)
    collator = collator_cls(llama_tokenizer, max_length=64)
    assert isinstance(collator, LMCollator)

tests/unit/collators/test_registry.py
KompleteAI/xllm
import pytest
import torch

from src.xllm.core.config import HuggingFaceConfig
from src.xllm.core.deepspeed_configs import STAGE_2


def test_check_hub():
    config = HuggingFaceConfig(hub_model_id="llama2", push_to_hub=False)
    config.check_hub()


def test_check_none_deepspeed():
    config = HuggingFaceConfig(deepspeed_stage=None)
    config.check_deepspeed()


def test_check_not_use_flash_attention_2():
    config = HuggingFaceConfig(use_flash_attention_2=False)
    config.check_flash_attention()


def test_check_hub_push_without_repo():
    config = HuggingFaceConfig(hub_model_id=None, push_to_hub=True)
    with pytest.raises(ValueError):
        config.check_hub()


def test_tokenizer_name():
    config = HuggingFaceConfig(tokenizer_name_or_path="llama1000", model_name_or_path="llama2")
    assert config.correct_tokenizer_name_or_path == "llama1000"


def test_tokenizer_name_from_model():
    config = HuggingFaceConfig(tokenizer_name_or_path=None, model_name_or_path="llama2")
    assert config.correct_tokenizer_name_or_path == "llama2"


def test_lora_target_modules():
    config = HuggingFaceConfig(raw_lora_target_modules="q,w,e")
    assert config.lora_target_modules == ["q", "w", "e"]


@pytest.mark.skipif(not torch.cuda.is_available(), reason="Running tests without GPU")
def test_dtype():
    config = HuggingFaceConfig()
    assert config.dtype == torch.float16


def test_deepspeed():
    config = HuggingFaceConfig(deepspeed_stage=2)
    assert config.deepspeed == STAGE_2


def test_check_deepspeed():
    config = HuggingFaceConfig(deepspeed_stage=0)
    config.check_deepspeed()


def test_check_deepspeed_exception():
    config = HuggingFaceConfig(deepspeed_stage=2)
    with pytest.raises(ImportError):
        config.check_deepspeed()


def test_check_flash_attention():
    config = HuggingFaceConfig(use_flash_attention_2=False)
    config.check_flash_attention()


def test_check_flash_attention_exception():
    config = HuggingFaceConfig(use_flash_attention_2=True)
    with pytest.raises(ImportError):
        config.check_flash_attention()


def test_check_auto_gptq_exception():
    config = HuggingFaceConfig()
    with pytest.raises(ImportError):
        config.check_auto_gptq()

tests/unit/core/test_config.py
KompleteAI/xllm
import pytest
from peft import PeftModel
from pytest import MonkeyPatch
from torch import Tensor
from transformers import (
    BitsAndBytesConfig,
    GPTQConfig,
    PreTrainedTokenizer,
    TrainingArguments,
)

from src.xllm.collators.lm import LMCollator
from src.xllm.collators.registry import collators_registry
from src.xllm.core.config import HuggingFaceConfig
from src.xllm.core.dependencies import (
    build_collator,
    build_dataset,
    build_model,
    build_quantization_config,
    build_tokenizer,
    build_trainer,
    build_training_arguments,
)
from src.xllm.datasets.registry import datasets_registry
from src.xllm.datasets.soda import SodaDataset
from src.xllm.trainers.registry import trainers_registry
from tests.helpers.constants import LLAMA_TOKENIZER_DIR
from tests.helpers.dummy_data import DATA, DummyDataset
from tests.helpers.patches import patch_from_pretrained_auto_causal_lm


def test_build_training_arguments(config: HuggingFaceConfig):
    arguments = build_training_arguments(config=config)
    assert arguments.per_device_train_batch_size == config.per_device_train_batch_size
    assert arguments.deepspeed is None


def test_build_dataset_train(path_to_train_dummy_data: str):
    datasets_registry.add(key="dummy", value=DummyDataset)
    config = HuggingFaceConfig(dataset_key="dummy", train_local_path_to_data=path_to_train_dummy_data)
    dataset = build_dataset(config=config, is_train=True)
    assert dataset[0] is not None


def test_build_dataset_eval(path_to_train_dummy_data: str):
    datasets_registry.add(key="dummy1", value=DummyDataset)
    config = HuggingFaceConfig(dataset_key="dummy1", eval_local_path_to_data=path_to_train_dummy_data)
    dataset = build_dataset(config=config, is_train=False)
    assert dataset[0] is not None


def test_build_dataset_eval_none(path_to_train_dummy_data: str):
    datasets_registry.add(key="dummy2", value=DummyDataset)
    config = HuggingFaceConfig(
        dataset_key="dummy2",
        train_local_path_to_data=path_to_train_dummy_data,
        eval_local_path_to_data=None,
    )
    dataset = build_dataset(config=config, is_train=False)
    assert dataset is None


def test_build_dataset_exception(path_to_train_dummy_data: str):
    datasets_registry.add(key="exc", value=HuggingFaceConfig)
    config = HuggingFaceConfig(dataset_key="exc", train_local_path_to_data=path_to_train_dummy_data)
    with pytest.raises(ValueError):
        build_dataset(config=config, is_train=True)


def test_build_tokenizer():
    config = HuggingFaceConfig(tokenizer_name_or_path=LLAMA_TOKENIZER_DIR)
    tokenizer = build_tokenizer(config=config)
    tokenizer("hello")


def test_build_tokenizer_use_fast():
    config = HuggingFaceConfig(tokenizer_name_or_path=LLAMA_TOKENIZER_DIR)
    tokenizer = build_tokenizer(config=config, use_fast=False)
    tokenizer("hello")


def test_build_collator(config: HuggingFaceConfig, llama_tokenizer: PreTrainedTokenizer):
    collator = build_collator(config=config, tokenizer=llama_tokenizer)
    batch = collator(DATA)
    for value in batch.values():
        assert isinstance(value, Tensor)


def test_build_collator_exception(llama_tokenizer: PreTrainedTokenizer):
    collators_registry.add(key="exc", value=HuggingFaceConfig)
    config = HuggingFaceConfig(collator_key="exc")
    with pytest.raises(ValueError):
        _ = build_collator(config=config, tokenizer=llama_tokenizer)


def test_build_quantization_config_bnb():
    config = HuggingFaceConfig(load_in_8bit=True)
    quantization_config = build_quantization_config(config=config)
    assert isinstance(quantization_config, BitsAndBytesConfig)
    assert quantization_config.load_in_8bit


def test_build_quantization_config_gptq():
    config = HuggingFaceConfig(gptq_bits=4, gptq_group_size=128, from_gptq=True)
    quantization_config = build_quantization_config(config=config)
    assert isinstance(quantization_config, GPTQConfig)
    assert quantization_config.bits == 4
    assert quantization_config.group_size == 128


def test_build_quantization_config_none():
    config = HuggingFaceConfig(from_gptq=False, load_in_4bit=False, load_in_8bit=False)
    quantization_config = build_quantization_config(config=config)
    assert quantization_config is None


@pytest.mark.parametrize("apply_lora", [False, True])
def test_build_model(monkeypatch: MonkeyPatch, apply_lora: bool):
    config = HuggingFaceConfig(apply_lora=apply_lora)
    with patch_from_pretrained_auto_causal_lm(monkeypatch=monkeypatch):
        _ = build_model(
            config=config,
            quantization_config=None,
        )


def test_build_trainer(
    config: HuggingFaceConfig,
    training_arguments: TrainingArguments,
    llama_lora_model: PeftModel,
    soda_dataset: SodaDataset,
    llama_lm_collator: LMCollator,
):
    trainer = build_trainer(
        config=config,
        pad_token_id=2,
        training_arguments=training_arguments,
        model=llama_lora_model,
        train_dataset=soda_dataset,
        collator=llama_lm_collator,
        eval_dataset=None,
    )
    dataloader = trainer.get_train_dataloader()
    batch = next(iter(dataloader))
    loss = trainer.compute_loss(model=llama_lora_model, inputs=batch, return_outputs=False)
    assert isinstance(loss, Tensor)
    assert len(loss.size()) == 0


def test_build_trainer_exception(
    training_arguments: TrainingArguments,
    llama_lora_model: PeftModel,
    soda_dataset: SodaDataset,
    llama_lm_collator: LMCollator,
):
    trainers_registry.add(key="exc", value=HuggingFaceConfig)
    config = HuggingFaceConfig(trainer_key="exc")
    with pytest.raises(ValueError):
        _ = build_trainer(
            config=config,
            pad_token_id=2,
            training_arguments=training_arguments,
            model=llama_lora_model,
            train_dataset=soda_dataset,
            collator=llama_lm_collator,
            eval_dataset=None,
        )

tests/unit/core/test_dependencies.py
KompleteAI/xllm
from copy import deepcopy
from typing import List, Optional, Tuple

import pytest

from src.xllm import enums
from src.xllm.core.config import HuggingFaceConfig
from src.xllm.datasets.base import BaseDataset
from src.xllm.types import RawSample
from tests.helpers.dummy_data import DATA


class TestDataset(BaseDataset):
    @classmethod
    def download(cls, config: HuggingFaceConfig) -> Tuple[List[RawSample], Optional[List[RawSample]]]:
        return DATA, DATA

    def get_sample(self, index: int) -> RawSample:
        sample = deepcopy(self.data[index])
        sample["something"] = 1
        del sample[enums.General.text_parts]
        return sample


def test_base_dataset_exception(config: HuggingFaceConfig):
    train_data, _ = TestDataset.download(config=config)
    dataset = TestDataset(data=train_data)
    with pytest.raises(ValueError):
        _ = dataset[0]


def test_prepare(path_to_empty_train_dummy_data: str, path_to_empty_eval_dummy_data: str):
    config = HuggingFaceConfig(
        train_local_path_to_data=path_to_empty_train_dummy_data,
        eval_local_path_to_data=path_to_empty_eval_dummy_data,
    )
    TestDataset.prepare(config=config)


def test_prepare_eval_path_is_none(path_to_empty_train_dummy_data: str):
    config = HuggingFaceConfig(
        train_local_path_to_data=path_to_empty_train_dummy_data,
        add_eval_to_train_if_no_path=True,
    )
    TestDataset.prepare(config=config)


def test_prepare_eval_max_samples(path_to_empty_train_dummy_data: str, path_to_empty_eval_dummy_data: str):
    config = HuggingFaceConfig(
        train_local_path_to_data=path_to_empty_train_dummy_data,
        eval_local_path_to_data=path_to_empty_eval_dummy_data,
        add_eval_to_train_if_no_path=True,
        max_eval_samples=1,
    )
    TestDataset.prepare(config=config)

tests/unit/datasets/test_base.py
KompleteAI/xllm
from src.xllm import enums
from src.xllm.datasets.registry import datasets_registry
from src.xllm.datasets.soda import SodaDataset
from tests.helpers.dummy_data import DATA


def test_get_soda_dataset() -> None:
    dataset_cls = datasets_registry.get(key=enums.Datasets.soda)
    dataset = dataset_cls(data=DATA)
    assert isinstance(dataset, SodaDataset)

tests/unit/datasets/test_registry.py
KompleteAI/xllm
from pytest import MonkeyPatch

from src.xllm.core.config import HuggingFaceConfig
from src.xllm.datasets.soda import SodaDataset
from tests.helpers.patches import patch_load_soda_dataset


def test_soda_download(monkeypatch: MonkeyPatch, config: HuggingFaceConfig):
    with patch_load_soda_dataset(monkeypatch=monkeypatch):
        train_data, _ = SodaDataset.download(config=config)

tests/unit/datasets/test_soda.py
KompleteAI/xllm
import os

from pytest import MonkeyPatch

from src.xllm.core.config import HuggingFaceConfig
from src.xllm.experiments.base import Experiment
from tests.helpers.constants import LLAMA_TOKENIZER_DIR
from tests.helpers.patches import patch_from_pretrained_auto_causal_lm, patch_trainer_train


def test_base_experiment_init(monkeypatch: MonkeyPatch, path_to_train_dummy_data: str):
    config = HuggingFaceConfig(
        push_to_hub=False,
        deepspeed_stage=0,
        train_local_path_to_data=path_to_train_dummy_data,
        tokenizer_name_or_path=LLAMA_TOKENIZER_DIR,
    )
    with patch_from_pretrained_auto_causal_lm(monkeypatch=monkeypatch):
        Experiment(config=config)


def test_base_experiment_train(monkeypatch: MonkeyPatch, path_to_train_prepared_dummy_data: str):
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    config = HuggingFaceConfig(
        push_to_hub=False,
        deepspeed_stage=0,
        train_local_path_to_data=path_to_train_prepared_dummy_data,
        report_to_wandb=False,
        save_total_limit=0,
        max_steps=2,
        tokenizer_name_or_path=LLAMA_TOKENIZER_DIR,
    )

    with patch_from_pretrained_auto_causal_lm(monkeypatch=monkeypatch):
        experiment = Experiment(config=config)
        with patch_trainer_train(monkeypatch=monkeypatch):
            experiment.build()
            experiment.run()

tests/unit/experiments/test_base.py
KompleteAI/xllm
from pytest import MonkeyPatch

from src.xllm.core.config import HuggingFaceConfig
from src.xllm.run.download import download
from tests.helpers.constants import LLAMA_TOKENIZER_DIR
from tests.helpers.patches import patch_from_pretrained_auto_causal_lm, patch_load_soda_dataset


def test_download(monkeypatch: MonkeyPatch, path_to_download_result: str):
    config = HuggingFaceConfig(
        tokenizer_name_or_path=LLAMA_TOKENIZER_DIR, train_local_path_to_data=path_to_download_result
    )
    with patch_load_soda_dataset(monkeypatch=monkeypatch):
        with patch_from_pretrained_auto_causal_lm(monkeypatch=monkeypatch):
            download(config=config)

tests/unit/run/test_download.py
KompleteAI/xllm
from pytest import MonkeyPatch

from src.xllm.core.config import HuggingFaceConfig
from src.xllm.run.fuse import fuse
from tests.helpers.patches import (
    patch_from_pretrained_auto_causal_lm,
    patch_peft_model_from_pretrained,
    patch_tokenizer_from_pretrained,
)


def test_fuse(monkeypatch: MonkeyPatch, path_to_fused_model_local_path: str):
    config = HuggingFaceConfig(
        push_to_hub=True,
        hub_model_id="KompleteAI/SomeModelLoRA",
        lora_hub_model_id="KompleteAI/SomeModelLoRA",
        fused_model_local_path=path_to_fused_model_local_path,
        push_to_hub_bos_add_bos_token=False,
    )

    with patch_tokenizer_from_pretrained(monkeypatch=monkeypatch):
        with patch_from_pretrained_auto_causal_lm(monkeypatch=monkeypatch):
            with patch_peft_model_from_pretrained(monkeypatch=monkeypatch):
                fuse(config=config)

tests/unit/run/test_fuse.py
KompleteAI/xllm
from pytest import MonkeyPatch

from src.xllm.core.config import HuggingFaceConfig
from src.xllm.run.train import train
from tests.helpers.constants import LLAMA_TOKENIZER_DIR
from tests.helpers.patches import patch_from_pretrained_auto_causal_lm, patch_trainer_train


def test_train(monkeypatch: MonkeyPatch, path_to_train_prepared_dummy_data: str):
    config = HuggingFaceConfig(
        push_to_hub=False,
        deepspeed_stage=0,
        train_local_path_to_data=path_to_train_prepared_dummy_data,
        report_to_wandb=False,
        save_total_limit=0,
        max_steps=2,
        tokenizer_name_or_path=LLAMA_TOKENIZER_DIR,
    )
    with patch_from_pretrained_auto_causal_lm(monkeypatch=monkeypatch):
        with patch_trainer_train(monkeypatch=monkeypatch):
            train(config=config)

tests/unit/run/test_train.py
KompleteAI/xllm
from peft import PeftModel
from torch import Tensor
from transformers import TrainingArguments

from src.xllm.collators.lm import LMCollator
from src.xllm.core.config import HuggingFaceConfig
from src.xllm.datasets.soda import SodaDataset
from src.xllm.trainers.lm import LMTrainer


def test_lm_trainer(
    config: HuggingFaceConfig,
    llama_lora_model: PeftModel,
    training_arguments: TrainingArguments,
    llama_lm_collator: LMCollator,
    soda_dataset: SodaDataset,
):
    trainer = LMTrainer(
        config=config,
        model=llama_lora_model,
        args=training_arguments,
        data_collator=llama_lm_collator,
        train_dataset=soda_dataset,
        ignore_index=0,
        eval_dataset=None,
    )
    dataloader = trainer.get_train_dataloader()
    batch = next(iter(dataloader))
    loss = trainer.compute_loss(model=llama_lora_model, inputs=batch, return_outputs=False)
    assert isinstance(loss, Tensor)
    assert len(loss.size()) == 0

tests/unit/trainers/test_lm.py
KompleteAI/xllm
from peft import PeftModel
from transformers import TrainingArguments

from src.xllm import enums
from src.xllm.collators.lm import LMCollator
from src.xllm.core.config import HuggingFaceConfig
from src.xllm.datasets.soda import SodaDataset
from src.xllm.trainers.lm import LMTrainer
from src.xllm.trainers.registry import trainers_registry


def test_get_lm_trainer(
    config: HuggingFaceConfig,
    llama_lora_model: PeftModel,
    training_arguments: TrainingArguments,
    llama_lm_collator: LMCollator,
    soda_dataset: SodaDataset,
):
    trainer_cls = trainers_registry.get(key=enums.Trainers.lm)
    trainer = trainer_cls(
        config=config,
        model=llama_lora_model,
        args=training_arguments,
        data_collator=llama_lm_collator,
        train_dataset=soda_dataset,
        ignore_index=0,
        eval_dataset=None,
    )
    assert isinstance(trainer, LMTrainer)

tests/unit/trainers/test_registry.py
KompleteAI/xllm
from src.xllm.core.config import HuggingFaceConfig
from src.xllm.utils.cli import setup_cli


def test_setup_cli(config: HuggingFaceConfig):
    setup_cli(config=config)

tests/unit/utils/test_cli.py
KompleteAI/xllm
import pytest

from src.xllm import enums
from src.xllm.utils.logger import dist_logger


@pytest.mark.parametrize("message", ["Hello world", "Hey!", "All good"])
def test_info(message: str):
    dist_logger.info(message=message)


@pytest.mark.parametrize("message", ["Hello world", "Hey!", "All good"])
def test_waring(message: str):
    dist_logger.warning(message=message)


@pytest.mark.parametrize("message", ["Hello world", "Hey!", "All good"])
def test_error(message: str):
    dist_logger.error(message=message)


@pytest.mark.parametrize("message", ["Hello world", "Hey!", "All good"])
def test_critical(message: str):
    dist_logger.critical(message=message)


@pytest.mark.parametrize("message", ["Hello world", "Hey!", "All good"])
@pytest.mark.parametrize(
    "level",
    [
        enums.LogLevel.info,
        enums.LogLevel.warning,
        enums.LogLevel.error,
        enums.LogLevel.critical,
        None,
        "some",
    ],
)
def test_log(message: str, level: str):
    dist_logger.log(message=message, level=level)


@pytest.mark.parametrize("message", ["Hello world", "Hey!", "All good"])
@pytest.mark.parametrize(
    "level",
    [
        enums.LogLevel.info,
        enums.LogLevel.warning,
        enums.LogLevel.error,
        enums.LogLevel.critical,
        None,
        "some",
    ],
)
def test_call(message: str, level: str):
    dist_logger(message=message, level=level)

tests/unit/utils/test_logger.py
KompleteAI/xllm
from typing import Any, Dict, List

import pytest

from src.xllm.utils.miscellaneous import have_missing_keys


@pytest.mark.parametrize(
    argnames=["data", "must_have_keys"],
    argvalues=[
        ({"a": 1}, ["a"]),
        ({"a": 1, "b": 2}, ["a", "b"]),
        ({"a": 1, "b": 2, "c": 3}, ["a", "b", "c"]),
    ],
)
def test_no_have_missing_keys(data: Dict[str, Any], must_have_keys: List[str]) -> None:
    flag, difference = have_missing_keys(data=data, must_have_keys=must_have_keys)
    assert not flag
    assert len(difference) == 0


@pytest.mark.parametrize("data", [{"a": 1}, {"a": 1, "b": 2}, {"a": 1, "b": 2, "c": 3}])
@pytest.mark.parametrize("must_have_keys", [["c", "d", "e"], ["1", "2", "3"]])
def test_have_missing_keys(data: Dict[str, Any], must_have_keys: List[str]) -> None:
    flag, difference = have_missing_keys(data=data, must_have_keys=must_have_keys)
    assert flag
    assert len(difference) > 0

tests/unit/utils/test_miscellaneous.py
KompleteAI/xllm
from peft import PeftModelForCausalLM, PeftType
from transformers import LlamaConfig, LlamaForCausalLM

from src.xllm.core.config import HuggingFaceConfig
from src.xllm.utils.nn import apply_lora, stabilize_training


def test_apply_lora(llama_model_config: LlamaConfig):
    config = HuggingFaceConfig(apply_lora=True, raw_lora_target_modules="all")
    model = LlamaForCausalLM(config=llama_model_config)
    peft_model, lora_config = apply_lora(config=config, model=model)
    assert isinstance(peft_model, PeftModelForCausalLM)
    assert peft_model.peft_type == PeftType.LORA


def test_stabilize_training(llama_model_config: LlamaConfig):
    model = LlamaForCausalLM(config=llama_model_config)
    stabilize_training(model=model)

tests/unit/utils/test_nn.py
KompleteAI/xllm
import pytest

from src.xllm.collators.lm import LMCollator
from src.xllm.utils.registry import Registry


def test_registry():
    registry = Registry(name="tmp", default_value=LMCollator)
    registry.add(key="a", value=1)
    a_value = registry.get(key="a")
    assert a_value == 1
    default_value = registry.get(key="b")
    default_value_from_none = registry.get(key=None)
    assert default_value == default_value_from_none


def test_registry_without_default():
    registry = Registry(name="tmp")
    with pytest.raises(ValueError):
        registry.get(key="a")

tests/unit/utils/test_registry.py
KompleteAI/xllm