diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..0002b90419b52774f2c3b45cdfa00b86e0819617 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +src/llamafactory/data/__pycache__/mm_plugin.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text +src/llamafactory/extras/__pycache__/constants.cpython-312.pyc filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..bff23355cf714397b5b0b93dd425a6e918ceb1fb --- /dev/null +++ b/README.md @@ -0,0 +1,966 @@ +![# LLaMA Factory](assets/logo.png) + +[![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/LLaMA-Factory?style=social)](https://github.com/hiyouga/LLaMA-Factory/stargazers) +[![GitHub last commit](https://img.shields.io/github/last-commit/hiyouga/LLaMA-Factory)](https://github.com/hiyouga/LLaMA-Factory/commits/main) +[![GitHub contributors](https://img.shields.io/github/contributors/hiyouga/LLaMA-Factory?color=orange)](https://github.com/hiyouga/LLaMA-Factory/graphs/contributors) +[![GitHub workflow](https://github.com/hiyouga/LLaMA-Factory/actions/workflows/tests.yml/badge.svg)](https://github.com/hiyouga/LLaMA-Factory/actions/workflows/tests.yml) +[![PyPI](https://img.shields.io/pypi/v/llamafactory)](https://pypi.org/project/llamafactory/) +[![Citation](https://img.shields.io/badge/citation-1000+-green)](https://scholar.google.com/scholar?cites=12620864006390196564) +[![Docker Pulls](https://img.shields.io/docker/pulls/hiyouga/llamafactory)](https://hub.docker.com/r/hiyouga/llamafactory/tags) + +[![Twitter](https://img.shields.io/twitter/follow/llamafactory_ai)](https://twitter.com/llamafactory_ai) +[![Discord](assets/thirdparty/discord.svg)](https://discord.gg/rKfvV9r9FK) +[![WeChat](https://img.shields.io/badge/WeChat-User%20Group-blue?logo=wechat)](https://github.com/hiyouga/llamafactory-community) +[![Blog](https://img.shields.io/badge/Hugo-Official%20Blog-blue?logo=hugo)](https://blog.llamafactory.net/en/) + +[![Open in Colab](assets/thirdparty/colab.svg)](https://colab.research.google.com/drive/1eRTPn37ltBbYsISy9Aw2NuI2Aq5CQrD9?usp=sharing) +[![Open in DSW](assets/thirdparty/dsw.svg)](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory) +[![Open in Spaces](https://img.shields.io/badge/🤗-Open%20in%20Spaces-blue)](https://huggingface.co/spaces/hiyouga/LLaMA-Board) +[![Open in Studios](https://img.shields.io/badge/ModelScope-Open%20in%20Studios-blue)](https://modelscope.cn/studios/hiyouga/LLaMA-Board) +[![Open in Novita](https://img.shields.io/badge/Novita-Deploy%20Template-blue)](https://novita.ai/templates-library/105981?sharer=88115474-394e-4bda-968e-b88e123d0c47) + +### Used by [Amazon](https://aws.amazon.com/cn/blogs/machine-learning/how-apoidea-group-enhances-visual-information-extraction-from-banking-documents-with-multimodal-models-using-llama-factory-on-amazon-sagemaker-hyperpod/), [NVIDIA](https://developer.nvidia.com/rtx/ai-toolkit), [Aliyun](https://help.aliyun.com/zh/pai/use-cases/fine-tune-a-llama-3-model-with-llama-factory), etc. + +
+ +### Supporters ❤️ + +|
Warp sponsorship
Warp, the agentic terminal for developers
Available for MacOS, Linux, & Windows | SerpAPI sponsorship | +| ---- | ---- | + +---- + +### Easily fine-tune 100+ large language models with zero-code [CLI](#quickstart) and [Web UI](#fine-tuning-with-llama-board-gui-powered-by-gradio) + +![GitHub Trend](https://trendshift.io/api/badge/repositories/4535) + +
+ +👋 Join our [WeChat](https://github.com/hiyouga/llamafactory-community/blob/main/wechat/main.jpg) and [NPU](https://github.com/hiyouga/llamafactory-community/blob/main/wechat/npu.jpg) user groups. + +\[ English | [中文](README_zh.md) \] + +**Fine-tuning a large language model can be easy as...** + +https://github.com/user-attachments/assets/3991a3a8-4276-4d30-9cab-4cb0c4b9b99e + +Start local training: +- Please refer to [usage](#getting-started) + +Start cloud training: +- **Colab (free)**: https://colab.research.google.com/drive/1eRTPn37ltBbYsISy9Aw2NuI2Aq5CQrD9?usp=sharing +- **PAI-DSW (free trial)**: https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory + +Read technical notes: +- **Documentation (WIP)**: https://llamafactory.readthedocs.io/en/latest/ +- **Documentation (AMD GPU)**: https://rocm.docs.amd.com/projects/ai-developer-hub/en/latest/notebooks/fine_tune/llama_factory_llama3.html +- **Documentation (ASCEND NPU)**: https://llamafactory.readthedocs.io/en/latest/multibackend/npu/index.html +- **Official Blog**: https://blog.llamafactory.net/en/ + +> [!NOTE] +> Except for the above links, all other websites are unauthorized third-party websites. Please carefully use them. + +## Table of Contents + +- [Features](#features) +- [Blogs](#blogs) +- [Changelog](#changelog) +- [Supported Models](#supported-models) +- [Supported Training Approaches](#supported-training-approaches) +- [Provided Datasets](#provided-datasets) +- [Requirement](#requirement) +- [Getting Started](#getting-started) + - [Installation](#installation) + - [Data Preparation](#data-preparation) + - [Quickstart](#quickstart) + - [Fine-Tuning with LLaMA Board GUI](#fine-tuning-with-llama-board-gui-powered-by-gradio) + - [Build Docker](#build-docker) + - [Deploy with OpenAI-style API and vLLM](#deploy-with-openai-style-api-and-vllm) + - [Download from ModelScope Hub](#download-from-modelscope-hub) + - [Download from Modelers Hub](#download-from-modelers-hub) + - [Use W&B Logger](#use-wb-logger) + - [Use SwanLab Logger](#use-swanlab-logger) +- [Projects using LLaMA Factory](#projects-using-llama-factory) +- [License](#license) +- [Citation](#citation) +- [Acknowledgement](#acknowledgement) + +## Features + +- **Various models**: LLaMA, LLaVA, Mistral, Mixtral-MoE, Qwen3, Qwen3-VL, DeepSeek, Gemma, GLM, Phi, etc. +- **Integrated methods**: (Continuous) pre-training, (multimodal) supervised fine-tuning, reward modeling, PPO, DPO, KTO, ORPO, etc. +- **Scalable resources**: 16-bit full-tuning, freeze-tuning, LoRA and 2/3/4/5/6/8-bit QLoRA via AQLM/AWQ/GPTQ/LLM.int8/HQQ/EETQ. +- **Advanced algorithms**: [GaLore](https://github.com/jiaweizzhao/GaLore), [BAdam](https://github.com/Ledzy/BAdam), [APOLLO](https://github.com/zhuhanqing/APOLLO), [Adam-mini](https://github.com/zyushun/Adam-mini), [Muon](https://github.com/KellerJordan/Muon), [OFT](https://github.com/huggingface/peft/tree/main/src/peft/tuners/oft), DoRA, LongLoRA, LLaMA Pro, Mixture-of-Depths, LoRA+, LoftQ and PiSSA. +- **Practical tricks**: [FlashAttention-2](https://github.com/Dao-AILab/flash-attention), [Unsloth](https://github.com/unslothai/unsloth), [Liger Kernel](https://github.com/linkedin/Liger-Kernel), [KTransformers](https://github.com/kvcache-ai/ktransformers/), RoPE scaling, NEFTune and rsLoRA. +- **Wide tasks**: Multi-turn dialogue, tool using, image understanding, visual grounding, video recognition, audio understanding, etc. +- **Experiment monitors**: LlamaBoard, TensorBoard, Wandb, MLflow, [SwanLab](https://github.com/SwanHubX/SwanLab), etc. +- **Faster inference**: OpenAI-style API, Gradio UI and CLI with [vLLM worker](https://github.com/vllm-project/vllm) or [SGLang worker](https://github.com/sgl-project/sglang). + +### Day-N Support for Fine-Tuning Cutting-Edge Models + +| Support Date | Model Name | +| ------------ | -------------------------------------------------------------------- | +| Day 0 | Qwen3 / Qwen2.5-VL / Gemma 3 / GLM-4.1V / InternLM 3 / MiniCPM-o-2.6 | +| Day 1 | Llama 3 / GLM-4 / Mistral Small / PaliGemma2 / Llama 4 | + +## Blogs + +> [!TIP] +> Now we have a dedicated blog for LLaMA Factory! +> +> Website: https://blog.llamafactory.net/en/ + +- 💡 [KTransformers Fine-Tuning × LLaMA Factory: Fine-tuning 1000 Billion models with 2 4090-GPU + CPU](https://blog.llamafactory.net/en/posts/ktransformers/) (English) +- 💡 [Easy Dataset × LLaMA Factory: Enabling LLMs to Efficiently Learn Domain Knowledge](https://buaa-act.feishu.cn/wiki/GVzlwYcRFiR8OLkHbL6cQpYin7g) (English) +- 💡 [DataFlow × LLaMA Factory: Producing High-Quality Data for LLM Training with a Data Preparation Pipeline](https://wcny4qa9krto.feishu.cn/wiki/LWkkwTDBfiiRKqkDSvucG6yjnbW) (English) | [中文](https://wcny4qa9krto.feishu.cn/wiki/LlMxweUAJimrmykRD5qcGuswnHd) +- 💡 [DataFlex × LLaMA Factory: A Data-Centric Dynamic Training System Built on LLaMA-Factory](https://wcny4qa9krto.feishu.cn/wiki/OlREwPQWdi9K6ZkJNHIciLhtnkv) (English) | [中文](https://wcny4qa9krto.feishu.cn/wiki/H2A9wSsbCinzavkT2oyc2C5Vn0e) +- [A One-Stop Code-Free Model Reinforcement Learning and Deployment Platform based on LLaMA-Factory and EasyR1](https://aws.amazon.com/cn/blogs/china/building-llm-model-hub-based-on-llamafactory-and-easyr1/) (Chinese) +- [How Apoidea Group enhances visual information extraction from banking documents with multimodal models using LLaMA-Factory on Amazon SageMaker HyperPod](https://aws.amazon.com/cn/blogs/machine-learning/how-apoidea-group-enhances-visual-information-extraction-from-banking-documents-with-multimodal-models-using-llama-factory-on-amazon-sagemaker-hyperpod/) (English) + +
All Blogs + +- [LLaMA Factory: Fine-tuning the DeepSeek-R1-Distill-Qwen-7B Model for News Classifier](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory_deepseek_r1_distill_7b) (Chinese) +- [A One-Stop Code-Free Model Fine-Tuning \& Deployment Platform based on SageMaker and LLaMA-Factory](https://aws.amazon.com/cn/blogs/china/a-one-stop-code-free-model-fine-tuning-deployment-platform-based-on-sagemaker-and-llama-factory/) (Chinese) +- [LLaMA Factory Multi-Modal Fine-Tuning Practice: Fine-Tuning Qwen2-VL for Personal Tourist Guide](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory_qwen2vl) (Chinese) +- [LLaMA Factory: Fine-tuning Llama3 for Role-Playing](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory) (Chinese) + +
+ +## Changelog + +[25/10/26] We support Megatron-core training backend with [**mcore_adapter**](https://github.com/alibaba/ROLL/tree/main/mcore_adapter). See [PR #9237](https://github.com/hiyouga/LLaMA-Factory/pull/9237) to get started. + +[25/08/22] We supported **[OFT](https://arxiv.org/abs/2306.07280)** and **[OFTv2](https://arxiv.org/abs/2506.19847)**. See [examples](examples/README.md) for usage. + +[25/08/20] We supported fine-tuning the **[Intern-S1-mini](https://huggingface.co/internlm/Intern-S1-mini)** models. See [PR #8976](https://github.com/hiyouga/LLaMA-Factory/pull/8976) to get started. + +[25/08/06] We supported fine-tuning the **[GPT-OSS](https://github.com/openai/gpt-oss)** models. See [PR #8826](https://github.com/hiyouga/LLaMA-Factory/pull/8826) to get started. + +
Full Changelog + +[25/07/02] We supported fine-tuning the **[GLM-4.1V-9B-Thinking](https://github.com/THUDM/GLM-4.1V-Thinking)** model. + +[25/04/28] We supported fine-tuning the **[Qwen3](https://qwenlm.github.io/blog/qwen3/)** model family. + +[25/04/21] We supported the **[Muon](https://github.com/KellerJordan/Muon)** optimizer. See [examples](examples/README.md) for usage. Thank [@tianshijing](https://github.com/tianshijing)'s PR. + +[25/04/16] We supported fine-tuning the **[InternVL3](https://huggingface.co/OpenGVLab/InternVL3-8B)** model. See [PR #7258](https://github.com/hiyouga/LLaMA-Factory/pull/7258) to get started. + +[25/04/14] We supported fine-tuning the **[GLM-Z1](https://huggingface.co/THUDM/GLM-Z1-9B-0414)** and **[Kimi-VL](https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct)** models. + +[25/04/06] We supported fine-tuning the **[Llama 4](https://ai.meta.com/blog/llama-4-multimodal-intelligence/)** model. See [PR #7611](https://github.com/hiyouga/LLaMA-Factory/pull/7611) to get started. + +[25/03/31] We supported fine-tuning the **[Qwen2.5 Omni](https://qwenlm.github.io/blog/qwen2.5-omni/)** model. See [PR #7537](https://github.com/hiyouga/LLaMA-Factory/pull/7537) to get started. + +[25/03/15] We supported **[SGLang](https://github.com/sgl-project/sglang)** as inference backend. Try `infer_backend: sglang` to accelerate inference. + +[25/03/12] We supported fine-tuning the **[Gemma 3](https://huggingface.co/blog/gemma3)** model. + +[25/02/24] Announcing **[EasyR1](https://github.com/hiyouga/EasyR1)**, an efficient, scalable and multi-modality RL training framework for efficient GRPO training. + +[25/02/11] We supported saving the **[Ollama](https://github.com/ollama/ollama)** modelfile when exporting the model checkpoints. See [examples](examples/README.md) for usage. + +[25/02/05] We supported fine-tuning the **[Qwen2-Audio](Qwen/Qwen2-Audio-7B-Instruct)** and **[MiniCPM-o-2.6](https://huggingface.co/openbmb/MiniCPM-o-2_6)** on audio understanding tasks. + +[25/01/31] We supported fine-tuning the **[DeepSeek-R1](https://huggingface.co/deepseek-ai/DeepSeek-R1)** and **[Qwen2.5-VL](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct)** models. + +[25/01/15] We supported **[APOLLO](https://arxiv.org/abs/2412.05270)** optimizer. See [examples](examples/README.md) for usage. + +[25/01/14] We supported fine-tuning the **[MiniCPM-o-2.6](https://huggingface.co/openbmb/MiniCPM-o-2_6)** and **[MiniCPM-V-2.6](https://huggingface.co/openbmb/MiniCPM-V-2_6)** models. Thank [@BUAADreamer](https://github.com/BUAADreamer)'s PR. + +[25/01/14] We supported fine-tuning the **[InternLM 3](https://huggingface.co/collections/internlm/)** models. Thank [@hhaAndroid](https://github.com/hhaAndroid)'s PR. + +[25/01/10] We supported fine-tuning the **[Phi-4](https://huggingface.co/microsoft/phi-4)** model. + +[24/12/21] We supported using **[SwanLab](https://github.com/SwanHubX/SwanLab)** for experiment tracking and visualization. See [this section](#use-swanlab-logger) for details. + +[24/11/27] We supported fine-tuning the **[Skywork-o1](https://huggingface.co/Skywork/Skywork-o1-Open-Llama-3.1-8B)** model and the **[OpenO1](https://huggingface.co/datasets/O1-OPEN/OpenO1-SFT)** dataset. + +[24/10/09] We supported downloading pre-trained models and datasets from the **[Modelers Hub](https://modelers.cn/models)**. See [this tutorial](#download-from-modelers-hub) for usage. + +[24/09/19] We supported fine-tuning the **[Qwen2.5](https://qwenlm.github.io/blog/qwen2.5/)** models. + +[24/08/30] We supported fine-tuning the **[Qwen2-VL](https://qwenlm.github.io/blog/qwen2-vl/)** models. Thank [@simonJJJ](https://github.com/simonJJJ)'s PR. + +[24/08/27] We supported **[Liger Kernel](https://github.com/linkedin/Liger-Kernel)**. Try `enable_liger_kernel: true` for efficient training. + +[24/08/09] We supported **[Adam-mini](https://github.com/zyushun/Adam-mini)** optimizer. See [examples](examples/README.md) for usage. Thank [@relic-yuexi](https://github.com/relic-yuexi)'s PR. + +[24/07/04] We supported [contamination-free packed training](https://github.com/MeetKai/functionary/tree/main/functionary/train/packing). Use `neat_packing: true` to activate it. Thank [@chuan298](https://github.com/chuan298)'s PR. + +[24/06/16] We supported **[PiSSA](https://arxiv.org/abs/2404.02948)** algorithm. See [examples](examples/README.md) for usage. + +[24/06/07] We supported fine-tuning the **[Qwen2](https://qwenlm.github.io/blog/qwen2/)** and **[GLM-4](https://github.com/THUDM/GLM-4)** models. + +[24/05/26] We supported **[SimPO](https://arxiv.org/abs/2405.14734)** algorithm for preference learning. See [examples](examples/README.md) for usage. + +[24/05/20] We supported fine-tuning the **PaliGemma** series models. Note that the PaliGemma models are pre-trained models, you need to fine-tune them with `paligemma` template for chat completion. + +[24/05/18] We supported **[KTO](https://arxiv.org/abs/2402.01306)** algorithm for preference learning. See [examples](examples/README.md) for usage. + +[24/05/14] We supported training and inference on the Ascend NPU devices. Check [installation](#installation) section for details. + +[24/04/26] We supported fine-tuning the **LLaVA-1.5** multimodal LLMs. See [examples](examples/README.md) for usage. + +[24/04/22] We provided a **[Colab notebook](https://colab.research.google.com/drive/1eRTPn37ltBbYsISy9Aw2NuI2Aq5CQrD9?usp=sharing)** for fine-tuning the Llama-3 model on a free T4 GPU. Two Llama-3-derived models fine-tuned using LLaMA Factory are available at Hugging Face, check [Llama3-8B-Chinese-Chat](https://huggingface.co/shenzhi-wang/Llama3-8B-Chinese-Chat) and [Llama3-Chinese](https://huggingface.co/zhichen/Llama3-Chinese) for details. + +[24/04/21] We supported **[Mixture-of-Depths](https://arxiv.org/abs/2404.02258)** according to [AstraMindAI's implementation](https://github.com/astramind-ai/Mixture-of-depths). See [examples](examples/README.md) for usage. + +[24/04/16] We supported **[BAdam](https://arxiv.org/abs/2404.02827)** optimizer. See [examples](examples/README.md) for usage. + +[24/04/16] We supported **[unsloth](https://github.com/unslothai/unsloth)**'s long-sequence training (Llama-2-7B-56k within 24GB). It achieves **117%** speed and **50%** memory compared with FlashAttention-2, more benchmarks can be found in [this page](https://github.com/hiyouga/LLaMA-Factory/wiki/Performance-comparison). + +[24/03/31] We supported **[ORPO](https://arxiv.org/abs/2403.07691)**. See [examples](examples/README.md) for usage. + +[24/03/21] Our paper "[LlamaFactory: Unified Efficient Fine-Tuning of 100+ Language Models](https://arxiv.org/abs/2403.13372)" is available at arXiv! + +[24/03/20] We supported **FSDP+QLoRA** that fine-tunes a 70B model on 2x24GB GPUs. See [examples](examples/README.md) for usage. + +[24/03/13] We supported **[LoRA+](https://arxiv.org/abs/2402.12354)**. See [examples](examples/README.md) for usage. + +[24/03/07] We supported **[GaLore](https://arxiv.org/abs/2403.03507)** optimizer. See [examples](examples/README.md) for usage. + +[24/03/07] We integrated **[vLLM](https://github.com/vllm-project/vllm)** for faster and concurrent inference. Try `infer_backend: vllm` to enjoy **270%** inference speed. + +[24/02/28] We supported weight-decomposed LoRA (**[DoRA](https://arxiv.org/abs/2402.09353)**). Try `use_dora: true` to activate DoRA training. + +[24/02/15] We supported **block expansion** proposed by [LLaMA Pro](https://github.com/TencentARC/LLaMA-Pro). See [examples](examples/README.md) for usage. + +[24/02/05] Qwen1.5 (Qwen2 beta version) series models are supported in LLaMA-Factory. Check this [blog post](https://qwenlm.github.io/blog/qwen1.5/) for details. + +[24/01/18] We supported **agent tuning** for most models, equipping model with tool using abilities by fine-tuning with `dataset: glaive_toolcall_en`. + +[23/12/23] We supported **[unsloth](https://github.com/unslothai/unsloth)**'s implementation to boost LoRA tuning for the LLaMA, Mistral and Yi models. Try `use_unsloth: true` argument to activate unsloth patch. It achieves **170%** speed in our benchmark, check [this page](https://github.com/hiyouga/LLaMA-Factory/wiki/Performance-comparison) for details. + +[23/12/12] We supported fine-tuning the latest MoE model **[Mixtral 8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1)** in our framework. See hardware requirement [here](#hardware-requirement). + +[23/12/01] We supported downloading pre-trained models and datasets from the **[ModelScope Hub](https://modelscope.cn/models)**. See [this tutorial](#download-from-modelscope-hub) for usage. + +[23/10/21] We supported **[NEFTune](https://arxiv.org/abs/2310.05914)** trick for fine-tuning. Try `neftune_noise_alpha: 5` argument to activate NEFTune. + +[23/09/27] We supported **$S^2$-Attn** proposed by [LongLoRA](https://github.com/dvlab-research/LongLoRA) for the LLaMA models. Try `shift_attn: true` argument to enable shift short attention. + +[23/09/23] We integrated MMLU, C-Eval and CMMLU benchmarks in this repo. See [examples](examples/README.md) for usage. + +[23/09/10] We supported **[FlashAttention-2](https://github.com/Dao-AILab/flash-attention)**. Try `flash_attn: fa2` argument to enable FlashAttention-2 if you are using RTX4090, A100 or H100 GPUs. + +[23/08/12] We supported **RoPE scaling** to extend the context length of the LLaMA models. Try `rope_scaling: linear` argument in training and `rope_scaling: dynamic` argument at inference to extrapolate the position embeddings. + +[23/08/11] We supported **[DPO training](https://arxiv.org/abs/2305.18290)** for instruction-tuned models. See [examples](examples/README.md) for usage. + +[23/07/31] We supported **dataset streaming**. Try `streaming: true` and `max_steps: 10000` arguments to load your dataset in streaming mode. + +[23/07/29] We released two instruction-tuned 13B models at Hugging Face. See these Hugging Face Repos ([LLaMA-2](https://huggingface.co/hiyouga/Llama-2-Chinese-13b-chat) / [Baichuan](https://huggingface.co/hiyouga/Baichuan-13B-sft)) for details. + +[23/07/18] We developed an **all-in-one Web UI** for training, evaluation and inference. Try `train_web.py` to fine-tune models in your Web browser. Thank [@KanadeSiina](https://github.com/KanadeSiina) and [@codemayq](https://github.com/codemayq) for their efforts in the development. + +[23/07/09] We released **[FastEdit](https://github.com/hiyouga/FastEdit)** ⚡🩹, an easy-to-use package for editing the factual knowledge of large language models efficiently. Please follow [FastEdit](https://github.com/hiyouga/FastEdit) if you are interested. + +[23/06/29] We provided a **reproducible example** of training a chat model using instruction-following datasets, see [Baichuan-7B-sft](https://huggingface.co/hiyouga/Baichuan-7B-sft) for details. + +[23/06/22] We aligned the [demo API](src/api_demo.py) with the [OpenAI's](https://platform.openai.com/docs/api-reference/chat) format where you can insert the fine-tuned model in **arbitrary ChatGPT-based applications**. + +[23/06/03] We supported quantized training and inference (aka **[QLoRA](https://github.com/artidoro/qlora)**). See [examples](examples/README.md) for usage. + +
+ +> [!TIP] +> If you cannot use the latest feature, please pull the latest code and install LLaMA-Factory again. + +## Supported Models + +| Model | Model size | Template | +| ----------------------------------------------------------------- | -------------------------------- | -------------------- | +| [BLOOM/BLOOMZ](https://huggingface.co/bigscience) | 560M/1.1B/1.7B/3B/7.1B/176B | - | +| [DeepSeek (LLM/Code/MoE)](https://huggingface.co/deepseek-ai) | 7B/16B/67B/236B | deepseek | +| [DeepSeek 3-3.2](https://huggingface.co/deepseek-ai) | 236B/671B | deepseek3 | +| [DeepSeek R1 (Distill)](https://huggingface.co/deepseek-ai) | 1.5B/7B/8B/14B/32B/70B/671B | deepseekr1 | +| [ERNIE-4.5](https://huggingface.co/baidu) | 0.3B/21B/300B | ernie_nothink | +| [Falcon/Falcon H1](https://huggingface.co/tiiuae) | 0.5B/1.5B/3B/7B/11B/34B/40B/180B | falcon/falcon_h1 | +| [Gemma/Gemma 2/CodeGemma](https://huggingface.co/google) | 2B/7B/9B/27B | gemma/gemma2 | +| [Gemma 3/Gemma 3n](https://huggingface.co/google) | 270M/1B/4B/6B/8B/12B/27B | gemma3/gemma3n | +| [GLM-4/GLM-4-0414/GLM-Z1](https://huggingface.co/zai-org) | 9B/32B | glm4/glmz1 | +| [GLM-4.5/GLM-4.5(6)V](https://huggingface.co/zai-org) | 9B/106B/355B | glm4_moe/glm4_5v | +| [GPT-2](https://huggingface.co/openai-community) | 0.1B/0.4B/0.8B/1.5B | - | +| [GPT-OSS](https://huggingface.co/openai) | 20B/120B | gpt_oss | +| [Granite 3-4](https://huggingface.co/ibm-granite) | 1B/2B/3B/7B/8B | granite3/granite4 | +| [Hunyuan/Hunyuan1.5 (MT)](https://huggingface.co/tencent/) | 0.5B/1.8B/4B/7B/13B | hunyuan/hunyuan_small| +| [InternLM 2-3](https://huggingface.co/internlm) | 7B/8B/20B | intern2 | +| [InternVL 2.5-3.5](https://huggingface.co/OpenGVLab) | 1B/2B/4B/8B/14B/30B/38B/78B/241B | intern_vl | +| [Intern-S1-mini](https://huggingface.co/internlm/) | 8B | intern_s1 | +| [Kimi-VL](https://huggingface.co/moonshotai) | 16B | kimi_vl | +| [Ling 2.0 (mini/flash)](https://huggingface.co/inclusionAI) | 16B/100B | bailing_v2 | +| [LFM 2.5 (VL)](https://huggingface.co/LiquidAI) | 1.2B/1.6B | lfm2/lfm2_vl | +| [Llama](https://github.com/facebookresearch/llama) | 7B/13B/33B/65B | - | +| [Llama 2](https://huggingface.co/meta-llama) | 7B/13B/70B | llama2 | +| [Llama 3-3.3](https://huggingface.co/meta-llama) | 1B/3B/8B/70B | llama3 | +| [Llama 4](https://huggingface.co/meta-llama) | 109B/402B | llama4 | +| [Llama 3.2 Vision](https://huggingface.co/meta-llama) | 11B/90B | mllama | +| [LLaVA-1.5](https://huggingface.co/llava-hf) | 7B/13B | llava | +| [LLaVA-NeXT](https://huggingface.co/llava-hf) | 7B/8B/13B/34B/72B/110B | llava_next | +| [LLaVA-NeXT-Video](https://huggingface.co/llava-hf) | 7B/34B | llava_next_video | +| [MiMo](https://huggingface.co/XiaomiMiMo) | 7B/309B | mimo/mimo_v2 | +| [MiniCPM 4/5](https://huggingface.co/openbmb) | 0.5B/1B/8B | cpm4/empty | +| [MiniCPM-o/MiniCPM-V 4.5](https://huggingface.co/openbmb) | 8B/9B | minicpm_o/minicpm_v | +| [MiniCPM-V 4.6](https://huggingface.co/openbmb) | 3B/8B | minicpm_v_4_6 | +| [MiniMax-M1/MiniMax-M2](https://huggingface.co/MiniMaxAI/models) | 229B/456B | minimax1/minimax2 | +| [Ministral 3](https://huggingface.co/mistralai) | 3B/8B/14B | ministral3 | +| [Mistral/Mixtral](https://huggingface.co/mistralai) | 7B/8x7B/8x22B | mistral | +| [PaliGemma/PaliGemma2](https://huggingface.co/google) | 3B/10B/28B | paligemma | +| [Phi-3/Phi-3.5](https://huggingface.co/microsoft) | 4B/14B | phi | +| [Phi-3-small](https://huggingface.co/microsoft) | 7B | phi_small | +| [Phi-4-mini/Phi-4](https://huggingface.co/microsoft) | 3.8B/14B | phi4_mini/phi4 | +| [Pixtral](https://huggingface.co/mistralai) | 12B | pixtral | +| [Qwen2 (Code/Math/MoE/QwQ)](https://huggingface.co/Qwen) | 0.5B/1.5B/3B/7B/14B/32B/72B/110B | qwen | +| [Qwen3 (MoE/Instruct/Thinking/Next)](https://huggingface.co/Qwen) | 0.6B/1.7B/4B/8B/14B/32B/80B/235B | qwen3/qwen3_nothink | +| [Qwen3.5](https://huggingface.co/Qwen) | 0.8B/2B/4B/9B/27B/35B/122B/397B | qwen3_5/qwen3_5_nothink | +| [Qwen3.6](https://huggingface.co/Qwen) | 27B/35B | qwen3_6 | +| [Qwen2-Audio](https://huggingface.co/Qwen) | 7B | qwen2_audio | +| [Qwen2.5-Omni](https://huggingface.co/Qwen) | 3B/7B | qwen2_omni | +| [Qwen3-Omni](https://huggingface.co/Qwen) | 30B | qwen3_omni | +| [Qwen2-VL/Qwen2.5-VL/QVQ](https://huggingface.co/Qwen) | 2B/3B/7B/32B/72B | qwen2_vl | +| [Qwen3-VL](https://huggingface.co/Qwen) | 2B/4B/8B/30B/32B/235B | qwen3_vl | +| [Seed (OSS/Coder)](https://huggingface.co/ByteDance-Seed) | 8B/36B | seed_oss/seed_coder | +| [StarCoder 2](https://huggingface.co/bigcode) | 3B/7B/15B | - | +| [TeleChat 2-2.5](https://huggingface.co/Tele-AI) | 3B/7B/35B/115B | telechat2 | +| [Yuan 2](https://huggingface.co/IEITYuan) | 2B/51B/102B | yuan | + +> [!NOTE] +> For the "base" models, the `template` argument can be chosen from `default`, `alpaca`, `vicuna` etc. But make sure to use the **corresponding template** for the "instruct/chat" models. +> +> If the model has both reasoning and non-reasoning versions, please use the `_nothink` suffix to distinguish between them. For example, `qwen3` and `qwen3_nothink`. +> +> Remember to use the **SAME** template in training and inference. +> +> \*: You should install the `transformers` from main branch and use `DISABLE_VERSION_CHECK=1` to skip version check. +> +> \*\*: You need to install a specific version of `transformers` to use the corresponding model. + +Please refer to [constants.py](src/llamafactory/extras/constants.py) for a full list of models we supported. + +You also can add a custom chat template to [template.py](src/llamafactory/data/template.py). + +## Supported Training Approaches + +| Approach | Full-tuning | Freeze-tuning | LoRA | QLoRA | OFT | QOFT | +| ---------------------- | ------------------ | ------------------ | ------------------ | ------------------ | ------------------ | ------------------ | +| Pre-Training | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| Supervised Fine-Tuning | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| Reward Modeling | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| PPO Training | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| DPO Training | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| KTO Training | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| ORPO Training | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | +| SimPO Training | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: | + +> [!TIP] +> The implementation details of PPO can be found in [this blog](https://newfacade.github.io/notes-on-reinforcement-learning/17-ppo-trl.html). + +## Provided Datasets + +
Pre-training datasets + +- [Wiki Demo (en)](data/wiki_demo.txt) +- [RefinedWeb (en)](https://huggingface.co/datasets/tiiuae/falcon-refinedweb) +- [RedPajama V2 (en)](https://huggingface.co/datasets/togethercomputer/RedPajama-Data-V2) +- [Wikipedia (en)](https://huggingface.co/datasets/olm/olm-wikipedia-20221220) +- [Wikipedia (zh)](https://huggingface.co/datasets/pleisto/wikipedia-cn-20230720-filtered) +- [Pile (en)](https://huggingface.co/datasets/EleutherAI/pile) +- [SkyPile (zh)](https://huggingface.co/datasets/Skywork/SkyPile-150B) +- [FineWeb (en)](https://huggingface.co/datasets/HuggingFaceFW/fineweb) +- [FineWeb-Edu (en)](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) +- [CCI3-HQ (zh)](https://huggingface.co/datasets/BAAI/CCI3-HQ) +- [CCI3-Data (zh)](https://huggingface.co/datasets/BAAI/CCI3-Data) +- [CCI4.0-M2-Base-v1 (en&zh)](https://huggingface.co/datasets/BAAI/CCI4.0-M2-Base-v1) +- [CCI4.0-M2-CoT-v1 (en&zh)](https://huggingface.co/datasets/BAAI/CCI4.0-M2-CoT-v1) +- [CCI4.0-M2-Extra-v1 (en&zh)](https://huggingface.co/datasets/BAAI/CCI4.0-M2-Extra-v1) +- [The Stack (en)](https://huggingface.co/datasets/bigcode/the-stack) +- [StarCoder (en)](https://huggingface.co/datasets/bigcode/starcoderdata) + +
+ +
Supervised fine-tuning datasets + +- [Identity (en&zh)](data/identity.json) +- [Stanford Alpaca (en)](https://github.com/tatsu-lab/stanford_alpaca) +- [Stanford Alpaca (zh)](https://github.com/ymcui/Chinese-LLaMA-Alpaca-3) +- [Alpaca GPT4 (en&zh)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM) +- [Glaive Function Calling V2 (en&zh)](https://huggingface.co/datasets/glaiveai/glaive-function-calling-v2) +- [LIMA (en)](https://huggingface.co/datasets/GAIR/lima) +- [Guanaco Dataset (multilingual)](https://huggingface.co/datasets/JosephusCheung/GuanacoDataset) +- [BELLE 2M (zh)](https://huggingface.co/datasets/BelleGroup/train_2M_CN) +- [BELLE 1M (zh)](https://huggingface.co/datasets/BelleGroup/train_1M_CN) +- [BELLE 0.5M (zh)](https://huggingface.co/datasets/BelleGroup/train_0.5M_CN) +- [BELLE Dialogue 0.4M (zh)](https://huggingface.co/datasets/BelleGroup/generated_chat_0.4M) +- [BELLE School Math 0.25M (zh)](https://huggingface.co/datasets/BelleGroup/school_math_0.25M) +- [BELLE Multiturn Chat 0.8M (zh)](https://huggingface.co/datasets/BelleGroup/multiturn_chat_0.8M) +- [UltraChat (en)](https://github.com/thunlp/UltraChat) +- [OpenPlatypus (en)](https://huggingface.co/datasets/garage-bAInd/Open-Platypus) +- [CodeAlpaca 20k (en)](https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k) +- [Alpaca CoT (multilingual)](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT) +- [OpenOrca (en)](https://huggingface.co/datasets/Open-Orca/OpenOrca) +- [SlimOrca (en)](https://huggingface.co/datasets/Open-Orca/SlimOrca) +- [MathInstruct (en)](https://huggingface.co/datasets/TIGER-Lab/MathInstruct) +- [Firefly 1.1M (zh)](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M) +- [Wiki QA (en)](https://huggingface.co/datasets/wiki_qa) +- [Web QA (zh)](https://huggingface.co/datasets/suolyer/webqa) +- [WebNovel (zh)](https://huggingface.co/datasets/zxbsmk/webnovel_cn) +- [Nectar (en)](https://huggingface.co/datasets/berkeley-nest/Nectar) +- [deepctrl (en&zh)](https://www.modelscope.cn/datasets/deepctrl/deepctrl-sft-data) +- [Advertise Generating (zh)](https://huggingface.co/datasets/HasturOfficial/adgen) +- [ShareGPT Hyperfiltered (en)](https://huggingface.co/datasets/totally-not-an-llm/sharegpt-hyperfiltered-3k) +- [ShareGPT4 (en&zh)](https://huggingface.co/datasets/shibing624/sharegpt_gpt4) +- [UltraChat 200k (en)](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) +- [Infinity Instruct (zh)](https://huggingface.co/datasets/BAAI/Infinity-Instruct) +- [AgentInstruct (en)](https://huggingface.co/datasets/THUDM/AgentInstruct) +- [LMSYS Chat 1M (en)](https://huggingface.co/datasets/lmsys/lmsys-chat-1m) +- [Evol Instruct V2 (en)](https://huggingface.co/datasets/WizardLM/WizardLM_evol_instruct_V2_196k) +- [Cosmopedia (en)](https://huggingface.co/datasets/HuggingFaceTB/cosmopedia) +- [STEM (zh)](https://huggingface.co/datasets/hfl/stem_zh_instruction) +- [Ruozhiba (zh)](https://huggingface.co/datasets/hfl/ruozhiba_gpt4_turbo) +- [Neo-sft (zh)](https://huggingface.co/datasets/m-a-p/neo_sft_phase2) +- [Magpie-Pro-300K-Filtered (en)](https://huggingface.co/datasets/Magpie-Align/Magpie-Pro-300K-Filtered) +- [Magpie-ultra-v0.1 (en)](https://huggingface.co/datasets/argilla/magpie-ultra-v0.1) +- [WebInstructSub (en)](https://huggingface.co/datasets/TIGER-Lab/WebInstructSub) +- [OpenO1-SFT (en&zh)](https://huggingface.co/datasets/O1-OPEN/OpenO1-SFT) +- [Open-Thoughts (en)](https://huggingface.co/datasets/open-thoughts/OpenThoughts-114k) +- [Open-R1-Math (en)](https://huggingface.co/datasets/open-r1/OpenR1-Math-220k) +- [Chinese-DeepSeek-R1-Distill (zh)](https://huggingface.co/datasets/Congliu/Chinese-DeepSeek-R1-Distill-data-110k-SFT) +- [LLaVA mixed (en&zh)](https://huggingface.co/datasets/BUAADreamer/llava-en-zh-300k) +- [Pokemon-gpt4o-captions (en&zh)](https://huggingface.co/datasets/jugg1024/pokemon-gpt4o-captions) +- [DLR-Web (en)](https://huggingface.co/datasets/Attention1115/DLR-Web) +- [Open Assistant (de)](https://huggingface.co/datasets/mayflowergmbh/oasst_de) +- [Dolly 15k (de)](https://huggingface.co/datasets/mayflowergmbh/dolly-15k_de) +- [Alpaca GPT4 (de)](https://huggingface.co/datasets/mayflowergmbh/alpaca-gpt4_de) +- [OpenSchnabeltier (de)](https://huggingface.co/datasets/mayflowergmbh/openschnabeltier_de) +- [Evol Instruct (de)](https://huggingface.co/datasets/mayflowergmbh/evol-instruct_de) +- [Dolphin (de)](https://huggingface.co/datasets/mayflowergmbh/dolphin_de) +- [Booksum (de)](https://huggingface.co/datasets/mayflowergmbh/booksum_de) +- [Airoboros (de)](https://huggingface.co/datasets/mayflowergmbh/airoboros-3.0_de) +- [Ultrachat (de)](https://huggingface.co/datasets/mayflowergmbh/ultra-chat_de) + +
+ +
Preference datasets + +- [DPO mixed (en&zh)](https://huggingface.co/datasets/hiyouga/DPO-En-Zh-20k) +- [UltraFeedback (en)](https://huggingface.co/datasets/HuggingFaceH4/ultrafeedback_binarized) +- [COIG-P (zh)](https://huggingface.co/datasets/m-a-p/COIG-P) +- [RLHF-V (en)](https://huggingface.co/datasets/openbmb/RLHF-V-Dataset) +- [VLFeedback (en)](https://huggingface.co/datasets/Zhihui/VLFeedback) +- [RLAIF-V (en)](https://huggingface.co/datasets/openbmb/RLAIF-V-Dataset) +- [Orca DPO Pairs (en)](https://huggingface.co/datasets/Intel/orca_dpo_pairs) +- [HH-RLHF (en)](https://huggingface.co/datasets/Anthropic/hh-rlhf) +- [Nectar (en)](https://huggingface.co/datasets/berkeley-nest/Nectar) +- [Orca DPO (de)](https://huggingface.co/datasets/mayflowergmbh/intel_orca_dpo_pairs_de) +- [KTO mixed (en)](https://huggingface.co/datasets/argilla/kto-mix-15k) + +
+ +Some datasets require confirmation before using them, so we recommend logging in with your Hugging Face account using these commands. + +```bash +pip install "huggingface_hub<1.0.0" +huggingface-cli login +``` + +## Requirement + +| Mandatory | Minimum | Recommend | +| ------------ | ------- | --------- | +| python | 3.11 | >=3.11 | +| torch | 2.0.0 | 2.6.0 | +| torchvision | 0.15.0 | 0.21.0 | +| transformers | 4.49.0 | 4.50.0 | +| datasets | 2.16.0 | 3.2.0 | +| accelerate | 0.34.0 | 1.2.1 | +| peft | 0.14.0 | 0.15.1 | +| trl | 0.8.6 | 0.9.6 | + +| Optional | Minimum | Recommend | +| ------------ | ------- | --------- | +| CUDA | 11.6 | 12.2 | +| deepspeed | 0.10.0 | 0.16.4 | +| bitsandbytes | 0.39.0 | 0.43.1 | +| vllm | 0.4.3 | 0.8.2 | +| flash-attn | 2.5.6 | 2.7.2 | + +### Hardware Requirement + +\* *estimated* + +| Method | Bits | 7B | 14B | 30B | 70B | `x`B | +| ----------------------------------- | ---- | ----- | ----- | ----- | ------ | ------- | +| Full (`bf16` or `fp16`) | 32 | 120GB | 240GB | 600GB | 1200GB | `18x`GB | +| Full (`pure_bf16`) | 16 | 60GB | 120GB | 300GB | 600GB | `8x`GB | +| Freeze/LoRA/GaLore/APOLLO/BAdam/OFT | 16 | 16GB | 32GB | 64GB | 160GB | `2x`GB | +| QLoRA / QOFT | 8 | 10GB | 20GB | 40GB | 80GB | `x`GB | +| QLoRA / QOFT | 4 | 6GB | 12GB | 24GB | 48GB | `x/2`GB | +| QLoRA / QOFT | 2 | 4GB | 8GB | 16GB | 24GB | `x/4`GB | + +## Getting Started + +### Installation + +> [!IMPORTANT] +> Installation is mandatory. + +#### Install from Source + +```bash +git clone --depth 1 https://github.com/hiyouga/LlamaFactory.git +cd LlamaFactory +pip install -e . +pip install -r requirements/metrics.txt +``` + +Optional dependencies available: `metrics`, `deepspeed`. Install with: `pip install -e . && pip install -r requirements/metrics.txt -r requirements/deepspeed.txt` + +Additional dependencies for specific features are available in `examples/requirements/`. + +#### Install from Docker Image + +```bash +docker run -it --rm --gpus=all --ipc=host hiyouga/llamafactory:latest +``` + +This image is built on Ubuntu 22.04 (x86\_64), CUDA 12.4, Python 3.11, PyTorch 2.6.0, and Flash-attn 2.7.4. + +Find the pre-built images: https://hub.docker.com/r/hiyouga/llamafactory/tags + +Please refer to [build docker](#build-docker) to build the image yourself. + +
Setting up a virtual environment with uv + +Create an isolated Python environment with [uv](https://github.com/astral-sh/uv): + +```bash +uv run llamafactory-cli webui +``` + +
+ +
For Windows users + +#### Install PyTorch + +You need to manually install the GPU version of PyTorch on the Windows platform. Please refer to the [official website](https://pytorch.org/get-started/locally/) and the following command to install PyTorch with CUDA support: + +```bash +pip uninstall torch torchvision torchaudio +pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126 +python -c "import torch; print(torch.cuda.is_available())" +``` + +If you see `True` then you have successfully installed PyTorch with CUDA support. + +Try `dataloader_num_workers: 0` if you encounter `Can't pickle local object` error. + +#### Install BitsAndBytes + +To enable Quantized LoRA (QLoRA) on Windows, you need to install bitsandbytes. + +For most users, it is recommended to install the latest official release: + +```bash +pip install bitsandbytes +``` + +If you are using uv to manage your virtual environment, it is recommended to install bitsandbytes after installing the GPU-enabled version of PyTorch: + +```bash +uv pip install bitsandbytes --no-deps +``` + +[!IMPORTANT] +Pay attention to the CUDA Toolkit version when installing bitsandbytes. Official bitsandbytes releases are built for specific CUDA Toolkit versions. On Windows x86-64, separate builds are currently provided for CUDA 11.8–12.6 and CUDA 12.8–12.9. Support for NVIDIA RTX 50 Series GPUs (e.g., RTX 5060 Ti, sm_120) requires the CUDA 12.8–12.9 builds. + +If your environment uses an older CUDA version, or you need compatibility with older Windows / PyTorch combinations, you can install the third-party precompiled Windows wheel: + +```bash +pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.41.2.post2-py3-none-win_amd64.whl +``` + +#### Install Flash Attention-2 + +To enable FlashAttention-2 on the Windows platform, please use the script from [flash-attention-windows-wheel](https://huggingface.co/lldacing/flash-attention-windows-wheel) to compile and install it by yourself. + +
+ +
For Ascend NPU users + +To install LLaMA Factory on Ascend NPU devices, please upgrade Python to version 3.10 or higher: `pip install -r requirements/npu.txt`. Additionally, you need to install the **Ascend CANN Toolkit and Kernels**. Please follow the [installation tutorial](https://llamafactory.readthedocs.io/en/latest/advanced/npu_installation.html). + + +You can also download the pre-built Docker images: + +```bash +# Docker Hub +docker pull hiyouga/llamafactory:latest-npu-a2 +docker pull hiyouga/llamafactory:latest-npu-a3 + +# quay.io +docker pull quay.io/ascend/llamafactory:latest-npu-a2 +docker pull quay.io/ascend/llamafactory:latest-npu-a3 +``` + +#### Install BitsAndBytes + +To use QLoRA based on bitsandbytes on Ascend NPU, please follow these 3 steps: + +1. Manually compile bitsandbytes: Refer to [the installation documentation](https://huggingface.co/docs/bitsandbytes/installation?backend=Ascend+NPU&platform=Ascend+NPU) for the NPU version of bitsandbytes to complete the compilation and installation. The compilation requires a cmake version of at least 3.22.1 and a g++ version of at least 12.x. + +```bash +# Install bitsandbytes from source +# Clone bitsandbytes repo, Ascend NPU backend is currently enabled on multi-backend-refactor branch +git clone -b multi-backend-refactor https://github.com/bitsandbytes-foundation/bitsandbytes.git +cd bitsandbytes/ + +# Install dependencies +pip install -r requirements-dev.txt + +# Install the dependencies for the compilation tools. Note that the commands for this step may vary depending on the operating system. The following are provided for reference +apt-get install -y build-essential cmake + +# Compile & install +cmake -DCOMPUTE_BACKEND=npu -S . +make +pip install . +``` + +2. Install transformers from the main branch. + +```bash +git clone -b main https://github.com/huggingface/transformers.git +cd transformers +pip install . +``` + +3. Set `double_quantization: false` in the configuration. You can refer to the [example](examples/train_qlora/qwen3_lora_sft_bnb_npu.yaml). + +
+ +### Data Preparation + +Please refer to [data/README.md](data/README.md) for checking the details about the format of dataset files. You can use datasets on HuggingFace / ModelScope / Modelers hub, load the dataset in local disk, or specify a path to s3/gcs cloud storage. + +> [!NOTE] +> Please update `data/dataset_info.json` to use your custom dataset. + +You can also use **[Easy Dataset](https://github.com/ConardLi/easy-dataset)**, **[DataFlow](https://github.com/OpenDCAI/DataFlow)** and **[GraphGen](https://github.com/open-sciencelab/GraphGen)** to create synthetic data for fine-tuning. + +### Quickstart + +Use the following 3 commands to run LoRA **fine-tuning**, **inference** and **merging** of the Qwen3-4B-Instruct model, respectively. + +```bash +llamafactory-cli train examples/train_lora/qwen3_lora_sft.yaml +llamafactory-cli chat examples/inference/qwen3_lora_sft.yaml +llamafactory-cli export examples/merge_lora/qwen3_lora_sft.yaml +``` + +See [examples/README.md](examples/README.md) for advanced usage (including distributed training). + +> [!TIP] +> Use `llamafactory-cli help` to show help information. +> +> Read [FAQs](https://github.com/hiyouga/LLaMA-Factory/issues/4614) first if you encounter any problems. + +### Fine-Tuning with LLaMA Board GUI (powered by [Gradio](https://github.com/gradio-app/gradio)) + +```bash +llamafactory-cli webui +``` + +### Build Docker + +For CUDA users: + +```bash +cd docker/docker-cuda/ +docker compose up -d +docker compose exec llamafactory bash +``` + +For Ascend NPU users: + +```bash +cd docker/docker-npu/ +docker compose up -d +docker compose exec llamafactory bash +``` + +For AMD ROCm users: + +```bash +cd docker/docker-rocm/ +docker compose up -d +docker compose exec llamafactory bash +``` + +
Build without Docker Compose + +For CUDA users: + +```bash +docker build -f ./docker/docker-cuda/Dockerfile \ + --build-arg PIP_INDEX=https://pypi.org/simple \ + -t llamafactory:latest . + +docker run -dit --ipc=host --gpus=all \ + -p 7860:7860 \ + -p 8000:8000 \ + --name llamafactory \ + llamafactory:latest + +docker exec -it llamafactory bash +``` + +For Ascend NPU users: + +```bash +docker build -f ./docker/docker-npu/Dockerfile \ + --build-arg PIP_INDEX=https://pypi.org/simple \ + -t llamafactory:latest . + +docker run -dit --ipc=host \ + -v /usr/local/dcmi:/usr/local/dcmi \ + -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ + -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \ + -v /etc/ascend_install.info:/etc/ascend_install.info \ + -p 7860:7860 \ + -p 8000:8000 \ + --device /dev/davinci0 \ + --device /dev/davinci_manager \ + --device /dev/devmm_svm \ + --device /dev/hisi_hdc \ + --name llamafactory \ + llamafactory:latest + +docker exec -it llamafactory bash +``` + +For AMD ROCm users: + +```bash +docker build -f ./docker/docker-rocm/Dockerfile \ + --build-arg PIP_INDEX=https://pypi.org/simple \ + -t llamafactory:latest . + +docker run -dit --ipc=host \ + -p 7860:7860 \ + -p 8000:8000 \ + --device /dev/kfd \ + --device /dev/dri \ + --name llamafactory \ + llamafactory:latest + +docker exec -it llamafactory bash +``` + +
+ +
Use Docker volumes + +You can uncomment `VOLUME [ "/root/.cache/huggingface", "/app/shared_data", "/app/output" ]` in the Dockerfile to use data volumes. + +When building the Docker image, use `-v ./hf_cache:/root/.cache/huggingface` argument to mount the local directory to the container. The following data volumes are available. + +- `hf_cache`: Utilize Hugging Face cache on the host machine. +- `shared_data`: The directionary to store datasets on the host machine. +- `output`: Set export dir to this location so that the merged result can be accessed directly on the host machine. + +
+ +### Deploy with OpenAI-style API and vLLM + +```bash +API_PORT=8000 llamafactory-cli api examples/inference/qwen3.yaml infer_backend=vllm vllm_enforce_eager=true +``` + +> [!TIP] +> Visit [this page](https://platform.openai.com/docs/api-reference/chat/create) for API document. +> +> Examples: [Image understanding](scripts/api_example/test_image.py) | [Function calling](scripts/api_example/test_toolcall.py) + +### Download from ModelScope Hub + +If you have trouble with downloading models and datasets from Hugging Face, you can use ModelScope. + +```bash +export USE_MODELSCOPE_HUB=1 # `set USE_MODELSCOPE_HUB=1` for Windows +``` + +Train the model by specifying a model ID of the ModelScope Hub as the `model_name_or_path`. You can find a full list of model IDs at [ModelScope Hub](https://modelscope.cn/models), e.g., `LLM-Research/Meta-Llama-3-8B-Instruct`. + +### Download from Modelers Hub + +You can also use Modelers Hub to download models and datasets. + +```bash +export USE_OPENMIND_HUB=1 # `set USE_OPENMIND_HUB=1` for Windows +``` + +Train the model by specifying a model ID of the Modelers Hub as the `model_name_or_path`. You can find a full list of model IDs at [Modelers Hub](https://modelers.cn/models), e.g., `TeleAI/TeleChat-7B-pt`. + +### Use W&B Logger + +To use [Weights & Biases](https://wandb.ai) for logging experimental results, you need to add the following arguments to yaml files. + +```yaml +report_to: wandb +run_name: test_run # optional +``` + +Set `WANDB_API_KEY` to [your key](https://wandb.ai/authorize) when launching training tasks to log in with your W&B account. + +### Use SwanLab Logger + +To use [SwanLab](https://github.com/SwanHubX/SwanLab) for logging experimental results, you need to add the following arguments to yaml files. + +```yaml +use_swanlab: true +swanlab_run_name: test_run # optional +``` + +When launching training tasks, you can log in to SwanLab in three ways: + +1. Add `swanlab_api_key=` to the yaml file, and set it to your [API key](https://swanlab.cn/settings). +2. Set the environment variable `SWANLAB_API_KEY` to your [API key](https://swanlab.cn/settings). +3. Use the `swanlab login` command to complete the login. + +## Projects using LLaMA Factory + +If you have a project that should be incorporated, please contact via email or create a pull request. + +
Click to show + +1. Wang et al. ESRL: Efficient Sampling-based Reinforcement Learning for Sequence Generation. 2023. [[arxiv]](https://arxiv.org/abs/2308.02223) +1. Yu et al. Open, Closed, or Small Language Models for Text Classification? 2023. [[arxiv]](https://arxiv.org/abs/2308.10092) +1. Wang et al. UbiPhysio: Support Daily Functioning, Fitness, and Rehabilitation with Action Understanding and Feedback in Natural Language. 2023. [[arxiv]](https://arxiv.org/abs/2308.10526) +1. Luceri et al. Leveraging Large Language Models to Detect Influence Campaigns in Social Media. 2023. [[arxiv]](https://arxiv.org/abs/2311.07816) +1. Zhang et al. Alleviating Hallucinations of Large Language Models through Induced Hallucinations. 2023. [[arxiv]](https://arxiv.org/abs/2312.15710) +1. Wang et al. Know Your Needs Better: Towards Structured Understanding of Marketer Demands with Analogical Reasoning Augmented LLMs. KDD 2024. [[arxiv]](https://arxiv.org/abs/2401.04319) +1. Wang et al. CANDLE: Iterative Conceptualization and Instantiation Distillation from Large Language Models for Commonsense Reasoning. ACL 2024. [[arxiv]](https://arxiv.org/abs/2401.07286) +1. Choi et al. FACT-GPT: Fact-Checking Augmentation via Claim Matching with LLMs. 2024. [[arxiv]](https://arxiv.org/abs/2402.05904) +1. Zhang et al. AutoMathText: Autonomous Data Selection with Language Models for Mathematical Texts. 2024. [[arxiv]](https://arxiv.org/abs/2402.07625) +1. Lyu et al. KnowTuning: Knowledge-aware Fine-tuning for Large Language Models. 2024. [[arxiv]](https://arxiv.org/abs/2402.11176) +1. Yang et al. LaCo: Large Language Model Pruning via Layer Collapse. 2024. [[arxiv]](https://arxiv.org/abs/2402.11187) +1. Bhardwaj et al. Language Models are Homer Simpson! Safety Re-Alignment of Fine-tuned Language Models through Task Arithmetic. 2024. [[arxiv]](https://arxiv.org/abs/2402.11746) +1. Yang et al. Enhancing Empathetic Response Generation by Augmenting LLMs with Small-scale Empathetic Models. 2024. [[arxiv]](https://arxiv.org/abs/2402.11801) +1. Yi et al. Generation Meets Verification: Accelerating Large Language Model Inference with Smart Parallel Auto-Correct Decoding. ACL 2024 Findings. [[arxiv]](https://arxiv.org/abs/2402.11809) +1. Cao et al. Head-wise Shareable Attention for Large Language Models. 2024. [[arxiv]](https://arxiv.org/abs/2402.11819) +1. Zhang et al. Enhancing Multilingual Capabilities of Large Language Models through Self-Distillation from Resource-Rich Languages. 2024. [[arxiv]](https://arxiv.org/abs/2402.12204) +1. Kim et al. Efficient and Effective Vocabulary Expansion Towards Multilingual Large Language Models. 2024. [[arxiv]](https://arxiv.org/abs/2402.14714) +1. Yu et al. KIEval: A Knowledge-grounded Interactive Evaluation Framework for Large Language Models. ACL 2024. [[arxiv]](https://arxiv.org/abs/2402.15043) +1. Huang et al. Key-Point-Driven Data Synthesis with its Enhancement on Mathematical Reasoning. 2024. [[arxiv]](https://arxiv.org/abs/2403.02333) +1. Duan et al. Negating Negatives: Alignment without Human Positive Samples via Distributional Dispreference Optimization. 2024. [[arxiv]](https://arxiv.org/abs/2403.03419) +1. Xie and Schwertfeger. Empowering Robotics with Large Language Models: osmAG Map Comprehension with LLMs. 2024. [[arxiv]](https://arxiv.org/abs/2403.08228) +1. Wu et al. Large Language Models are Parallel Multilingual Learners. 2024. [[arxiv]](https://arxiv.org/abs/2403.09073) +1. Zhang et al. EDT: Improving Large Language Models' Generation by Entropy-based Dynamic Temperature Sampling. 2024. [[arxiv]](https://arxiv.org/abs/2403.14541) +1. Weller et al. FollowIR: Evaluating and Teaching Information Retrieval Models to Follow Instructions. 2024. [[arxiv]](https://arxiv.org/abs/2403.15246) +1. Hongbin Na. CBT-LLM: A Chinese Large Language Model for Cognitive Behavioral Therapy-based Mental Health Question Answering. COLING 2024. [[arxiv]](https://arxiv.org/abs/2403.16008) +1. Zan et al. CodeS: Natural Language to Code Repository via Multi-Layer Sketch. 2024. [[arxiv]](https://arxiv.org/abs/2403.16443) +1. Liu et al. Extensive Self-Contrast Enables Feedback-Free Language Model Alignment. 2024. [[arxiv]](https://arxiv.org/abs/2404.00604) +1. Luo et al. BAdam: A Memory Efficient Full Parameter Training Method for Large Language Models. 2024. [[arxiv]](https://arxiv.org/abs/2404.02827) +1. Du et al. Chinese Tiny LLM: Pretraining a Chinese-Centric Large Language Model. 2024. [[arxiv]](https://arxiv.org/abs/2404.04167) +1. Ma et al. Parameter Efficient Quasi-Orthogonal Fine-Tuning via Givens Rotation. ICML 2024. [[arxiv]](https://arxiv.org/abs/2404.04316) +1. Liu et al. Dynamic Generation of Personalities with Large Language Models. 2024. [[arxiv]](https://arxiv.org/abs/2404.07084) +1. Shang et al. How Far Have We Gone in Stripped Binary Code Understanding Using Large Language Models. 2024. [[arxiv]](https://arxiv.org/abs/2404.09836) +1. Huang et al. LLMTune: Accelerate Database Knob Tuning with Large Language Models. 2024. [[arxiv]](https://arxiv.org/abs/2404.11581) +1. Deng et al. Text-Tuple-Table: Towards Information Integration in Text-to-Table Generation via Global Tuple Extraction. 2024. [[arxiv]](https://arxiv.org/abs/2404.14215) +1. Acikgoz et al. Hippocrates: An Open-Source Framework for Advancing Large Language Models in Healthcare. 2024. [[arxiv]](https://arxiv.org/abs/2404.16621) +1. Zhang et al. Small Language Models Need Strong Verifiers to Self-Correct Reasoning. ACL 2024 Findings. [[arxiv]](https://arxiv.org/abs/2404.17140) +1. Zhou et al. FREB-TQA: A Fine-Grained Robustness Evaluation Benchmark for Table Question Answering. NAACL 2024. [[arxiv]](https://arxiv.org/abs/2404.18585) +1. Xu et al. Large Language Models for Cyber Security: A Systematic Literature Review. 2024. [[arxiv]](https://arxiv.org/abs/2405.04760) +1. Dammu et al. "They are uncultured": Unveiling Covert Harms and Social Threats in LLM Generated Conversations. 2024. [[arxiv]](https://arxiv.org/abs/2405.05378) +1. Yi et al. A safety realignment framework via subspace-oriented model fusion for large language models. 2024. [[arxiv]](https://arxiv.org/abs/2405.09055) +1. Lou et al. SPO: Multi-Dimensional Preference Sequential Alignment With Implicit Reward Modeling. 2024. [[arxiv]](https://arxiv.org/abs/2405.12739) +1. Zhang et al. Getting More from Less: Large Language Models are Good Spontaneous Multilingual Learners. 2024. [[arxiv]](https://arxiv.org/abs/2405.13816) +1. Zhang et al. TS-Align: A Teacher-Student Collaborative Framework for Scalable Iterative Finetuning of Large Language Models. 2024. [[arxiv]](https://arxiv.org/abs/2405.20215) +1. Zihong Chen. Sentence Segmentation and Sentence Punctuation Based on XunziALLM. 2024. [[paper]](https://aclanthology.org/2024.lt4hala-1.30) +1. Gao et al. The Best of Both Worlds: Toward an Honest and Helpful Large Language Model. 2024. [[arxiv]](https://arxiv.org/abs/2406.00380) +1. Wang and Song. MARS: Benchmarking the Metaphysical Reasoning Abilities of Language Models with a Multi-task Evaluation Dataset. 2024. [[arxiv]](https://arxiv.org/abs/2406.02106) +1. Hu et al. Computational Limits of Low-Rank Adaptation (LoRA) for Transformer-Based Models. 2024. [[arxiv]](https://arxiv.org/abs/2406.03136) +1. Ge et al. Time Sensitive Knowledge Editing through Efficient Finetuning. ACL 2024. [[arxiv]](https://arxiv.org/abs/2406.04496) +1. Tan et al. Peer Review as A Multi-Turn and Long-Context Dialogue with Role-Based Interactions. 2024. [[arxiv]](https://arxiv.org/abs/2406.05688) +1. Song et al. Turbo Sparse: Achieving LLM SOTA Performance with Minimal Activated Parameters. 2024. [[arxiv]](https://arxiv.org/abs/2406.05955) +1. Gu et al. RWKV-CLIP: A Robust Vision-Language Representation Learner. 2024. [[arxiv]](https://arxiv.org/abs/2406.06973) +1. Chen et al. Advancing Tool-Augmented Large Language Models: Integrating Insights from Errors in Inference Trees. 2024. [[arxiv]](https://arxiv.org/abs/2406.07115) +1. Zhu et al. Are Large Language Models Good Statisticians?. 2024. [[arxiv]](https://arxiv.org/abs/2406.07815) +1. Li et al. Know the Unknown: An Uncertainty-Sensitive Method for LLM Instruction Tuning. 2024. [[arxiv]](https://arxiv.org/abs/2406.10099) +1. Ding et al. IntentionQA: A Benchmark for Evaluating Purchase Intention Comprehension Abilities of Language Models in E-commerce. 2024. [[arxiv]](https://arxiv.org/abs/2406.10173) +1. He et al. COMMUNITY-CROSS-INSTRUCT: Unsupervised Instruction Generation for Aligning Large Language Models to Online Communities. 2024. [[arxiv]](https://arxiv.org/abs/2406.12074) +1. Lin et al. FVEL: Interactive Formal Verification Environment with Large Language Models via Theorem Proving. 2024. [[arxiv]](https://arxiv.org/abs/2406.14408) +1. Treutlein et al. Connecting the Dots: LLMs can Infer and Verbalize Latent Structure from Disparate Training Data. 2024. [[arxiv]](https://arxiv.org/abs/2406.14546) +1. Feng et al. SS-Bench: A Benchmark for Social Story Generation and Evaluation. 2024. [[arxiv]](https://arxiv.org/abs/2406.15695) +1. Feng et al. Self-Constructed Context Decompilation with Fined-grained Alignment Enhancement. 2024. [[arxiv]](https://arxiv.org/abs/2406.17233) +1. Liu et al. Large Language Models for Cuffless Blood Pressure Measurement From Wearable Biosignals. 2024. [[arxiv]](https://arxiv.org/abs/2406.18069) +1. Iyer et al. Exploring Very Low-Resource Translation with LLMs: The University of Edinburgh's Submission to AmericasNLP 2024 Translation Task. AmericasNLP 2024. [[paper]](https://aclanthology.org/2024.americasnlp-1.25) +1. Li et al. Calibrating LLMs with Preference Optimization on Thought Trees for Generating Rationale in Science Question Scoring. 2024. [[arxiv]](https://arxiv.org/abs/2406.19949) +1. Yang et al. Financial Knowledge Large Language Model. 2024. [[arxiv]](https://arxiv.org/abs/2407.00365) +1. Lin et al. DogeRM: Equipping Reward Models with Domain Knowledge through Model Merging. 2024. [[arxiv]](https://arxiv.org/abs/2407.01470) +1. Bako et al. Evaluating the Semantic Profiling Abilities of LLMs for Natural Language Utterances in Data Visualization. 2024. [[arxiv]](https://arxiv.org/abs/2407.06129) +1. Huang et al. RoLoRA: Fine-tuning Rotated Outlier-free LLMs for Effective Weight-Activation Quantization. 2024. [[arxiv]](https://arxiv.org/abs/2407.08044) +1. Jiang et al. LLM-Collaboration on Automatic Science Journalism for the General Audience. 2024. [[arxiv]](https://arxiv.org/abs/2407.09756) +1. Inouye et al. Applied Auto-tuning on LoRA Hyperparameters. 2024. [[paper]](https://scholarcommons.scu.edu/cseng_senior/272/) +1. Qi et al. Research on Tibetan Tourism Viewpoints information generation system based on LLM. 2024. [[arxiv]](https://arxiv.org/abs/2407.13561) +1. Xu et al. Course-Correction: Safety Alignment Using Synthetic Preferences. 2024. [[arxiv]](https://arxiv.org/abs/2407.16637) +1. Sun et al. LAMBDA: A Large Model Based Data Agent. 2024. [[arxiv]](https://arxiv.org/abs/2407.17535) +1. Zhu et al. CollectiveSFT: Scaling Large Language Models for Chinese Medical Benchmark with Collective Instructions in Healthcare. 2024. [[arxiv]](https://arxiv.org/abs/2407.19705) +1. Yu et al. Correcting Negative Bias in Large Language Models through Negative Attention Score Alignment. 2024. [[arxiv]](https://arxiv.org/abs/2408.00137) +1. Xie et al. The Power of Personalized Datasets: Advancing Chinese Composition Writing for Elementary School through Targeted Model Fine-Tuning. IALP 2024. [[paper]](https://www.asianlp.sg/conferences/ialp2024/proceedings/papers/IALP2024_P055.pdf) +1. Liu et al. Instruct-Code-Llama: Improving Capabilities of Language Model in Competition Level Code Generation by Online Judge Feedback. ICIC 2024. [[paper]](https://link.springer.com/chapter/10.1007/978-981-97-5669-8_11) +1. Wang et al. Cybernetic Sentinels: Unveiling the Impact of Safety Data Selection on Model Security in Supervised Fine-Tuning. ICIC 2024. [[paper]](https://link.springer.com/chapter/10.1007/978-981-97-5669-8_23) +1. Xia et al. Understanding the Performance and Estimating the Cost of LLM Fine-Tuning. 2024. [[arxiv]](https://arxiv.org/abs/2408.04693) +1. Zeng et al. Perceive, Reflect, and Plan: Designing LLM Agent for Goal-Directed City Navigation without Instructions. 2024. [[arxiv]](https://arxiv.org/abs/2408.04168) +1. Xia et al. Using Pre-trained Language Model for Accurate ESG Prediction. FinNLP 2024. [[paper]](https://aclanthology.org/2024.finnlp-2.1/) +1. Liang et al. I-SHEEP: Self-Alignment of LLM from Scratch through an Iterative Self-Enhancement Paradigm. 2024. [[arxiv]](https://arxiv.org/abs/2408.08072) +1. Bai et al. Aligning Large Language Model with Direct Multi-Preference Optimization for Recommendation. CIKM 2024. [[paper]](https://dl.acm.org/doi/10.1145/3627673.3679611) +1. Zhang et al. CPsyCoun: A Report-based Multi-turn Dialogue Reconstruction and Evaluation Framework for Chinese Psychological Counseling. ACL 2024. [[paper]](https://aclanthology.org/2024.findings-acl.830.pdf) +1. **[StarWhisper](https://github.com/Yu-Yang-Li/StarWhisper)**: A large language model for Astronomy, based on ChatGLM2-6B and Qwen-14B. +1. **[DISC-LawLLM](https://github.com/FudanDISC/DISC-LawLLM)**: A large language model specialized in Chinese legal domain, based on Baichuan-13B, is capable of retrieving and reasoning on legal knowledge. +1. **[Sunsimiao](https://github.com/X-D-Lab/Sunsimiao)**: A large language model specialized in Chinese medical domain, based on Baichuan-7B and ChatGLM-6B. +1. **[CareGPT](https://github.com/WangRongsheng/CareGPT)**: A series of large language models for Chinese medical domain, based on LLaMA2-7B and Baichuan-13B. +1. **[MachineMindset](https://github.com/PKU-YuanGroup/Machine-Mindset/)**: A series of MBTI Personality large language models, capable of giving any LLM 16 different personality types based on different datasets and training methods. +1. **[Luminia-13B-v3](https://huggingface.co/Nekochu/Luminia-13B-v3)**: A large language model specialized in generate metadata for stable diffusion. [[demo]](https://huggingface.co/spaces/Nekochu/Luminia-13B_SD_Prompt) +1. **[Chinese-LLaVA-Med](https://github.com/BUAADreamer/Chinese-LLaVA-Med)**: A multimodal large language model specialized in Chinese medical domain, based on LLaVA-1.5-7B. +1. **[AutoRE](https://github.com/THUDM/AutoRE)**: A document-level relation extraction system based on large language models. +1. **[NVIDIA RTX AI Toolkit](https://github.com/NVIDIA/RTX-AI-Toolkit)**: SDKs for fine-tuning LLMs on Windows PC for NVIDIA RTX. +1. **[LazyLLM](https://github.com/LazyAGI/LazyLLM)**: An easy and lazy way for building multi-agent LLMs applications and supports model fine-tuning via LLaMA Factory. +1. **[RAG-Retrieval](https://github.com/NLPJCL/RAG-Retrieval)**: A full pipeline for RAG retrieval model fine-tuning, inference, and distillation. [[blog]](https://zhuanlan.zhihu.com/p/987727357) +1. **[360-LLaMA-Factory](https://github.com/Qihoo360/360-LLaMA-Factory)**: A modified library that supports long sequence SFT & DPO using ring attention. +1. **[Sky-T1](https://novasky-ai.github.io/posts/sky-t1/)**: An o1-like model fine-tuned by NovaSky AI with very small cost. +1. **[WeClone](https://github.com/xming521/WeClone)**: One-stop solution for creating your digital avatar from chat logs. +1. **[EmoLLM](https://github.com/SmartFlowAI/EmoLLM)**: A project about large language models (LLMs) and mental health. +
+ +## License + +This repository is licensed under the [Apache-2.0 License](LICENSE). + +Please follow the model licenses to use the corresponding model weights: [BLOOM](https://huggingface.co/spaces/bigscience/license) / [DeepSeek](https://github.com/deepseek-ai/DeepSeek-LLM/blob/main/LICENSE-MODEL) / [Falcon](https://huggingface.co/tiiuae/falcon-180B/blob/main/LICENSE.txt) / [Gemma](https://ai.google.dev/gemma/terms) / [GLM-4](https://huggingface.co/THUDM/glm-4-9b/blob/main/LICENSE) / [GPT-2](https://github.com/openai/gpt-2/blob/master/LICENSE) / [Granite](LICENSE) / [InternLM](https://github.com/InternLM/InternLM#license) / [Llama](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) / [Llama 2](https://ai.meta.com/llama/license/) / [Llama 3](https://llama.meta.com/llama3/license/) / [Llama 4](https://github.com/meta-llama/llama-models/blob/main/models/llama4/LICENSE) / [MiniCPM](https://github.com/OpenBMB/MiniCPM/blob/main/MiniCPM%20Model%20License.md) / [Mistral/Mixtral/Pixtral](LICENSE) / [Phi-3/Phi-4](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/LICENSE) / [Qwen](https://github.com/QwenLM/Qwen/blob/main/Tongyi%20Qianwen%20LICENSE%20AGREEMENT) / [StarCoder 2](https://huggingface.co/spaces/bigcode/bigcode-model-license-agreement) / [TeleChat2](https://huggingface.co/Tele-AI/telechat-7B/blob/main/TeleChat%E6%A8%A1%E5%9E%8B%E7%A4%BE%E5%8C%BA%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf) / [Yuan 2](https://github.com/IEIT-Yuan/Yuan-2.0/blob/main/LICENSE-Yuan) + +## Citation + +If this work is helpful, please kindly cite as: + +```bibtex +@inproceedings{zheng2024llamafactory, + title={LlamaFactory: Unified Efficient Fine-Tuning of 100+ Language Models}, + author={Yaowei Zheng and Richong Zhang and Junhao Zhang and Yanhan Ye and Zheyan Luo and Zhangchi Feng and Yongqiang Ma}, + booktitle={Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 3: System Demonstrations)}, + address={Bangkok, Thailand}, + publisher={Association for Computational Linguistics}, + year={2024}, + url={http://arxiv.org/abs/2403.13372} +} +``` + +## Acknowledgement + +This repo benefits from [PEFT](https://github.com/huggingface/peft), [TRL](https://github.com/huggingface/trl), [QLoRA](https://github.com/artidoro/qlora) and [FastChat](https://github.com/lm-sys/FastChat). Thanks for their wonderful works. + +## Star History + +![Star History Chart](https://api.star-history.com/svg?repos=hiyouga/LLaMA-Factory&type=Date) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..6842d5aeed9456b93b1667195eba44ff9e0c5e24 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,141 @@ +accelerate==1.11.0 +aiofiles==24.1.0 +aiohappyeyeballs==2.7.1 +aiohttp==3.14.1 +aiosignal==1.4.0 +annotated-doc==0.0.4 +annotated-types==0.7.0 +antlr4-python3-runtime==4.9.3 +anyio==4.14.1 +attrs==26.1.0 +av==16.0.0 +bitsandbytes==0.49.2 +brotli==1.2.0 +certifi==2026.6.17 +charset-normalizer==3.4.9 +click==8.4.2 +contourpy==1.3.3 +cuda-bindings==13.3.1 +cuda-pathfinder==1.5.6 +cuda-toolkit==13.0.3.0 +cycler==0.12.1 +datasets==4.0.0 +dill==0.3.8 +docstring_parser==0.18.0 +einops==0.8.2 +fastapi==0.139.0 +ffmpy==1.0.0 +filelock==3.29.7 +fire==0.7.1 +fonttools==4.63.0 +frozenlist==1.8.0 +fsspec==2025.3.0 +gradio==5.50.0 +gradio_client==1.14.0 +groovy==0.1.2 +h11==0.16.0 +hf-xet==1.5.1 +hf_transfer==0.1.9 +httpcore==1.0.9 +httpx==0.28.1 +huggingface_hub==1.23.0 +idna==3.18 +Jinja2==3.1.6 +kiwisolver==1.5.0 +-e git+https://github.com/hiyouga/LLaMA-Factory.git@ea31c43d806162a7fd98065abfef2d974fff5766#egg=llamafactory +markdown-it-py==4.2.0 +MarkupSafe==3.0.3 +matplotlib==3.11.0 +mdurl==0.1.2 +modelscope==1.38.1 +modelscope-hub==0.1.7 +mpmath==1.3.0 +multidict==6.7.1 +multiprocess==0.70.16 +networkx==3.6.1 +numpy==2.5.1 +nvidia-cublas==13.1.1.3 +nvidia-cublas-cu12==12.4.5.8 +nvidia-cuda-cupti==13.0.85 +nvidia-cuda-cupti-cu12==12.4.127 +nvidia-cuda-nvrtc==13.0.88 +nvidia-cuda-nvrtc-cu12==12.4.127 +nvidia-cuda-runtime==13.0.96 +nvidia-cuda-runtime-cu12==12.4.127 +nvidia-cudnn-cu12==9.1.0.70 +nvidia-cudnn-cu13==9.20.0.48 +nvidia-cufft==12.0.0.61 +nvidia-cufft-cu12==11.2.1.3 +nvidia-cufile==1.15.1.6 +nvidia-curand==10.4.0.35 +nvidia-curand-cu12==10.3.5.147 +nvidia-cusolver==12.0.4.66 +nvidia-cusolver-cu12==11.6.1.9 +nvidia-cusparse==12.6.3.3 +nvidia-cusparse-cu12==12.3.1.170 +nvidia-cusparselt-cu12==0.6.2 +nvidia-cusparselt-cu13==0.8.1 +nvidia-nccl-cu12==2.21.5 +nvidia-nccl-cu13==2.29.7 +nvidia-nvjitlink==13.3.33 +nvidia-nvjitlink-cu12==12.4.127 +nvidia-nvshmem-cu13==3.4.5 +nvidia-nvtx==13.0.85 +nvidia-nvtx-cu12==12.4.127 +omegaconf==2.3.1 +orjson==3.11.9 +packaging==26.2 +pandas==2.3.3 +peft==0.18.1 +pillow==11.3.0 +propcache==0.5.2 +protobuf==7.35.1 +psutil==7.2.2 +pyarrow==25.0.0 +pydantic==2.12.3 +pydantic_core==2.41.4 +pydub==0.25.1 +Pygments==2.20.0 +pyparsing==3.3.2 +python-dateutil==2.9.0.post0 +python-multipart==0.0.32 +pytz==2026.2 +PyYAML==6.0.3 +regex==2026.6.28 +requests==2.34.2 +rich==15.0.0 +ruff==0.15.21 +safehttpx==0.1.7 +safetensors==0.8.0 +scipy==1.18.0 +semantic-version==2.10.0 +sentencepiece==0.2.1 +setuptools==83.0.0 +shellingham==1.5.4 +shtab==1.8.1 +six==1.17.0 +sse-starlette==3.4.5 +starlette==0.52.1 +sympy==1.13.1 +termcolor==3.3.0 +tiktoken==0.13.0 +tokenizers==0.22.2 +tomlkit==0.13.3 +torch==2.6.0+cu124 +torchaudio==2.6.0+cu124 +torchdata==0.11.0 +torchvision==0.21.0+cu124 +tqdm==4.68.4 +transformers==5.7.0 +triton==3.2.0 +trl==0.24.0 +typer==0.26.8 +typing-inspection==0.4.2 +typing_extensions==4.16.0 +tyro==0.8.14 +tzdata==2026.3 +urllib3==2.7.0 +uvicorn==0.51.0 +websockets==15.0.1 +xxhash==3.8.1 +yarl==1.24.2 diff --git a/scripts/api_example/test_image.py b/scripts/api_example/test_image.py new file mode 100644 index 0000000000000000000000000000000000000000..afd2b69c4ec951bcc6b08b4d5e50f11048f7f7d8 --- /dev/null +++ b/scripts/api_example/test_image.py @@ -0,0 +1,65 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 openai import OpenAI +from transformers.utils.versions import require_version + + +require_version("openai>=1.5.0", "To fix: pip install openai>=1.5.0") + + +def main(): + client = OpenAI( + api_key="{}".format(os.getenv("API_KEY", "0")), + base_url="http://localhost:{}/v1".format(os.getenv("API_PORT", 8000)), + ) + messages = [] + messages.append( + { + "role": "user", + "content": [ + {"type": "text", "text": "Output the color and number of each box."}, + { + "type": "image_url", + "image_url": {"url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-VL/boxes.png"}, + }, + ], + } + ) + result = client.chat.completions.create(messages=messages, model="test") + messages.append(result.choices[0].message) + print("Round 1:", result.choices[0].message.content) + # The image shows a pyramid of colored blocks with numbers on them. Here are the colors and numbers of ... + messages.append( + { + "role": "user", + "content": [ + {"type": "text", "text": "What kind of flower is this?"}, + { + "type": "image_url", + "image_url": {"url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-VL/flowers.jpg"}, + }, + ], + } + ) + result = client.chat.completions.create(messages=messages, model="test") + messages.append(result.choices[0].message) + print("Round 2:", result.choices[0].message.content) + # The image shows a cluster of forget-me-not flowers. Forget-me-nots are small ... + + +if __name__ == "__main__": + main() diff --git a/scripts/api_example/test_toolcall.py b/scripts/api_example/test_toolcall.py new file mode 100644 index 0000000000000000000000000000000000000000..e291ba693df025673d14198fd07fbbe5e8552421 --- /dev/null +++ b/scripts/api_example/test_toolcall.py @@ -0,0 +1,77 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 openai import OpenAI +from transformers.utils.versions import require_version + + +require_version("openai>=1.5.0", "To fix: pip install openai>=1.5.0") + + +def calculate_gpa(grades: list[str], hours: list[int]) -> float: + grade_to_score = {"A": 4, "B": 3, "C": 2} + total_score, total_hour = 0, 0 + for grade, hour in zip(grades, hours): + total_score += grade_to_score[grade] * hour + total_hour += hour + return round(total_score / total_hour, 2) + + +def main(): + client = OpenAI( + api_key="{}".format(os.getenv("API_KEY", "0")), + base_url="http://localhost:{}/v1".format(os.getenv("API_PORT", 8000)), + ) + tools = [ + { + "type": "function", + "function": { + "name": "calculate_gpa", + "description": "Calculate the Grade Point Average (GPA) based on grades and credit hours", + "parameters": { + "type": "object", + "properties": { + "grades": {"type": "array", "items": {"type": "string"}, "description": "The grades"}, + "hours": {"type": "array", "items": {"type": "integer"}, "description": "The credit hours"}, + }, + "required": ["grades", "hours"], + }, + }, + } + ] + tool_map = {"calculate_gpa": calculate_gpa} + + messages = [] + messages.append({"role": "user", "content": "My grades are A, A, B, and C. The credit hours are 3, 4, 3, and 2."}) + result = client.chat.completions.create(messages=messages, model="test", tools=tools) + if result.choices[0].message.tool_calls is None: + raise ValueError("Cannot retrieve function call from the response.") + + messages.append(result.choices[0].message) + tool_call = result.choices[0].message.tool_calls[0].function + print(tool_call) + # Function(arguments='{"grades": ["A", "A", "B", "C"], "hours": [3, 4, 3, 2]}', name='calculate_gpa') + name, arguments = tool_call.name, json.loads(tool_call.arguments) + tool_result = tool_map[name](**arguments) + messages.append({"role": "tool", "content": json.dumps({"gpa": tool_result}, ensure_ascii=False)}) + result = client.chat.completions.create(messages=messages, model="test", tools=tools) + print(result.choices[0].message.content) + # Based on the grades and credit hours you provided, your Grade Point Average (GPA) is 3.42. + + +if __name__ == "__main__": + main() diff --git a/scripts/bench_qwen.py b/scripts/bench_qwen.py new file mode 100644 index 0000000000000000000000000000000000000000..4ab48dda2c42b4e3ea291ee539c2e4afe7b51c99 --- /dev/null +++ b/scripts/bench_qwen.py @@ -0,0 +1,173 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 dataclasses import dataclass +from typing import Any + +import fire +import torch +from peft import PeftModel +from torch.utils.data import Dataset +from transformers import DataCollatorForSeq2Seq, Qwen2_5_VLProcessor + +from llamafactory.extras.constants import IGNORE_INDEX +from llamafactory.hparams import get_train_args +from llamafactory.model import load_model, load_tokenizer +from llamafactory.train.callbacks import LogCallback +from llamafactory.train.sft.trainer import CustomSeq2SeqTrainer + + +class DummyDataset(Dataset): + def __init__(self, size: int = 1000, seq_length: int = 1024, processor: Qwen2_5_VLProcessor = None): + self.size = size + self.seq_length = seq_length + self.vocab_size = 32768 + self.processor = processor + + image_token_num = 18 * 18 // (2 * 2) + image_t = 2 + + self.text_seqlen = seq_length // 4 # 25% text + video_seq_length = self.seq_length - self.text_seqlen - image_t * image_token_num + video_t = video_seq_length // image_token_num + + self.image_size = [18 * 18 * image_t, 1176] + self.image_grid_thw = torch.tensor([[1, 18, 18]] * image_t, dtype=torch.long) + self.image_seqlen = image_t * image_token_num + + self.video_size = [18 * 18 * video_t, 1176] + self.video_grid_thw = torch.tensor([[video_t, 18, 18]], dtype=torch.long) + self.video_seqlen = video_t * image_token_num + + def __len__(self): + return self.size + + def __getitem__(self, index: int): + input_ids = torch.randint(low=0, high=self.vocab_size, size=(self.seq_length,)) + input_ids[: self.image_seqlen] = self.processor.image_token_id + input_ids[self.image_seqlen : self.image_seqlen + self.video_seqlen] = self.processor.video_token_id + + attention_mask = torch.ones((self.seq_length,), dtype=torch.long) + labels = input_ids.clone() + labels[: self.image_seqlen + self.video_seqlen] = IGNORE_INDEX + pixel_values = torch.rand(self.image_size, dtype=torch.float32) + pixel_values_videos = torch.rand(self.video_size, dtype=torch.float32) + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + "pixel_values": pixel_values, + "pixel_values_videos": pixel_values_videos, + "image_grid_thw": self.image_grid_thw, + "video_grid_thw": self.video_grid_thw, + } + + +@dataclass +class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq): + def __post_init__(self): + if isinstance(self.model, PeftModel): + self.model = self.model.base_model.model + + if self.model is not None and hasattr(self.model, "get_rope_index"): # for qwen2vl mrope + self.get_rope_func = self.model.get_rope_index # transformers < 4.52.0 or qwen2.5 omni + elif self.model is not None and hasattr(self.model, "model") and hasattr(self.model.model, "get_rope_index"): + self.get_rope_func = self.model.model.get_rope_index # transformers >= 4.52.0 + else: + self.get_rope_func = None + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]: + batch_pixel_values = [feature.pop("pixel_values") for feature in features] + batch_pixel_values_videos = [feature.pop("pixel_values_videos") for feature in features] + batch_image_grid_thw = [feature.pop("image_grid_thw") for feature in features] + batch_video_grid_thw = [feature.pop("video_grid_thw") for feature in features] + + batch: dict[str, torch.Tensor] = super().__call__(features) + + batch["pixel_values"] = torch.cat(batch_pixel_values, dim=0) + batch["pixel_values_videos"] = torch.cat(batch_pixel_values_videos, dim=0) + batch["image_grid_thw"] = torch.cat(batch_image_grid_thw, dim=0) + batch["video_grid_thw"] = torch.cat(batch_video_grid_thw, dim=0) + + if self.get_rope_func is not None: + rope_index_kwargs = { + "input_ids": batch["input_ids"], + "image_grid_thw": batch["image_grid_thw"], + "video_grid_thw": batch["video_grid_thw"], + "attention_mask": (batch["attention_mask"] >= 1).float(), + } + batch["position_ids"], batch["rope_deltas"] = self.get_rope_func(**rope_index_kwargs) + + if "position_ids" not in batch or batch["position_ids"].dim() != 3: + raise ValueError("Qwen2VL requires 3D position ids for mrope.") + + return batch + + +def bench_qwen( + model_name_or_path: str = "Qwen/Qwen2-VL-7B-Instruct", + batch_size: int = 1, + seq_length: int = 2048, + liger_kernel: bool = False, + deepspeed_stage: int = 3, +): + os.environ["LLAMABOARD_ENABLED"] = "true" + os.environ["LLAMABOARD_WORKDIR"] = "output/dummy_dir" + args = { + "model_name_or_path": model_name_or_path, + "enable_liger_kernel": liger_kernel, + "stage": "sft", + "do_train": True, + "finetuning_type": "full", + "dataset": "alpaca_en_demo", + "template": "qwen2_vl", + "cutoff_len": seq_length, + "output_dir": "output/dummy_dir", + "logging_steps": 10, + "save_strategy": "no", + "save_only_model": True, + "overwrite_output_dir": True, + "per_device_train_batch_size": batch_size, + "max_steps": 1000, + "bf16": True, + "include_num_input_tokens_seen": True, + "report_to": "none", + } + if deepspeed_stage in [2, 3]: + args["deepspeed"] = f"examples/deepspeed/ds_z{deepspeed_stage}_config.json" + + model_args, _, training_args, finetuning_args, _ = get_train_args(args) + tokenizer_module = load_tokenizer(model_args) + tokenizer = tokenizer_module["tokenizer"] + trainset = DummyDataset(size=100000, seq_length=seq_length, processor=tokenizer_module["processor"]) + model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train) + data_collator = MultiModalDataCollatorForSeq2Seq( + tokenizer=tokenizer, model=model, pad_to_multiple_of=8, label_pad_token_id=IGNORE_INDEX + ) + + trainer = CustomSeq2SeqTrainer( + model=model, + args=training_args, + finetuning_args=finetuning_args, + data_collator=data_collator, + callbacks=[LogCallback()], + train_dataset=trainset, + **tokenizer_module, + ) + trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint) + + +if __name__ == "__main__": + fire.Fire(bench_qwen) diff --git a/scripts/convert_ckpt/llamafy_baichuan2.py b/scripts/convert_ckpt/llamafy_baichuan2.py new file mode 100644 index 0000000000000000000000000000000000000000..62dc6a51ac38e5fa9dc06a2136bb9dd2154af055 --- /dev/null +++ b/scripts/convert_ckpt/llamafy_baichuan2.py @@ -0,0 +1,112 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 collections import OrderedDict +from typing import Any + +import fire +import torch +from huggingface_hub import split_torch_state_dict_into_shards +from safetensors.torch import save_file +from tqdm import tqdm +from transformers.modeling_utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME + + +CONFIG_NAME = "config.json" + + +def save_weight(input_dir: str, output_dir: str, shard_size: str, save_safetensors: bool): + baichuan2_state_dict: dict[str, torch.Tensor] = OrderedDict() + for filepath in tqdm(os.listdir(input_dir), desc="Load weights"): + if os.path.isfile(os.path.join(input_dir, filepath)) and filepath.endswith(".bin"): + shard_weight = torch.load(os.path.join(input_dir, filepath), map_location="cpu", weights_only=True) + baichuan2_state_dict.update(shard_weight) + + llama_state_dict: dict[str, torch.Tensor] = OrderedDict() + for key, value in tqdm(baichuan2_state_dict.items(), desc="Convert format"): + if "W_pack" in key: + proj_size = value.size(0) // 3 + llama_state_dict[key.replace("W_pack", "q_proj")] = value[:proj_size, :] + llama_state_dict[key.replace("W_pack", "k_proj")] = value[proj_size : 2 * proj_size, :] + llama_state_dict[key.replace("W_pack", "v_proj")] = value[2 * proj_size :, :] + elif "lm_head" in key: + llama_state_dict[key] = torch.nn.functional.normalize(value) + else: + llama_state_dict[key] = value + + weights_name = SAFE_WEIGHTS_NAME if save_safetensors else WEIGHTS_NAME + filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors") + state_dict_split = split_torch_state_dict_into_shards( + llama_state_dict, filename_pattern=filename_pattern, max_shard_size=shard_size + ) + for shard_file, tensors in tqdm(state_dict_split.filename_to_tensors.items(), desc="Save weights"): + shard = {tensor: llama_state_dict[tensor].contiguous() for tensor in tensors} + if save_safetensors: + save_file(shard, os.path.join(output_dir, shard_file), metadata={"format": "pt"}) + else: + torch.save(shard, os.path.join(output_dir, shard_file)) + + if not state_dict_split.is_sharded: + print(f"Model weights saved in {os.path.join(output_dir, weights_name)}.") + else: + index = { + "metadata": state_dict_split.metadata, + "weight_map": state_dict_split.tensor_to_filename, + } + index_name = SAFE_WEIGHTS_INDEX_NAME if save_safetensors else WEIGHTS_INDEX_NAME + with open(os.path.join(output_dir, index_name), "w", encoding="utf-8") as f: + json.dump(index, f, indent=2, sort_keys=True) + + print(f"Model weights saved in {output_dir}.") + + +def save_config(input_dir: str, output_dir: str): + with open(os.path.join(input_dir, CONFIG_NAME), encoding="utf-8") as f: + llama2_config_dict: dict[str, Any] = json.load(f) + + llama2_config_dict["architectures"] = ["LlamaForCausalLM"] + llama2_config_dict.pop("auto_map", None) + llama2_config_dict.pop("tokenizer_class", None) + llama2_config_dict["model_type"] = "llama" + + with open(os.path.join(output_dir, CONFIG_NAME), "w", encoding="utf-8") as f: + json.dump(llama2_config_dict, f, indent=2) + + print(f"Model config saved in {os.path.join(output_dir, CONFIG_NAME)}") + + +def llamafy_baichuan2( + input_dir: str, + output_dir: str, + shard_size: str = "2GB", + save_safetensors: bool = True, +): + r"""Convert the Baichuan2-7B model in the same format as LLaMA2-7B. + + Usage: python llamafy_baichuan2.py --input_dir input --output_dir output + Converted model: https://huggingface.co/hiyouga/Baichuan2-7B-Base-LLaMAfied + """ + try: + os.makedirs(output_dir, exist_ok=False) + except Exception as e: + raise print("Output dir already exists", e) + + save_weight(input_dir, output_dir, shard_size, save_safetensors) + save_config(input_dir, output_dir) + + +if __name__ == "__main__": + fire.Fire(llamafy_baichuan2) diff --git a/scripts/convert_ckpt/llamafy_qwen.py b/scripts/convert_ckpt/llamafy_qwen.py new file mode 100644 index 0000000000000000000000000000000000000000..599b0f1226cde2e3de6f009c056b8b89c84838b4 --- /dev/null +++ b/scripts/convert_ckpt/llamafy_qwen.py @@ -0,0 +1,165 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 collections import OrderedDict +from typing import Any + +import fire +import torch +from huggingface_hub import split_torch_state_dict_into_shards +from safetensors import safe_open +from safetensors.torch import save_file +from tqdm import tqdm +from transformers.modeling_utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME +from transformers.utils import check_min_version + + +try: + check_min_version("4.34.0") +except Exception: + raise ValueError("Please upgrade `transformers` to 4.34.0") + + +CONFIG_NAME = "config.json" + + +def save_weight(input_dir: str, output_dir: str, shard_size: str, save_safetensors: bool) -> str: + qwen_state_dict: dict[str, torch.Tensor] = OrderedDict() + for filepath in tqdm(os.listdir(input_dir), desc="Load weights"): + if os.path.isfile(os.path.join(input_dir, filepath)) and filepath.endswith(".safetensors"): + with safe_open(os.path.join(input_dir, filepath), framework="pt", device="cpu") as f: + for key in f.keys(): + qwen_state_dict[key] = f.get_tensor(key) + + llama_state_dict: dict[str, torch.Tensor] = OrderedDict() + torch_dtype = None + for key, value in tqdm(qwen_state_dict.items(), desc="Convert format"): + if torch_dtype is None: + torch_dtype = value.dtype + if "wte" in key: + llama_state_dict["model.embed_tokens.weight"] = value + elif "ln_f" in key: + llama_state_dict["model.norm.weight"] = value + else: + key = key.replace("transformer.h", "model.layers") + if "attn.c_attn" in key: + proj_size = value.size(0) // 3 + llama_state_dict[key.replace("attn.c_attn", "self_attn.q_proj")] = value[:proj_size, ...] + llama_state_dict[key.replace("attn.c_attn", "self_attn.k_proj")] = value[ + proj_size : 2 * proj_size, ... + ] + llama_state_dict[key.replace("attn.c_attn", "self_attn.v_proj")] = value[2 * proj_size :, ...] + elif "attn.c_proj" in key: + llama_state_dict[key.replace("attn.c_proj", "self_attn.o_proj")] = value + llama_state_dict[key.replace("attn.c_proj.weight", "self_attn.o_proj.bias")] = torch.zeros_like( + value[:, 0] + ).squeeze() + elif "ln_1" in key: + llama_state_dict[key.replace("ln_1", "input_layernorm")] = value + elif "ln_2" in key: + llama_state_dict[key.replace("ln_2", "post_attention_layernorm")] = value + elif "mlp.w1" in key: + llama_state_dict[key.replace("mlp.w1", "mlp.up_proj")] = value + elif "mlp.w2" in key: + llama_state_dict[key.replace("mlp.w2", "mlp.gate_proj")] = value + elif "mlp.c_proj" in key: + llama_state_dict[key.replace("mlp.c_proj", "mlp.down_proj")] = value + elif "lm_head" in key: + llama_state_dict[key] = value + else: + raise KeyError(f"Unable to process key {key}") + + weights_name = SAFE_WEIGHTS_NAME if save_safetensors else WEIGHTS_NAME + filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors") + state_dict_split = split_torch_state_dict_into_shards( + llama_state_dict, filename_pattern=filename_pattern, max_shard_size=shard_size + ) + for shard_file, tensors in tqdm(state_dict_split.filename_to_tensors.items(), desc="Save weights"): + shard = {tensor: llama_state_dict[tensor].contiguous() for tensor in tensors} + if save_safetensors: + save_file(shard, os.path.join(output_dir, shard_file), metadata={"format": "pt"}) + else: + torch.save(shard, os.path.join(output_dir, shard_file)) + + if not state_dict_split.is_sharded: + print(f"Model weights saved in {os.path.join(output_dir, weights_name)}.") + else: + index = { + "metadata": state_dict_split.metadata, + "weight_map": state_dict_split.tensor_to_filename, + } + index_name = SAFE_WEIGHTS_INDEX_NAME if save_safetensors else WEIGHTS_INDEX_NAME + with open(os.path.join(output_dir, index_name), "w", encoding="utf-8") as f: + json.dump(index, f, indent=2, sort_keys=True) + + print(f"Model weights saved in {output_dir}.") + + return str(torch_dtype).replace("torch.", "") + + +def save_config(input_dir: str, output_dir: str, torch_dtype: str): + with open(os.path.join(input_dir, CONFIG_NAME), encoding="utf-8") as f: + qwen_config_dict: dict[str, Any] = json.load(f) + + llama2_config_dict: dict[str, Any] = OrderedDict() + llama2_config_dict["architectures"] = ["LlamaForCausalLM"] + llama2_config_dict["hidden_act"] = "silu" + llama2_config_dict["hidden_size"] = qwen_config_dict["hidden_size"] + llama2_config_dict["initializer_range"] = qwen_config_dict["initializer_range"] + llama2_config_dict["intermediate_size"] = qwen_config_dict["intermediate_size"] // 2 + llama2_config_dict["max_position_embeddings"] = qwen_config_dict["max_position_embeddings"] + llama2_config_dict["model_type"] = "llama" + llama2_config_dict["num_attention_heads"] = qwen_config_dict["num_attention_heads"] + llama2_config_dict["num_hidden_layers"] = qwen_config_dict["num_hidden_layers"] + llama2_config_dict["num_key_value_heads"] = qwen_config_dict["hidden_size"] // qwen_config_dict["kv_channels"] + llama2_config_dict["pretraining_tp"] = 1 + llama2_config_dict["rms_norm_eps"] = qwen_config_dict["layer_norm_epsilon"] + llama2_config_dict["rope_scaling"] = None + llama2_config_dict["tie_word_embeddings"] = qwen_config_dict["tie_word_embeddings"] + llama2_config_dict["torch_dtype"] = torch_dtype + llama2_config_dict["transformers_version"] = "4.34.0" + llama2_config_dict["use_cache"] = True + llama2_config_dict["vocab_size"] = qwen_config_dict["vocab_size"] + llama2_config_dict["attention_bias"] = True + + with open(os.path.join(output_dir, CONFIG_NAME), "w", encoding="utf-8") as f: + json.dump(llama2_config_dict, f, indent=2) + + print(f"Model config saved in {os.path.join(output_dir, CONFIG_NAME)}") + + +def llamafy_qwen( + input_dir: str, + output_dir: str, + shard_size: str = "2GB", + save_safetensors: bool = False, +): + r"""Convert the Qwen models in the same format as LLaMA2. + + Usage: python llamafy_qwen.py --input_dir input --output_dir output + Converted model: https://huggingface.co/hiyouga/Qwen-14B-Chat-LLaMAfied + """ + try: + os.makedirs(output_dir, exist_ok=False) + except Exception as e: + raise print("Output dir already exists", e) + + torch_dtype = save_weight(input_dir, output_dir, shard_size, save_safetensors) + save_config(input_dir, output_dir, torch_dtype) + + +if __name__ == "__main__": + fire.Fire(llamafy_qwen) diff --git a/scripts/convert_ckpt/tiny_llama4.py b/scripts/convert_ckpt/tiny_llama4.py new file mode 100644 index 0000000000000000000000000000000000000000..2a96cfa60f30660e572baa436c7f7b67e2c907f1 --- /dev/null +++ b/scripts/convert_ckpt/tiny_llama4.py @@ -0,0 +1,39 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 transformers import Llama4Config, Llama4ForConditionalGeneration, Llama4TextConfig, Llama4VisionConfig + + +if __name__ == "__main__": + vision_config = Llama4VisionConfig( + hidden_size=1408, + image_size=336, + intermediate_size=5632, + num_attention_heads=16, + num_hidden_layers=4, + vision_output_dim=4096, + ) + text_config = Llama4TextConfig( + hidden_size=512, + intermediate_size=1024, + intermediate_size_mlp=1024, + num_hidden_layers=4, + num_attention_heads=8, + num_key_value_heads=2, + head_dim=512 // 8, + num_local_experts=2, + ) + config = Llama4Config(vision_config=vision_config, text_config=text_config) + model = Llama4ForConditionalGeneration._from_config(config) + model.save_pretrained("tiny-llama4") diff --git a/scripts/convert_ckpt/tiny_qwen3.py b/scripts/convert_ckpt/tiny_qwen3.py new file mode 100644 index 0000000000000000000000000000000000000000..902c0ec09b10d456f8936cbdef60fffac1678720 --- /dev/null +++ b/scripts/convert_ckpt/tiny_qwen3.py @@ -0,0 +1,32 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM + + +if __name__ == "__main__": + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Instruct-2507") + config = Qwen3Config( + hidden_size=1408, + image_size=336, + intermediate_size=5632, + num_attention_heads=16, + num_hidden_layers=4, + vision_output_dim=4096, + ) + model = Qwen3ForCausalLM.from_config(config) + model.save_pretrained("tiny-qwen3") + tokenizer.save_pretrained("tiny-qwen3") + model.push_to_hub("llamafactory/tiny-random-qwen3") + tokenizer.push_to_hub("llamafactory/tiny-random-qwen3") diff --git a/scripts/dcp2hf.py b/scripts/dcp2hf.py new file mode 100644 index 0000000000000000000000000000000000000000..8e3256bbb7c2fb9c47ffcbf76b85822f39b28c2f --- /dev/null +++ b/scripts/dcp2hf.py @@ -0,0 +1,76 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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. + +"""Convert a DCP checkpoint to HuggingFace model format. + +Usage: + python scripts/dcp2hf.py convert --dcp_path=/path/to/dcp --hf_path=/path/to/hf --config_path=/path/to/config + +Arguments: + dcp_path: Path to the DCP checkpoint directory. + hf_path: Output path (directory) for HuggingFace model. + config_path: Path to the HuggingFace model directory containing config.json. +""" + +import fire +import torch +import torch.distributed.checkpoint as dcp +import transformers +from transformers import AutoConfig + + +def convert(dcp_path: str, hf_path: str, config_path: str) -> None: + """Convert DCP model weights to HF. + + Note: this script is used to convert a DCP checkpoint to HuggingFace model format, + it will just convert the DCP checkpoint to a HuggingFace model format, for the tokenizer, + you may need to copy from the original model. + + Args: + dcp_path: DCP checkpoint directory. + hf_path: Output path (directory) for HuggingFace model. + config_path: Path to the HuggingFace model directory containing config.json. + """ + if not dcp_path or not hf_path or not config_path: + raise ValueError("All 'dcp_path', 'hf_path', and 'config_path' are required.") + + print(f"Loading config from {config_path}...") + config = AutoConfig.from_pretrained(config_path) + architectures = getattr(config, "architectures", []) + if architectures: + model_cls = getattr(transformers, architectures[0], transformers.AutoModelForCausalLM) + else: + model_cls = transformers.AutoModelForCausalLM + + print("Initializing model on CPU...") + model = model_cls(config).to(torch.bfloat16) + + print(f"Loading DCP from {dcp_path}...") + state_dict = model.state_dict() + dcp.load(state_dict, checkpoint_id=dcp_path) + model.load_state_dict(state_dict) + + print(f"Saving to HF format at {hf_path}...") + model.save_pretrained(hf_path) + config.save_pretrained(hf_path) + print("Done!") + + +def help() -> None: + """Show help message.""" + print(__doc__) + + +if __name__ == "__main__": + fire.Fire({"convert": convert, "help": help, "--convert": convert}) diff --git a/scripts/eval_bleu_rouge.py b/scripts/eval_bleu_rouge.py new file mode 100644 index 0000000000000000000000000000000000000000..4ff96dd8ecbdd15946020dbcd15625ba801d43c4 --- /dev/null +++ b/scripts/eval_bleu_rouge.py @@ -0,0 +1,79 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 logging +import time + +import fire +from datasets import load_dataset + + +try: + import jieba # type: ignore + from nltk.translate.bleu_score import SmoothingFunction, sentence_bleu # type: ignore + from rouge_chinese import Rouge # type: ignore + + jieba.setLogLevel(logging.CRITICAL) + jieba.initialize() +except ImportError: + print("Please install llamafactory with `pip install -r requirements/metrics.txt`.") + raise + + +def compute_metrics(sample): + hypothesis = list(jieba.cut(sample["predict"])) + reference = list(jieba.cut(sample["label"])) + + bleu_score = sentence_bleu( + [list(sample["label"])], + list(sample["predict"]), + smoothing_function=SmoothingFunction().method3, + ) + + if len(" ".join(hypothesis).split()) == 0 or len(" ".join(reference).split()) == 0: + result = {"rouge-1": {"f": 0.0}, "rouge-2": {"f": 0.0}, "rouge-l": {"f": 0.0}} + else: + rouge = Rouge() + scores = rouge.get_scores(" ".join(hypothesis), " ".join(reference)) + result = scores[0] + + metric_result = {} + for k, v in result.items(): + metric_result[k] = round(v["f"] * 100, 4) + + metric_result["bleu-4"] = round(bleu_score * 100, 4) + + return metric_result + + +def main(filename: str): + start_time = time.time() + dataset = load_dataset("json", data_files=filename, split="train") + dataset = dataset.map(compute_metrics, num_proc=8, remove_columns=dataset.column_names) + score_dict = dataset.to_dict() + + average_score = {} + for task, scores in sorted(score_dict.items(), key=lambda x: x[0]): + print(f"{task}: {sum(scores) / len(scores):.4f}") + average_score[task] = sum(scores) / len(scores) + + with open("predictions_score.json", "w", encoding="utf-8") as f: + json.dump(average_score, f, indent=4) + + print(f"\nDone in {time.time() - start_time:.3f}s.\nScore file saved to predictions_score.json") + + +if __name__ == "__main__": + fire.Fire(main) diff --git a/scripts/hf2dcp.py b/scripts/hf2dcp.py new file mode 100644 index 0000000000000000000000000000000000000000..da51580b6203efd6060ec5c5c94d3c9fa798a027 --- /dev/null +++ b/scripts/hf2dcp.py @@ -0,0 +1,63 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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. + +"""Convert a HuggingFace model to DCP checkpoint format. + +Usage: + python scripts/hf2dcp.py convert --hf_path=/path/to/hf --dcp_path=/path/to/dcp + +Arguments: + hf_path: Path to the HuggingFace model directory. + dcp_path: Output path (directory) for DCP checkpoint. +""" + +import fire +import torch +import torch.distributed.checkpoint as dcp +import transformers +from transformers import AutoConfig + + +def convert(hf_path: str, dcp_path: str) -> None: + """Convert HF model weights to DCP. + + Args: + hf_path: HuggingFace model directory. + dcp_path: Output path (directory) for DCP checkpoint. + """ + if not hf_path or not dcp_path: + raise ValueError("Both 'hf_path' and 'dcp_path' are required.") + + print(f"Loading HF model from {hf_path}...") + config = AutoConfig.from_pretrained(hf_path) + architectures = getattr(config, "architectures", []) + if architectures: + model_cls = getattr(transformers, architectures[0], transformers.AutoModelForCausalLM) + else: + model_cls = transformers.AutoModelForCausalLM + + model = model_cls.from_pretrained(hf_path, device_map="cpu", torch_dtype=torch.bfloat16) + + print(f"Saving to DCP format at {dcp_path}...") + dcp.save(model.state_dict(), checkpoint_id=dcp_path) + print("Done!") + + +def help() -> None: + """Show help message.""" + print(__doc__) + + +if __name__ == "__main__": + fire.Fire({"convert": convert, "help": help, "--convert": convert}) diff --git a/scripts/llama_pro.py b/scripts/llama_pro.py new file mode 100644 index 0000000000000000000000000000000000000000..7e4b9448505769104f5155ce7bc4c3ef9ec01bc6 --- /dev/null +++ b/scripts/llama_pro.py @@ -0,0 +1,129 @@ +# Copyright 2025 Tencent Inc. and the LlamaFactory team. +# +# This code is inspired by the Tencent's LLaMA-Pro library. +# https://github.com/TencentARC/LLaMA-Pro/blob/main/scripts/block_expansion.py +# +# 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 collections import OrderedDict +from typing import TYPE_CHECKING + +import fire +import torch +from huggingface_hub import split_torch_state_dict_into_shards +from safetensors.torch import save_file +from tqdm import tqdm +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, PreTrainedModel +from transformers.modeling_utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME + + +if TYPE_CHECKING: + from transformers import PretrainedConfig + + +def change_name(name: str, old_index: int, new_index: int) -> str: + return name.replace(f".{old_index:d}.", f".{new_index:d}.") + + +def block_expansion( + model_name_or_path: str, + output_dir: str, + num_expand: int, + shard_size: str = "5GB", + save_safetensors: bool = True, +): + r"""Perform block expansion for LLaMA, Mistral, Qwen2 or Yi models. + + Usage: python llama_pro.py --model_name_or_path meta-llama/Llama-2-7b-hf --output_dir llama2_pro --num_expand 8 + """ + config: PretrainedConfig = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True) + num_layers = getattr(config, "num_hidden_layers") + if num_layers % num_expand != 0: + raise ValueError(f"`num_layers` {num_layers} should be divisible by `num_expand` {num_expand}.") + + setattr(config, "num_hidden_layers", num_layers + num_expand) + config.save_pretrained(output_dir) + + tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True) + tokenizer.save_pretrained(output_dir) + + print(f"Expanding model of {num_layers} layers to {num_layers + num_expand} layers.") + model = AutoModelForCausalLM.from_pretrained( + model_name_or_path, torch_dtype="auto", device_map="cpu", trust_remote_code=True, low_cpu_mem_usage=True + ) + assert isinstance(model, PreTrainedModel) # type hint + if save_safetensors and getattr(model.config, "tie_word_embeddings", False): + del model.lm_head # safetensors does not allow shared weights + + split = num_layers // num_expand + layer_cnt = 0 + state_dict = model.state_dict() + output_state_dict: dict[str, torch.Tensor] = OrderedDict() + for i in range(num_layers): + for key, value in state_dict.items(): + if f".{i:d}." in key: + output_state_dict[change_name(key, i, layer_cnt)] = value + + print(f"Add layer {layer_cnt} copied from layer {i}.") + layer_cnt += 1 + if (i + 1) % split == 0: + for key, value in state_dict.items(): + if f".{i:d}." in key: + if "down_proj" in key or "o_proj" in key: + output_state_dict[change_name(key, i, layer_cnt)] = torch.zeros_like(value) + else: + output_state_dict[change_name(key, i, layer_cnt)] = torch.clone(value) + + print(f"Add layer {layer_cnt} expanded from layer {i}.") + layer_cnt += 1 + + for key, value in state_dict.items(): + if key not in output_state_dict: + output_state_dict[key] = value + + weights_name = SAFE_WEIGHTS_NAME if save_safetensors else WEIGHTS_NAME + filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors") + state_dict_split = split_torch_state_dict_into_shards( + output_state_dict, filename_pattern=filename_pattern, max_shard_size=shard_size + ) + for shard_file, tensors in tqdm(state_dict_split.filename_to_tensors.items(), desc="Save weights"): + shard = {tensor: output_state_dict[tensor].contiguous() for tensor in tensors} + if save_safetensors: + save_file(shard, os.path.join(output_dir, shard_file), metadata={"format": "pt"}) + else: + torch.save(shard, os.path.join(output_dir, shard_file)) + + if not state_dict_split.is_sharded: + print(f"Model weights saved in {os.path.join(output_dir, weights_name)}.") + else: + index = { + "metadata": state_dict_split.metadata, + "weight_map": state_dict_split.tensor_to_filename, + } + index_name = SAFE_WEIGHTS_INDEX_NAME if save_safetensors else WEIGHTS_INDEX_NAME + with open(os.path.join(output_dir, index_name), "w", encoding="utf-8") as f: + json.dump(index, f, indent=2, sort_keys=True) + + print(f"Model weights saved in {output_dir}.") + + print("- Fine-tune this model with:") + print(f"model_name_or_path: {output_dir}") + print("finetuning_type: freeze") + print(f"freeze_trainable_layers: {num_expand}") + print("use_llama_pro: true") + + +if __name__ == "__main__": + fire.Fire(block_expansion) diff --git a/scripts/loftq_init.py b/scripts/loftq_init.py new file mode 100644 index 0000000000000000000000000000000000000000..3a7933889be55a254b2417e9dca2ce2b7d691401 --- /dev/null +++ b/scripts/loftq_init.py @@ -0,0 +1,88 @@ +# Copyright 2025 HuggingFace Inc. and the LlamaFactory team. +# +# This code is based on the HuggingFace's PEFT library. +# https://github.com/huggingface/peft/blob/v0.10.0/examples/loftq_finetuning/quantize_save_load.py +# +# 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 TYPE_CHECKING + +import fire +from peft import LoftQConfig, LoraConfig, TaskType, get_peft_model +from transformers import AutoModelForCausalLM, AutoTokenizer + + +if TYPE_CHECKING: + from transformers import PreTrainedModel + + +def quantize_loftq( + model_name_or_path: str, + output_dir: str, + loftq_bits: int = 4, + loftq_iter: int = 4, + lora_alpha: int = None, + lora_rank: int = 16, + lora_dropout: float = 0, + lora_target: tuple = ("q_proj", "v_proj"), + save_safetensors: bool = True, +): + r"""Initialize LoRA weights with LoRA-fine-tuning-aware Quantization (LoftQ). + + Usage: python loftq_init.py --model_name_or_path path_to_model --output_dir output_dir + """ + if isinstance(lora_target, str): + lora_target = [name.strip() for name in lora_target.split(",")] + + tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True, torch_dtype="auto") + + loftq_config = LoftQConfig(loftq_bits=loftq_bits, loftq_iter=loftq_iter) + lora_config = LoraConfig( + task_type=TaskType.CAUSAL_LM, + inference_mode=True, + r=lora_rank, + lora_alpha=lora_alpha if lora_alpha is not None else lora_rank * 2, + lora_dropout=lora_dropout, + target_modules=lora_target, + init_lora_weights="loftq", + loftq_config=loftq_config, + ) + + # Init LoftQ model + print("Initializing LoftQ weights, it may be take several minutes, wait patiently.") + peft_model = get_peft_model(model, lora_config) + loftq_dir = os.path.join(output_dir, "loftq_init") + + # Save LoftQ model + setattr(peft_model.peft_config["default"], "base_model_name_or_path", os.path.abspath(output_dir)) + setattr(peft_model.peft_config["default"], "init_lora_weights", True) # don't apply loftq again + peft_model.save_pretrained(loftq_dir, safe_serialization=save_safetensors) + print(f"Adapter weights saved in {loftq_dir}") + + # Save base model + base_model: PreTrainedModel = peft_model.unload() + base_model.save_pretrained(output_dir, safe_serialization=save_safetensors) + tokenizer.save_pretrained(output_dir) + print(f"Model weights saved in {output_dir}") + + print("- Fine-tune this model with:") + print(f"model_name_or_path: {output_dir}") + print(f"adapter_name_or_path: {loftq_dir}") + print("finetuning_type: lora") + print(f"quantization_bit: {loftq_bits}") + + +if __name__ == "__main__": + fire.Fire(quantize_loftq) diff --git a/scripts/megatron_merge.py b/scripts/megatron_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..e8e9e12a316942375a0a2dde3d2e2637992ed963 --- /dev/null +++ b/scripts/megatron_merge.py @@ -0,0 +1,130 @@ +# Copyright 2025 the ROLL team and the LlamaFactory team. +# +# This code is modified from the ROLL library. +# https://github.com/alibaba/ROLL/blob/main/mcore_adapter/tools/convert.py +# +# 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 fire +import torch +from mcore_adapter.models.converter.post_converter import convert_checkpoint_to_hf, convert_checkpoint_to_mca +from mcore_adapter.training_args import DistributingParallelArguments +from mcore_adapter.utils import get_logger +from transformers import AutoConfig + + +logger = get_logger(__name__) + + +def convert_mca_to_hf( + checkpoint_path: str, + output_path: str = "./output", + bf16: bool = False, + fp16: bool = False, + convert_model_max_length: int | None = None, +): + """Convert megatron checkpoint to HuggingFace format. + + Args: + checkpoint_path: Path to the checkpoint to convert + output_path: Path to save the converted checkpoint + bf16: Use bfloat16 precision + fp16: Use float16 precision + convert_model_max_length: Change the model_max_length in hf config.json + """ + if bf16 and fp16: + raise ValueError("bf16 and fp16 cannot be both True.") + + torch_dtype = None + if bf16: + torch_dtype = torch.bfloat16 + elif fp16: + torch_dtype = torch.float16 + + convert_checkpoint_to_hf(checkpoint_path, output_path, torch_dtype=torch_dtype) + + if convert_model_max_length is not None: + config = AutoConfig.from_pretrained(output_path, trust_remote_code=True) + config.model_max_length = convert_model_max_length + config.save_pretrained(output_path) + + +def convert( + checkpoint_path: str, + output_path: str = "./output", + bf16: bool = False, + fp16: bool = False, + convert_model_max_length: int | None = None, + tensor_model_parallel_size: int = 1, + pipeline_model_parallel_size: int = 1, + expert_model_parallel_size: int = 1, + virtual_pipeline_model_parallel_size: int | None = None, + moe_grouped_gemm: bool | None = None, +): + """Convert checkpoint between MCA and HuggingFace formats. + + Args: + checkpoint_path: Path to the checkpoint to convert + output_path: Path to save the converted checkpoint + bf16: Use bfloat16 precision + fp16: Use float16 precision + convert_model_max_length: Change the model_max_length in hf config.json + tensor_model_parallel_size: Tensor model parallel size + pipeline_model_parallel_size: Pipeline model parallel size + expert_model_parallel_size: Expert model parallel size + virtual_pipeline_model_parallel_size: Virtual pipeline model parallel size + moe_grouped_gemm: Use grouped gemm for MoE experts. When enabled, expert + weights are stored in a flattened format (linear_fc1.weight0, weight1, ...) + rather than per-expert format (local_experts.0.linear_fc1.weight, ...). + Must match the format used when saving the checkpoint. + """ + if bf16 and fp16: + raise ValueError("bf16 and fp16 cannot be both True.") + + mca_config_path = os.path.join(checkpoint_path, "mca_config.json") + from_mca = os.path.exists(mca_config_path) + + if not from_mca: + dist_args = DistributingParallelArguments( + tensor_model_parallel_size=tensor_model_parallel_size, + pipeline_model_parallel_size=pipeline_model_parallel_size, + expert_model_parallel_size=expert_model_parallel_size, + virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size, + moe_grouped_gemm=moe_grouped_gemm, + transformer_impl="transformer_engine", # hard code here since we default using te for training + ) + convert_checkpoint_to_mca( + checkpoint_path, + output_path, + dist_args, + bf16=bf16, + fp16=fp16, + ) + else: + convert_mca_to_hf( + checkpoint_path=checkpoint_path, + output_path=output_path, + bf16=bf16, + fp16=fp16, + convert_model_max_length=convert_model_max_length, + ) + + +def main(): + fire.Fire(convert) + + +if __name__ == "__main__": + main() diff --git a/scripts/pissa_init.py b/scripts/pissa_init.py new file mode 100644 index 0000000000000000000000000000000000000000..405a1472b42b8981a8ae52b612acdf07c60c66ad --- /dev/null +++ b/scripts/pissa_init.py @@ -0,0 +1,86 @@ +# Copyright 2025 HuggingFace Inc. and the LlamaFactory team. +# +# This code is based on the HuggingFace's PEFT library. +# https://github.com/huggingface/peft/blob/v0.11.0/examples/pissa_finetuning/preprocess.py +# +# 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 TYPE_CHECKING + +import fire +from peft import LoraConfig, TaskType, get_peft_model +from transformers import AutoModelForCausalLM, AutoTokenizer + + +if TYPE_CHECKING: + from transformers import PreTrainedModel + + +def quantize_pissa( + model_name_or_path: str, + output_dir: str, + pissa_iter: int = 16, + lora_alpha: int = None, + lora_rank: int = 16, + lora_dropout: float = 0, + lora_target: tuple = ("q_proj", "v_proj"), + save_safetensors: bool = True, +): + r"""Initialize LoRA weights with Principal Singular values and Singular vectors Adaptation (PiSSA). + + Usage: python pissa_init.py --model_name_or_path path_to_model --output_dir output_dir + """ + if isinstance(lora_target, str): + lora_target = [name.strip() for name in lora_target.split(",")] + + tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True, torch_dtype="auto") + + lora_config = LoraConfig( + task_type=TaskType.CAUSAL_LM, + r=lora_rank, + lora_alpha=lora_alpha if lora_alpha is not None else lora_rank * 2, + lora_dropout=lora_dropout, + target_modules=lora_target, + init_lora_weights="pissa" if pissa_iter == -1 else f"pissa_niter_{pissa_iter}", + ) + + # Init PiSSA model + peft_model = get_peft_model(model, lora_config) + pissa_dir = os.path.join(output_dir, "pissa_init") + + # Save PiSSA model + setattr(peft_model.peft_config["default"], "base_model_name_or_path", os.path.abspath(output_dir)) + setattr(peft_model.peft_config["default"], "init_lora_weights", True) # don't apply pissa again + peft_model.save_pretrained(pissa_dir, safe_serialization=save_safetensors) + print(f"Adapter weights saved in {pissa_dir}") + + # Save base model + base_model: PreTrainedModel = peft_model.unload() + base_model.save_pretrained(output_dir, safe_serialization=save_safetensors) + tokenizer.save_pretrained(output_dir) + print(f"Model weights saved in {output_dir}") + + print("- Fine-tune this model with:") + print(f"model_name_or_path: {output_dir}") + print(f"adapter_name_or_path: {pissa_dir}") + print("finetuning_type: lora") + print("pissa_init: false") + print("pissa_convert: true") + print("- and optionally with:") + print("quantization_bit: 4") + + +if __name__ == "__main__": + fire.Fire(quantize_pissa) diff --git a/scripts/qwen_omni_merge.py b/scripts/qwen_omni_merge.py new file mode 100644 index 0000000000000000000000000000000000000000..7236d23c3c023d4d454dfefeaaa96229ddb0afa5 --- /dev/null +++ b/scripts/qwen_omni_merge.py @@ -0,0 +1,140 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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. + +"""Why we need this script for qwen_omni? + +Because the qwen_omni model is constructed by two parts: +1. [Thinker]:[audio_encoder, vision_encoder, LLM backbone], which our repository does support to post-training. +2. [Talker]: [audio_decoder, wave_model], which is not supported to post-training without specific tokenizer. +When we post-training the model, we exactly train the [Thinker] part, and the [Talker] part is dropped. +So, to get the complete model, we need to merge the [Talker] part back to the [Thinker] part. +LoRA mode: [Thinker + LoRA weights] + [Original Talker] -> [Omni model] +Full mode: [Thinker] + [Original Talker] -> [Omni model] +For Processor, we do saved the processor from trained model instead of the original model. +""" + +import os +import shutil + +import fire +from peft import PeftModel +from transformers import AutoConfig, AutoModelForTextToWaveform, AutoProcessor +from transformers.utils import cached_file + + +def merge_lora( + model_path: str, + lora_path: str, + save_path: str = "./merged_model_checkpoint", + extra_file: str = "spk_dict.pt", + submodule_name: str = "thinker", +): + """Load the original model, merge the LoRA weights. + + For a specified submodule, and save the final merged model along with its configurations. + + Args: + model_path (str): Path to the original model directory. + lora_path (str): Path to the directory containing LoRA weights. + save_path (str): Directory where the merged model and configurations will be saved. + extra_file (str): Name of the extra file to be copied (default: "spk_dict.pt"). + submodule_name (str): Name of the submodule to merge (default: "thinker"). + """ + # 1. Load the original model + model = AutoModelForTextToWaveform.from_pretrained(model_path, torch_dtype="auto", device_map="cpu") + print("Successfully loaded the original model.") + + # 2. Extract the submodule to be merged (e.g., model.thinker) + if not hasattr(model, submodule_name): + raise AttributeError(f"The model does not have a submodule named '{submodule_name}'.") + + base_submodule = getattr(model, submodule_name) + print(f"Successfully extracted submodule: {submodule_name}.") + + # 3. Load the LoRA weights onto the extracted submodule + lora_model = PeftModel.from_pretrained(base_submodule, lora_path) + processor = AutoProcessor.from_pretrained(lora_path) + print("Successfully loaded LoRA weights and processor.") + + # 4. Merge the LoRA weights into the submodule and unload the LoRA modules + merged_submodule = lora_model.merge_and_unload() + print("Successfully merged LoRA weights.") + + # 5. Replace the original submodule with the merged submodule in the model + setattr(model, submodule_name, merged_submodule) + + # 6. Save the final merged model along with the tokenizer and processor configuration + model.save_pretrained(save_path) + processor.save_pretrained(save_path) + print(f"Merged model and processor saved to {save_path}.") + + try: + source_file = cached_file(path_or_repo_id=model_path, filename=extra_file) + shutil.copy(source_file, os.path.join(save_path, extra_file)) + print(f"File '{extra_file}' copied from {model_path} to {save_path}.") + except Exception: + print(f"File '{extra_file}' not found in {model_path}, skipping copy.") + + +def save_full_model( + model_path: str, + thinker_path: str, + save_path: str = "./merged_model_checkpoint", + extra_file: str = "spk_dict.pt", +): + """Load the saved thinker module and the original model, replace the thinker in the original model. + + Then save the complete model along with its tokenizer and processor configuration. + + Args: + model_path (str): Directory path of the original model. + thinker_path (str): Path to the saved thinker weights. + save_path (str): Directory where the merged model and configurations will be saved. + extra_file (str): Name of the extra file to be copied (default: "spk_dict.pt"). + """ + # 1. Load the saved thinker module and the original model + config = AutoConfig.from_pretrained(model_path) + if getattr(config, "model_type") == "qwen2_5_omni": + from transformers.models.qwen2_5_omni import Qwen2_5OmniThinkerForConditionalGeneration # type: ignore + + ThinkerClass = Qwen2_5OmniThinkerForConditionalGeneration + elif getattr(config, "model_type") == "qwen3_omni_moe": + from transformers.models.qwen3_omni_moe import Qwen3OmniMoeThinkerForConditionalGeneration # type: ignore + + ThinkerClass = Qwen3OmniMoeThinkerForConditionalGeneration + else: + raise ValueError(f"Unsupported model type: {getattr(config, 'model_type')}.") + + thinker = ThinkerClass.from_pretrained(thinker_path, torch_dtype="auto", device_map="cpu") + base_model = AutoModelForTextToWaveform.from_pretrained(model_path, torch_dtype="auto", device_map="cpu") + base_model.thinker = thinker + processor = AutoProcessor.from_pretrained(thinker_path) + print("Successfully loaded model weights and processor.") + + # 2. Save the complete model along with its tokenizer and processor configuration + base_model.save_pretrained(save_path) + processor.save_pretrained(save_path) + print(f"Merged model and processor saved to {save_path}.") + + # 3. Copy the extra file from the base model directory to the save_path + try: + source_file = cached_file(path_or_repo_id=model_path, filename=extra_file) + shutil.copy(source_file, os.path.join(save_path, extra_file)) + print(f"File '{extra_file}' copied from {model_path} to {save_path}.") + except Exception: + print(f"File '{extra_file}' not found in {model_path}, skipping copy.") + + +if __name__ == "__main__": + fire.Fire({"save_full": save_full_model, "merge_lora": merge_lora}) diff --git a/scripts/stat_utils/cal_flops.py b/scripts/stat_utils/cal_flops.py new file mode 100644 index 0000000000000000000000000000000000000000..3dc049959bb057737d4320a1960b1a43465b4662 --- /dev/null +++ b/scripts/stat_utils/cal_flops.py @@ -0,0 +1,49 @@ +# Copyright 2025 Microsoft Corporation and the LlamaFactory team. +# +# This code is inspired by the Microsoft's DeepSpeed library. +# https://www.deepspeed.ai/tutorials/flops-profiler/ +# +# 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 fire +import torch +from deepspeed.accelerator import get_accelerator # type: ignore +from deepspeed.profiling.flops_profiler import get_model_profile # type: ignore + +from llamafactory.chat import ChatModel + + +def calculate_flops( + model_name_or_path: str, + batch_size: int = 1, + seq_length: int = 512, + flash_attn: str = "auto", +): + r"""Calculate the flops of pre-trained models. + + Usage: python cal_flops.py --model_name_or_path path_to_model --batch_size 1 --seq_length 512 + """ + with get_accelerator().device(0): + chat_model = ChatModel(dict(model_name_or_path=model_name_or_path, template="empty", flash_attn=flash_attn)) + fake_input = torch.ones((batch_size, seq_length), dtype=torch.long, device=chat_model.engine.model.device) + input_dict = {"input_ids": fake_input, "labels": fake_input.clone()} + flops, macs, params = get_model_profile( + chat_model.engine.model, kwargs=input_dict, print_profile=True, detailed=True + ) + print("FLOPs:", flops) + print("MACs:", macs) + print("Params:", params) + + +if __name__ == "__main__": + fire.Fire(calculate_flops) diff --git a/scripts/stat_utils/cal_lr.py b/scripts/stat_utils/cal_lr.py new file mode 100644 index 0000000000000000000000000000000000000000..eb35c47e7ab6b196e5b95e9310e0ad7d27e12899 --- /dev/null +++ b/scripts/stat_utils/cal_lr.py @@ -0,0 +1,98 @@ +# Copyright 2025 imoneoi and the LlamaFactory team. +# +# This code is inspired by the imoneoi's OpenChat library. +# https://github.com/imoneoi/openchat/blob/3.6.0/ochat/training_deepspeed/train.py +# +# 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 math +from typing import Literal + +import fire +import torch +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import DataCollatorForLanguageModeling + +from llamafactory.data import MultiModalDataCollatorForSeq2Seq, get_dataset, get_template_and_fix_tokenizer +from llamafactory.extras.constants import IGNORE_INDEX +from llamafactory.hparams import get_train_args +from llamafactory.model import load_tokenizer + + +BASE_LR = 3e-4 # 1.5e-4 for 30B-70B models +BASE_BS = 4_000_000 # from llama paper + + +def calculate_lr( + model_name_or_path: str, + batch_size: int, # total batch size, namely (batch size * gradient accumulation * world size) + stage: Literal["pt", "sft"] = "sft", + dataset: str = "alpaca_en_demo", + dataset_dir: str = "data", + template: str = "default", + cutoff_len: int = 2048, # i.e. maximum input length during training + is_mistral_or_gemma: bool = False, # mistral and gemma models opt for a smaller learning rate, + packing: bool = False, +): + r"""Calculate the optimal learning rate for 7B/13B models using LLaMA's hyper-parameters. + + Usage: + python cal_lr.py --model_name_or_path path_to_model --dataset alpaca_en_demo --cutoff_len 1024 --batch_size 16 + """ + model_args, data_args, training_args, _, _ = get_train_args( + dict( + stage=stage, + model_name_or_path=model_name_or_path, + dataset=dataset, + dataset_dir=dataset_dir, + template=template, + cutoff_len=cutoff_len, + packing=packing, + preprocessing_num_workers=16, + output_dir="dummy_dir", + overwrite_cache=True, + do_train=True, + ) + ) + tokenizer_module = load_tokenizer(model_args) + tokenizer = tokenizer_module["tokenizer"] + template = get_template_and_fix_tokenizer(tokenizer, data_args) + trainset = get_dataset(template, model_args, data_args, training_args, stage, **tokenizer_module)["train_dataset"] + if stage == "pt": + data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) + elif stage == "sft": + data_collator = MultiModalDataCollatorForSeq2Seq( + template=template, tokenizer=tokenizer, label_pad_token_id=IGNORE_INDEX + ) + else: + raise NotImplementedError(f"Stage does not supported: {stage}.") + + dataloader = DataLoader(trainset, batch_size, shuffle=False, collate_fn=data_collator, pin_memory=True) + valid_tokens, total_tokens = 0, 0 + for batch in tqdm(dataloader, desc="Collecting valid tokens"): + valid_tokens += torch.sum(batch["labels"] != IGNORE_INDEX).item() + total_tokens += torch.numel(batch["labels"]) + + valid_ratio = valid_tokens / total_tokens + token_batch_size = cutoff_len * batch_size * valid_ratio + lr = BASE_LR * math.sqrt(token_batch_size / BASE_BS) # lr ~ sqrt(batch_size) + lr = lr / 6.0 if is_mistral_or_gemma else lr + print( + f"Optimal learning rate is {lr:.2e} for valid ratio% {valid_ratio * 100:.2f} " + f"and effective token batch size {token_batch_size:.2f}" + ) + + +if __name__ == "__main__": + fire.Fire(calculate_lr) diff --git a/scripts/stat_utils/cal_mfu.py b/scripts/stat_utils/cal_mfu.py new file mode 100644 index 0000000000000000000000000000000000000000..f1d4446eef71daa3ffa143fc456a0f58c928d491 --- /dev/null +++ b/scripts/stat_utils/cal_mfu.py @@ -0,0 +1,161 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 fire +import torch +import torch.distributed as dist +from transformers import AutoConfig + +from llamafactory.train.tuner import run_exp + + +BASE = 2 # gemm (add + mul) + + +def compute_model_flops( + model_name_or_path: str, + total_batch_size: int, + seq_length: int, + include_backward: bool = True, + include_recompute: bool = False, + include_flashattn: bool = False, +) -> int: + r"""Calculate the FLOPs of model per forward/backward pass.""" + config = AutoConfig.from_pretrained(model_name_or_path) + hidden_size = getattr(config, "hidden_size", None) + vocab_size = getattr(config, "vocab_size", None) + intermediate_size = getattr(config, "intermediate_size", None) + num_attention_heads = getattr(config, "num_attention_heads", None) + num_key_value_heads = getattr(config, "num_key_value_heads", None) + num_hidden_layers = getattr(config, "num_hidden_layers", None) + tie_word_embeddings = getattr(config, "tie_word_embeddings", False) + + # mlp module + mlp_flops_per_token = 3 * BASE * hidden_size * intermediate_size # up, gate, down + mlp_flops = total_batch_size * seq_length * num_hidden_layers * mlp_flops_per_token + + # attn projector module + q_flops_per_token = BASE * hidden_size * hidden_size + o_flops_per_token = BASE * hidden_size * hidden_size + k_flops_per_token = BASE * hidden_size * hidden_size * num_key_value_heads // num_attention_heads + v_flops_per_token = BASE * hidden_size * hidden_size * num_key_value_heads // num_attention_heads + attn_proj_flops_per_token = q_flops_per_token + o_flops_per_token + k_flops_per_token + v_flops_per_token + attn_proj_flops = total_batch_size * seq_length * num_hidden_layers * attn_proj_flops_per_token + + # attn sdpa module + sdpa_flops_per_layer = 2 * BASE * hidden_size * seq_length * seq_length # (q * k^T) * v + sdpa_flops = total_batch_size * num_hidden_layers * sdpa_flops_per_layer + + # embedding module + embedding_flops_per_token = hidden_size * vocab_size + embedding_flops = total_batch_size * seq_length * embedding_flops_per_token + if tie_word_embeddings is False: + embedding_flops *= 2 + + non_embedding_flops = mlp_flops + attn_proj_flops + sdpa_flops + non_embedding_coeff, embedding_coeff = 1, 1 + if include_backward: + non_embedding_coeff += 2 + embedding_coeff += 2 + + if include_recompute: + non_embedding_coeff += 1 + + total_flops = non_embedding_coeff * non_embedding_flops + embedding_coeff * embedding_flops + + if include_flashattn: + total_flops += sdpa_flops + + return total_flops + + +def compute_device_flops(world_size: int) -> float: + r"""Calculate the FLOPs of the device capability per second.""" + device_name = torch.cuda.get_device_name() + if "H100" in device_name or "H800" in device_name: + return 989 * 1e12 * world_size + elif "A100" in device_name or "A800" in device_name: + return 312 * 1e12 * world_size + elif "V100" in device_name: + return 125 * 1e12 * world_size + elif "4090" in device_name: + return 98 * 1e12 * world_size + else: + raise NotImplementedError(f"Device not supported: {device_name}.") + + +def calculate_mfu( + model_name_or_path: str, + batch_size: int = 1, + seq_length: int = 1024, + num_steps: int = 100, + finetuning_type: str = "lora", + flash_attn: str = "auto", + deepspeed_stage: int = 0, + disable_gc: bool = False, + liger_kernel: bool = False, + unsloth_gc: bool = False, +) -> float: + r"""Calculate MFU for given model and hyper-params. + + Usage: python cal_mfu.py --model_name_or_path path_to_model --batch_size 1 --seq_length 1024 + """ + args = { + "model_name_or_path": model_name_or_path, + "flash_attn": flash_attn, + "disable_gradient_checkpointing": disable_gc, + "enable_liger_kernel": liger_kernel, + "use_unsloth_gc": unsloth_gc, + "stage": "pt", + "do_train": True, + "finetuning_type": finetuning_type, + "dataset": "c4_demo", + "cutoff_len": seq_length, + "output_dir": os.path.join("saves", "test_mfu"), + "logging_strategy": "no", + "save_strategy": "no", + "save_only_model": True, + "overwrite_output_dir": True, + "per_device_train_batch_size": batch_size, + "max_steps": num_steps, + "bf16": True, + } + if deepspeed_stage in [2, 3]: + args["deepspeed"] = f"examples/deepspeed/ds_z{deepspeed_stage}_config.json" + + run_exp(args) + if dist.is_initialized(): + dist.barrier() + world_size = dist.get_world_size() + else: + world_size = 1 + + if int(os.getenv("LOCAL_RANK", "0")) == 0: + with open(os.path.join("saves", "test_mfu", "all_results.json"), encoding="utf-8") as f: + result = json.load(f) + + total_batch_size = batch_size * world_size + mfu_value = ( + result["train_steps_per_second"] + * compute_model_flops(model_name_or_path, total_batch_size, seq_length) + / compute_device_flops(world_size) + ) + print(f"MFU: {mfu_value * 100:.2f}%") + + +if __name__ == "__main__": + fire.Fire(calculate_mfu) diff --git a/scripts/stat_utils/cal_ppl.py b/scripts/stat_utils/cal_ppl.py new file mode 100644 index 0000000000000000000000000000000000000000..56b3c8d1177c4eb8139eb2b38de329a190ca3023 --- /dev/null +++ b/scripts/stat_utils/cal_ppl.py @@ -0,0 +1,134 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 dataclasses import dataclass +from typing import Any, Literal + +import fire +import torch +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers import DataCollatorForLanguageModeling + +from llamafactory.data import MultiModalDataCollatorForSeq2Seq, get_dataset, get_template_and_fix_tokenizer +from llamafactory.extras.constants import IGNORE_INDEX +from llamafactory.hparams import get_train_args +from llamafactory.model import load_model, load_tokenizer + + +@dataclass +class PairwiseDataCollatorWithPadding(MultiModalDataCollatorForSeq2Seq): + r"""Data collator for pairwise data.""" + + train_on_prompt: bool = False + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, torch.Tensor]: + r"""Pad batched data to the longest sequence in the batch.""" + chosen_features = [] + for feature in features: + chosen_features.append( + { + "input_ids": feature["chosen_input_ids"], + "attention_mask": feature["chosen_attention_mask"], + "labels": feature["chosen_input_ids"] if self.train_on_prompt else feature["chosen_labels"], + "images": feature["images"], + "videos": feature["videos"], + "audios": feature["audios"], + } + ) + + return super().__call__(chosen_features) + + +def calculate_ppl( + model_name_or_path: str, + save_name: str = "ppl.json", + batch_size: int = 4, + stage: Literal["pt", "sft", "rm"] = "sft", + dataset: str = "alpaca_en_demo", + dataset_dir: str = "data", + template: str = "default", + cutoff_len: int = 2048, + max_samples: int | None = None, + train_on_prompt: bool = False, +): + r"""Calculate the ppl on the dataset of the pre-trained models. + + Usage: export CUDA_VISIBLE_DEVICES=0 + python cal_ppl.py --model_name_or_path path_to_model --dataset alpaca_en_demo --save_name ppl.json + """ + model_args, data_args, training_args, finetuning_args, _ = get_train_args( + dict( + stage=stage, + model_name_or_path=model_name_or_path, + dataset=dataset, + dataset_dir=dataset_dir, + template=template, + cutoff_len=cutoff_len, + max_samples=max_samples, + train_on_prompt=train_on_prompt, + preprocessing_num_workers=16, + output_dir="dummy_dir", + overwrite_cache=True, + do_train=True, + ) + ) + tokenizer_module = load_tokenizer(model_args) + tokenizer = tokenizer_module["tokenizer"] + template = get_template_and_fix_tokenizer(tokenizer, data_args) + trainset = get_dataset(template, model_args, data_args, training_args, stage, **tokenizer_module)["train_dataset"] + model = load_model(tokenizer, model_args, finetuning_args, is_trainable=False) + if stage == "pt": + data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) + elif stage == "sft": + data_collator = MultiModalDataCollatorForSeq2Seq( + template=template, tokenizer=tokenizer, label_pad_token_id=IGNORE_INDEX + ) + elif stage == "rm": + data_collator = PairwiseDataCollatorWithPadding( + template=template, tokenizer=tokenizer, label_pad_token_id=IGNORE_INDEX, train_on_prompt=train_on_prompt + ) + else: + raise NotImplementedError(f"Stage does not supported: {stage}.") + + dataloader = DataLoader(trainset, batch_size, shuffle=False, collate_fn=data_collator, pin_memory=True) + criterion = torch.nn.CrossEntropyLoss(reduction="none") + total_ppl = 0 + perplexities = [] + batch: dict[str, torch.Tensor] + with torch.no_grad(): + for batch in tqdm(dataloader, desc="Computing perplexities"): + batch = batch.to(model.device) + outputs = model(**batch) + shift_logits: torch.Tensor = outputs["logits"][..., :-1, :] + shift_labels: torch.Tensor = batch["labels"][..., 1:] + loss_mask = shift_labels != IGNORE_INDEX + flatten_logits = shift_logits.contiguous().view(shift_labels.size(0) * shift_labels.size(1), -1) + flatten_labels = shift_labels.contiguous().view(-1) + token_logps: torch.Tensor = criterion(flatten_logits, flatten_labels) + token_logps = token_logps.contiguous().view(shift_logits.size(0), -1) + sentence_logps = (token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + total_ppl += sentence_logps.exp().sum().item() + perplexities.extend(sentence_logps.exp().tolist()) + + with open(save_name, "w", encoding="utf-8") as f: + json.dump(perplexities, f, indent=2) + + print(f"Average perplexity is {total_ppl / len(perplexities):.2f}") + print(f"Perplexities have been saved at {save_name}.") + + +if __name__ == "__main__": + fire.Fire(calculate_ppl) diff --git a/scripts/stat_utils/length_cdf.py b/scripts/stat_utils/length_cdf.py new file mode 100644 index 0000000000000000000000000000000000000000..c459c8fa0911b86bcce5a6b47ae5b2da1b670b2f --- /dev/null +++ b/scripts/stat_utils/length_cdf.py @@ -0,0 +1,69 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 collections import defaultdict + +import fire +from tqdm import tqdm + +from llamafactory.data import get_dataset, get_template_and_fix_tokenizer +from llamafactory.hparams import get_train_args +from llamafactory.model import load_tokenizer + + +def length_cdf( + model_name_or_path: str, + dataset: str = "alpaca_en_demo", + dataset_dir: str = "data", + template: str = "default", + interval: int = 1000, +): + r"""Calculate the distribution of the input lengths in the dataset. + + Usage: export CUDA_VISIBLE_DEVICES=0 + python length_cdf.py --model_name_or_path path_to_model --dataset alpaca_en_demo --template default + """ + model_args, data_args, training_args, _, _ = get_train_args( + dict( + stage="sft", + model_name_or_path=model_name_or_path, + dataset=dataset, + dataset_dir=dataset_dir, + template=template, + cutoff_len=1_000_000, + preprocessing_num_workers=16, + output_dir="dummy_dir", + overwrite_cache=True, + do_train=True, + ) + ) + tokenizer_module = load_tokenizer(model_args) + template = get_template_and_fix_tokenizer(tokenizer_module["tokenizer"], data_args) + trainset = get_dataset(template, model_args, data_args, training_args, "sft", **tokenizer_module)["train_dataset"] + total_num = len(trainset) + length_dict = defaultdict(int) + for sample in tqdm(trainset["input_ids"], desc="Collecting lengths"): + length_dict[len(sample) // interval * interval] += 1 + + length_tuples = list(length_dict.items()) + length_tuples.sort() + count_accu, prob_accu = 0, 0 + for length, count in length_tuples: + count_accu += count + prob_accu += count / total_num * 100 + print(f"{count_accu:d} ({prob_accu:.2f}%) samples have length < {length + interval}.") + + +if __name__ == "__main__": + fire.Fire(length_cdf) diff --git a/scripts/vllm_infer.py b/scripts/vllm_infer.py new file mode 100644 index 0000000000000000000000000000000000000000..44bdbed83422cbd9616d5e9af47b2e6f3a0966fc --- /dev/null +++ b/scripts/vllm_infer.py @@ -0,0 +1,280 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 gc +import json +import time + +import av +import fire +from datasets import load_dataset +from eval_bleu_rouge import compute_metrics +from tqdm import tqdm +from transformers import Seq2SeqTrainingArguments + +from llamafactory.data import get_dataset, get_template_and_fix_tokenizer +from llamafactory.extras.constants import IGNORE_INDEX +from llamafactory.extras.misc import get_device_count +from llamafactory.extras.packages import is_vllm_available +from llamafactory.hparams import get_infer_args +from llamafactory.model import load_tokenizer + + +if is_vllm_available(): + from vllm import LLM, SamplingParams + from vllm.lora.request import LoRARequest + + +def _need_video_kwargs(template): + NEEDED_TEMPLATE = ["qwen3_vl", "glm4v"] + if any(t in template for t in NEEDED_TEMPLATE): + return True + + return False + + +def vllm_infer( + model_name_or_path: str, + adapter_name_or_path: str = None, + dataset: str = "alpaca_en_demo", + dataset_dir: str = "data", + template: str = "default", + cutoff_len: int = 2048, + max_samples: int | None = None, + vllm_config: str = "{}", + save_name: str = "generated_predictions.jsonl", + matrix_save_name: str = None, + temperature: float = 0.95, + top_p: float = 0.7, + top_k: int = 50, + max_new_tokens: int = 1024, + repetition_penalty: float = 1.0, + skip_special_tokens: bool = True, + default_system: str | None = None, + enable_thinking: bool = True, + seed: int | None = None, + pipeline_parallel_size: int = 1, + image_max_pixels: int = 768 * 768, + image_min_pixels: int = 32 * 32, + video_fps: float = 2.0, + video_maxlen: int = 128, + batch_size: int = 1024, +): + r"""Perform batch generation using vLLM engine, which supports tensor parallelism. + + Usage: python vllm_infer.py --model_name_or_path meta-llama/Llama-2-7b-hf --template llama --dataset alpaca_en_demo + """ + if pipeline_parallel_size > get_device_count(): + raise ValueError("Pipeline parallel size should be smaller than the number of gpus.") + + model_args, data_args, _, generating_args = get_infer_args( + dict( + model_name_or_path=model_name_or_path, + adapter_name_or_path=adapter_name_or_path, + dataset=dataset, + dataset_dir=dataset_dir, + template=template, + cutoff_len=cutoff_len, + max_samples=max_samples, + preprocessing_num_workers=16, + default_system=default_system, + enable_thinking=enable_thinking, + vllm_config=vllm_config, + temperature=temperature, + top_p=top_p, + top_k=top_k, + max_new_tokens=max_new_tokens, + repetition_penalty=repetition_penalty, + ) + ) + + training_args = Seq2SeqTrainingArguments(output_dir="dummy_dir") + tokenizer_module = load_tokenizer(model_args) + tokenizer = tokenizer_module["tokenizer"] + template_obj = get_template_and_fix_tokenizer(tokenizer, data_args) + template_obj.mm_plugin.expand_mm_tokens = False # for vllm generate + + engine_args = { + "model": model_args.model_name_or_path, + "trust_remote_code": True, + "dtype": model_args.infer_dtype, + "max_model_len": cutoff_len + max_new_tokens, + "tensor_parallel_size": (get_device_count() // pipeline_parallel_size) or 1, + "pipeline_parallel_size": pipeline_parallel_size, + "disable_log_stats": True, + "enable_lora": model_args.adapter_name_or_path is not None, + } + if template_obj.mm_plugin.__class__.__name__ != "BasePlugin": + engine_args["limit_mm_per_prompt"] = {"image": 4, "video": 2, "audio": 2} + + if isinstance(model_args.vllm_config, dict): + engine_args.update(model_args.vllm_config) + + model_preparation_start_time = time.time() + llm = LLM(**engine_args) + + # load datasets + dataset_module = get_dataset(template_obj, model_args, data_args, training_args, "ppo", **tokenizer_module) + train_dataset = dataset_module["train_dataset"] + + sampling_params = SamplingParams( + repetition_penalty=generating_args.repetition_penalty or 1.0, # repetition_penalty must > 0 + temperature=generating_args.temperature, + top_p=generating_args.top_p or 1.0, # top_p must > 0 + top_k=generating_args.top_k or -1, # top_k must > 0 + stop_token_ids=template_obj.get_stop_token_ids(tokenizer), + max_tokens=generating_args.max_new_tokens, + skip_special_tokens=skip_special_tokens, + seed=seed, + ) + if model_args.adapter_name_or_path is not None: + lora_request = LoRARequest("default", 1, model_args.adapter_name_or_path[0]) + else: + lora_request = None + + # Store all results in these lists + all_prompts, all_preds, all_labels = [], [], [] + need_video_kwargs = _need_video_kwargs(template) + + model_predict_start_time = time.time() + # Add batch process to avoid the issue of too many files opened + for i in tqdm(range(0, len(train_dataset), batch_size), desc="Processing batched inference"): + vllm_inputs, prompts, labels = [], [], [] + batch = train_dataset[i : min(i + batch_size, len(train_dataset))] + + for j in range(len(batch["input_ids"])): + multi_modal_data = {} + video_metadata_kwargs = None + + if batch["images"][j] is not None: + image = batch["images"][j] + multi_modal_data["image"] = template_obj.mm_plugin._regularize_images( + image, image_max_pixels=image_max_pixels, image_min_pixels=image_min_pixels + )["images"] + + if batch["videos"][j] is not None: + video = batch["videos"][j] + multi_modal_data["video"] = template_obj.mm_plugin._regularize_videos( + video, + image_max_pixels=image_max_pixels, + image_min_pixels=image_min_pixels, + video_fps=video_fps, + video_maxlen=video_maxlen, + )["videos"] + if need_video_kwargs: + container = av.open(video[0], "r") + video_stream = next(stream for stream in container.streams if stream.type == "video") + sampling_indices = template_obj.mm_plugin._get_video_sample_indices( + video_stream, video_fps, video_maxlen + ) + total_frames = video_stream.frames + video_metadata_kwargs = { + "fps": getattr(tokenizer_module["processor"], "video_fps", 24.0), + "do_sample_frames": False, + "total_num_frames": total_frames, + } + video_metadata = dict( + fps=video_fps, + frames_indices=sampling_indices, + total_num_frames=total_frames, + video_backend="opencv", + ) + multi_modal_data["video"] = (multi_modal_data["video"], video_metadata) + + if batch["audios"][j] is not None: + audio = batch["audios"][j] + audio_data = template_obj.mm_plugin._regularize_audios( + audio, + sampling_rate=16000, + ) + multi_modal_data["audio"] = zip(audio_data["audios"], audio_data["sampling_rates"]) + + vllm_input_data = {"prompt_token_ids": batch["input_ids"][j], "multi_modal_data": multi_modal_data or None} + if video_metadata_kwargs is not None: + vllm_input_data["mm_processor_kwargs"] = video_metadata_kwargs + + vllm_inputs.append(vllm_input_data) + prompts.append(tokenizer.decode(batch["input_ids"][j], skip_special_tokens=skip_special_tokens)) + labels.append( + tokenizer.decode( + list(filter(lambda x: x != IGNORE_INDEX, batch["labels"][j])), + skip_special_tokens=skip_special_tokens, + ) + ) + + results = llm.generate(vllm_inputs, sampling_params, lora_request=lora_request) + preds = [result.outputs[0].text for result in results] + + # Accumulate results + all_prompts.extend(prompts) + all_preds.extend(preds) + all_labels.extend(labels) + gc.collect() + + model_predict_end_time = time.time() + # Write all results at once outside the loop + with open(save_name, "w", encoding="utf-8") as f: + for text, pred, label in zip(all_prompts, all_preds, all_labels): + f.write(json.dumps({"prompt": text, "predict": pred, "label": label}, ensure_ascii=False) + "\n") + + print("*" * 70) + print(f"{len(all_prompts)} total generated results have been saved at {save_name}.") + print("*" * 70) + + # Write all matrix results when matrix_save_name is not None, + # The result matrix is referencing src.llamafactory.train.sft.workflow.run_sft # 127~132 + # trainer.save_metrics("predict", predict_results.metrics) + # + # { + # "predict_bleu-4": 4.349975, + # "predict_model_preparation_time": 0.0128, + # "predict_rouge-1": 21.873359375, + # "predict_rouge-2": 4.144340625, + # "predict_rouge-l": 10.83949375, + # "predict_runtime": 131.664, + # "predict_samples_per_second": 0.076, + # "predict_steps_per_second": 0.008 + # } + # + if matrix_save_name is not None: + predict_time = model_predict_end_time - model_predict_start_time + preparation_time = model_predict_start_time - model_preparation_start_time + + start_time = time.time() + dataset = load_dataset("json", data_files=save_name, split="train") + dataset = dataset.map(compute_metrics, num_proc=8, remove_columns=dataset.column_names) + score_dict = dataset.to_dict() + + average_score = {} + for task, scores in sorted(score_dict.items(), key=lambda x: x[0]): + score = sum(scores) / len(scores) if scores else 0.0 + print(f"predict_{task}: {score:.4f}") + average_score["predict_" + task] = score + + average_score["predict_model_preparation_time"] = preparation_time + average_score["predict_runtime"] = predict_time + num_steps = len(range(0, len(train_dataset), batch_size)) + average_score["predict_samples_per_second"] = len(dataset) / predict_time if predict_time > 0 else 0.0 + average_score["predict_steps_per_second"] = num_steps / predict_time if predict_time > 0 else 0.0 + + with open(matrix_save_name, "w", encoding="utf-8") as f: + json.dump(average_score, f, indent=4) + + print("*" * 70) + print(f"\nDone in {time.time() - start_time:.3f}s.\nScore file saved to {matrix_save_name}.") + print("*" * 70) + + +if __name__ == "__main__": + fire.Fire(vllm_infer) diff --git a/src/api.py b/src/api.py new file mode 100644 index 0000000000000000000000000000000000000000..61215459ed91c6fa529a719cb9dac57223754d2e --- /dev/null +++ b/src/api.py @@ -0,0 +1,33 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 uvicorn + +from llamafactory.api.app import create_app +from llamafactory.chat import ChatModel + + +def main(): + chat_model = ChatModel() + app = create_app(chat_model) + api_host = os.getenv("API_HOST", "0.0.0.0") + api_port = int(os.getenv("API_PORT", "8000")) + print(f"Visit http://localhost:{api_port}/docs for API document.") + uvicorn.run(app, host=api_host, port=api_port) + + +if __name__ == "__main__": + main() diff --git a/src/llamafactory/__init__.py b/src/llamafactory/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b1567ef572714881cc464db25d3da3d08a460963 --- /dev/null +++ b/src/llamafactory/__init__.py @@ -0,0 +1,31 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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. + +r"""Efficient fine-tuning of large language models. + +Level: + api, webui > chat, eval, train > data, model > hparams > extras + +Disable version checking: DISABLE_VERSION_CHECK=1 +Enable VRAM recording: RECORD_VRAM=1 +Force using torchrun: FORCE_TORCHRUN=1 +Set logging verbosity: LLAMAFACTORY_VERBOSITY=WARN +Use modelscope: USE_MODELSCOPE_HUB=1 +Use openmind: USE_OPENMIND_HUB=1 +""" + +from .extras.env import VERSION + + +__version__ = VERSION diff --git a/src/llamafactory/__pycache__/__init__.cpython-312.pyc b/src/llamafactory/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..546dc2dc07aa71ad6ba0b6c5184a1b708b3149a5 Binary files /dev/null and b/src/llamafactory/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/llamafactory/__pycache__/cli.cpython-312.pyc b/src/llamafactory/__pycache__/cli.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee225edba911205e4f7b2c756d69f7688ff3f2f8 Binary files /dev/null and b/src/llamafactory/__pycache__/cli.cpython-312.pyc differ diff --git a/src/llamafactory/__pycache__/launcher.cpython-312.pyc b/src/llamafactory/__pycache__/launcher.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..376d1807ff965cc7f536826573174d77bde0dd67 Binary files /dev/null and b/src/llamafactory/__pycache__/launcher.cpython-312.pyc differ diff --git a/src/llamafactory/api/__init__.py b/src/llamafactory/api/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/llamafactory/api/app.py b/src/llamafactory/api/app.py new file mode 100644 index 0000000000000000000000000000000000000000..8ec0679cb7e053058f52bdbf947cb13e554c5ca8 --- /dev/null +++ b/src/llamafactory/api/app.py @@ -0,0 +1,133 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 asyncio +import os +from contextlib import asynccontextmanager +from functools import partial +from typing import Annotated + +from ..chat import ChatModel +from ..extras.constants import EngineName +from ..extras.misc import torch_gc +from ..extras.packages import is_fastapi_available, is_starlette_available, is_uvicorn_available +from .chat import ( + create_chat_completion_response, + create_score_evaluation_response, + create_stream_chat_completion_response, +) +from .protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ModelCard, + ModelList, + ScoreEvaluationRequest, + ScoreEvaluationResponse, +) + + +if is_fastapi_available(): + from fastapi import Depends, FastAPI, HTTPException, status + from fastapi.middleware.cors import CORSMiddleware + from fastapi.security.http import HTTPAuthorizationCredentials, HTTPBearer + + +if is_starlette_available(): + from sse_starlette import EventSourceResponse + + +if is_uvicorn_available(): + import uvicorn + + +async def sweeper() -> None: + while True: + torch_gc() + await asyncio.sleep(300) + + +@asynccontextmanager +async def lifespan(app: "FastAPI", chat_model: "ChatModel"): # collects GPU memory + if chat_model.engine.name == EngineName.HF: + asyncio.create_task(sweeper()) + + yield + torch_gc() + + +def create_app(chat_model: "ChatModel") -> "FastAPI": + root_path = os.getenv("FASTAPI_ROOT_PATH", "") + app = FastAPI(lifespan=partial(lifespan, chat_model=chat_model), root_path=root_path) + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + api_key = os.getenv("API_KEY") + security = HTTPBearer(auto_error=False) + + async def verify_api_key(auth: Annotated[HTTPAuthorizationCredentials | None, Depends(security)]): + if api_key and (auth is None or auth.credentials != api_key): + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key.") + + @app.get( + "/v1/models", + response_model=ModelList, + status_code=status.HTTP_200_OK, + dependencies=[Depends(verify_api_key)], + ) + async def list_models(): + model_card = ModelCard(id=os.getenv("API_MODEL_NAME", "gpt-3.5-turbo")) + return ModelList(data=[model_card]) + + @app.post( + "/v1/chat/completions", + response_model=ChatCompletionResponse, + status_code=status.HTTP_200_OK, + dependencies=[Depends(verify_api_key)], + ) + async def create_chat_completion(request: ChatCompletionRequest): + if not chat_model.engine.can_generate: + raise HTTPException(status_code=status.HTTP_405_METHOD_NOT_ALLOWED, detail="Not allowed") + + if request.stream: + generate = create_stream_chat_completion_response(request, chat_model) + return EventSourceResponse(generate, media_type="text/event-stream", sep="\n") + else: + return await create_chat_completion_response(request, chat_model) + + @app.post( + "/v1/score/evaluation", + response_model=ScoreEvaluationResponse, + status_code=status.HTTP_200_OK, + dependencies=[Depends(verify_api_key)], + ) + async def create_score_evaluation(request: ScoreEvaluationRequest): + if chat_model.engine.can_generate: + raise HTTPException(status_code=status.HTTP_405_METHOD_NOT_ALLOWED, detail="Not allowed") + + return await create_score_evaluation_response(request, chat_model) + + return app + + +def run_api() -> None: + chat_model = ChatModel() + app = create_app(chat_model) + api_host = os.getenv("API_HOST", "0.0.0.0") + api_port = int(os.getenv("API_PORT", "8000")) + print(f"Visit http://localhost:{api_port}/docs for API document.") + uvicorn.run(app, host=api_host, port=api_port) diff --git a/src/llamafactory/api/chat.py b/src/llamafactory/api/chat.py new file mode 100644 index 0000000000000000000000000000000000000000..06e15251872cb49180587c89c2ed504807de39be --- /dev/null +++ b/src/llamafactory/api/chat.py @@ -0,0 +1,294 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 base64 +import io +import json +import os +import re +import uuid +from collections.abc import AsyncGenerator +from typing import TYPE_CHECKING, Optional + +from ..data import Role as DataRole +from ..extras import logging +from ..extras.constants import AUDIO_PLACEHOLDER, IMAGE_PLACEHOLDER, VIDEO_PLACEHOLDER +from ..extras.misc import is_env_enabled +from ..extras.packages import is_fastapi_available, is_pillow_available, is_requests_available +from .common import check_lfi_path, check_ssrf_url, dictify, jsonify +from .protocol import ( + ChatCompletionMessage, + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionResponseUsage, + ChatCompletionStreamResponse, + ChatCompletionStreamResponseChoice, + Finish, + Function, + FunctionCall, + Role, + ScoreEvaluationResponse, +) + + +if is_fastapi_available(): + from fastapi import HTTPException, status + + +if is_pillow_available(): + from PIL import Image + + +if is_requests_available(): + import requests + + +if TYPE_CHECKING: + from ..chat import ChatModel + from ..data.mm_plugin import AudioInput, ImageInput, VideoInput + from .protocol import ChatCompletionRequest, ScoreEvaluationRequest + + +logger = logging.get_logger(__name__) +ROLE_MAPPING = { + Role.USER: DataRole.USER.value, + Role.ASSISTANT: DataRole.ASSISTANT.value, + Role.SYSTEM: DataRole.SYSTEM.value, + Role.FUNCTION: DataRole.FUNCTION.value, + Role.TOOL: DataRole.OBSERVATION.value, +} + + +def _process_request( + request: "ChatCompletionRequest", +) -> tuple[ + list[dict[str, str]], + Optional[str], + Optional[str], + Optional[list["ImageInput"]], + Optional[list["VideoInput"]], + Optional[list["AudioInput"]], +]: + if is_env_enabled("API_VERBOSE", "1"): + logger.info_rank0(f"==== request ====\n{json.dumps(dictify(request), indent=2, ensure_ascii=False)}") + + if len(request.messages) == 0: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid length") + + if request.messages[0].role == Role.SYSTEM: + content = request.messages.pop(0).content + if isinstance(content, list): + system = content[0].text if content else "" + else: + system = content + else: + system = None + + if len(request.messages) % 2 == 0: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Only supports u/a/u/a/u...") + + input_messages = [] + images, videos, audios = [], [], [] + for i, message in enumerate(request.messages): + if i % 2 == 0 and message.role not in [Role.USER, Role.TOOL]: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid role") + elif i % 2 == 1 and message.role not in [Role.ASSISTANT, Role.FUNCTION]: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid role") + + if message.role == Role.ASSISTANT and isinstance(message.tool_calls, list) and len(message.tool_calls): + tool_calls = [ + {"name": tool_call.function.name, "arguments": tool_call.function.arguments} + for tool_call in message.tool_calls + ] + content = json.dumps(tool_calls, ensure_ascii=False) + input_messages.append({"role": ROLE_MAPPING[Role.FUNCTION], "content": content}) + elif isinstance(message.content, list): + text_content = "" + for input_item in message.content: + if input_item.type == "text": + text_content += input_item.text + elif input_item.type == "image_url": + text_content += IMAGE_PLACEHOLDER + image_url = input_item.image_url.url + if re.match(r"^data:image\/(png|jpg|jpeg|gif|bmp);base64,(.+)$", image_url): # base64 image + image_stream = io.BytesIO(base64.b64decode(image_url.split(",", maxsplit=1)[1])) + elif os.path.isfile(image_url): # local file + check_lfi_path(image_url) + image_stream = open(image_url, "rb") + else: # web uri + check_ssrf_url(image_url) + image_stream = requests.get(image_url, stream=True).raw + + images.append(Image.open(image_stream).convert("RGB")) + elif input_item.type == "video_url": + text_content += VIDEO_PLACEHOLDER + video_url = input_item.video_url.url + if re.match(r"^data:video\/(mp4|mkv|avi|mov);base64,(.+)$", video_url): # base64 video + video_stream = io.BytesIO(base64.b64decode(video_url.split(",", maxsplit=1)[1])) + elif os.path.isfile(video_url): # local file + check_lfi_path(video_url) + video_stream = video_url + else: # web uri + check_ssrf_url(video_url) + video_stream = requests.get(video_url, stream=True).raw + + videos.append(video_stream) + elif input_item.type == "audio_url": + text_content += AUDIO_PLACEHOLDER + audio_url = input_item.audio_url.url + if re.match(r"^data:audio\/(mpeg|mp3|wav|ogg);base64,(.+)$", audio_url): # base64 audio + audio_stream = io.BytesIO(base64.b64decode(audio_url.split(",", maxsplit=1)[1])) + elif os.path.isfile(audio_url): # local file + check_lfi_path(audio_url) + audio_stream = audio_url + else: # web uri + check_ssrf_url(audio_url) + audio_stream = requests.get(audio_url, stream=True).raw + + audios.append(audio_stream) + else: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid input type {input_item.type}." + ) + + input_messages.append({"role": ROLE_MAPPING[message.role], "content": text_content}) + else: + input_messages.append({"role": ROLE_MAPPING[message.role], "content": message.content}) + + tool_list = request.tools + if isinstance(tool_list, list) and len(tool_list): + try: + tools = json.dumps([dictify(tool.function) for tool in tool_list], ensure_ascii=False) + except json.JSONDecodeError: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid tools") + else: + tools = None + + return input_messages, system, tools, images or None, videos or None, audios or None + + +def _create_stream_chat_completion_chunk( + completion_id: str, + model: str, + delta: "ChatCompletionMessage", + index: Optional[int] = 0, + finish_reason: Optional["Finish"] = None, +) -> str: + choice_data = ChatCompletionStreamResponseChoice(index=index, delta=delta, finish_reason=finish_reason) + chunk = ChatCompletionStreamResponse(id=completion_id, model=model, choices=[choice_data]) + return jsonify(chunk) + + +async def create_chat_completion_response( + request: "ChatCompletionRequest", chat_model: "ChatModel" +) -> "ChatCompletionResponse": + completion_id = f"chatcmpl-{uuid.uuid4().hex}" + input_messages, system, tools, images, videos, audios = _process_request(request) + responses = await chat_model.achat( + input_messages, + system, + tools, + images, + videos, + audios, + do_sample=request.do_sample, + temperature=request.temperature, + top_p=request.top_p, + max_new_tokens=request.max_tokens, + num_return_sequences=request.n, + repetition_penalty=request.presence_penalty, + stop=request.stop, + ) + + prompt_length, response_length = 0, 0 + choices = [] + for i, response in enumerate(responses): + if tools: + result = chat_model.engine.template.extract_tool(response.response_text) + else: + result = response.response_text + + if isinstance(result, list): + tool_calls = [] + for tool in result: + function = Function(name=tool.name, arguments=tool.arguments) + tool_calls.append(FunctionCall(id=f"call_{uuid.uuid4().hex}", function=function)) + + response_message = ChatCompletionMessage(role=Role.ASSISTANT, tool_calls=tool_calls) + finish_reason = Finish.TOOL + else: + response_message = ChatCompletionMessage(role=Role.ASSISTANT, content=result) + finish_reason = Finish.STOP if response.finish_reason == "stop" else Finish.LENGTH + + choices.append(ChatCompletionResponseChoice(index=i, message=response_message, finish_reason=finish_reason)) + prompt_length = response.prompt_length + response_length += response.response_length + + usage = ChatCompletionResponseUsage( + prompt_tokens=prompt_length, + completion_tokens=response_length, + total_tokens=prompt_length + response_length, + ) + + return ChatCompletionResponse(id=completion_id, model=request.model, choices=choices, usage=usage) + + +async def create_stream_chat_completion_response( + request: "ChatCompletionRequest", chat_model: "ChatModel" +) -> AsyncGenerator[str, None]: + completion_id = f"chatcmpl-{uuid.uuid4().hex}" + input_messages, system, tools, images, videos, audios = _process_request(request) + if tools: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Cannot stream function calls.") + + if request.n > 1: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Cannot stream multiple responses.") + + yield _create_stream_chat_completion_chunk( + completion_id=completion_id, model=request.model, delta=ChatCompletionMessage(role=Role.ASSISTANT, content="") + ) + async for new_token in chat_model.astream_chat( + input_messages, + system, + tools, + images, + videos, + audios, + do_sample=request.do_sample, + temperature=request.temperature, + top_p=request.top_p, + max_new_tokens=request.max_tokens, + repetition_penalty=request.presence_penalty, + stop=request.stop, + ): + if len(new_token) != 0: + yield _create_stream_chat_completion_chunk( + completion_id=completion_id, model=request.model, delta=ChatCompletionMessage(content=new_token) + ) + + yield _create_stream_chat_completion_chunk( + completion_id=completion_id, model=request.model, delta=ChatCompletionMessage(), finish_reason=Finish.STOP + ) + yield "[DONE]" + + +async def create_score_evaluation_response( + request: "ScoreEvaluationRequest", chat_model: "ChatModel" +) -> "ScoreEvaluationResponse": + score_id = f"scoreval-{uuid.uuid4().hex}" + if len(request.messages) == 0: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid request") + + scores = await chat_model.aget_scores(request.messages, max_length=request.max_length) + return ScoreEvaluationResponse(id=score_id, model=request.model, scores=scores) diff --git a/src/llamafactory/api/common.py b/src/llamafactory/api/common.py new file mode 100644 index 0000000000000000000000000000000000000000..7b4e9602de7ebc10b4f15c68ad9167cb9d80d8ef --- /dev/null +++ b/src/llamafactory/api/common.py @@ -0,0 +1,96 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 ipaddress +import json +import os +import socket +from typing import TYPE_CHECKING, Any +from urllib.parse import urlparse + +from ..extras.misc import is_env_enabled +from ..extras.packages import is_fastapi_available + + +if is_fastapi_available(): + from fastapi import HTTPException, status + + +if TYPE_CHECKING: + from pydantic import BaseModel + + +SAFE_MEDIA_PATH = os.environ.get("SAFE_MEDIA_PATH", os.path.join(os.path.dirname(__file__), "safe_media")) +ALLOW_LOCAL_FILES = is_env_enabled("ALLOW_LOCAL_FILES", "1") + + +def dictify(data: "BaseModel") -> dict[str, Any]: + try: # pydantic v2 + return data.model_dump(exclude_unset=True) + except AttributeError: # pydantic v1 + return data.dict(exclude_unset=True) + + +def jsonify(data: "BaseModel") -> str: + try: # pydantic v2 + return json.dumps(data.model_dump(exclude_unset=True), ensure_ascii=False) + except AttributeError: # pydantic v1 + return data.json(exclude_unset=True, ensure_ascii=False) + + +def check_lfi_path(path: str) -> None: + """Checks if a given path is vulnerable to LFI. Raises HTTPException if unsafe.""" + if not ALLOW_LOCAL_FILES: + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Local file access is disabled.") + + try: + os.makedirs(SAFE_MEDIA_PATH, exist_ok=True) + real_path = os.path.realpath(path) + safe_path = os.path.realpath(SAFE_MEDIA_PATH) + + if not real_path.startswith(safe_path): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail="File access is restricted to the safe media directory." + ) + except Exception: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid or inaccessible file path.") + + +def check_ssrf_url(url: str) -> None: + """Checks if a given URL is vulnerable to SSRF. Raises HTTPException if unsafe.""" + try: + parsed_url = urlparse(url) + if parsed_url.scheme not in ["http", "https"]: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Only HTTP/HTTPS URLs are allowed.") + + hostname = parsed_url.hostname + if not hostname: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid URL hostname.") + + ip_info = socket.getaddrinfo(hostname, parsed_url.port) + ip_address_str = ip_info[0][4][0] + ip = ipaddress.ip_address(ip_address_str) + + if not ip.is_global: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Access to private or reserved IP addresses is not allowed.", + ) + + except socket.gaierror: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=f"Could not resolve hostname: {parsed_url.hostname}" + ) + except Exception as e: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid URL: {e}") diff --git a/src/llamafactory/api/protocol.py b/src/llamafactory/api/protocol.py new file mode 100644 index 0000000000000000000000000000000000000000..19e5279dcbd387cca14c76058a4acf8f5155a221 --- /dev/null +++ b/src/llamafactory/api/protocol.py @@ -0,0 +1,156 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 time +from enum import StrEnum, unique +from typing import Any, Literal + +from pydantic import BaseModel, Field + + +@unique +class Role(StrEnum): + USER = "user" + ASSISTANT = "assistant" + SYSTEM = "system" + FUNCTION = "function" + TOOL = "tool" + + +@unique +class Finish(StrEnum): + STOP = "stop" + LENGTH = "length" + TOOL = "tool_calls" + + +class ModelCard(BaseModel): + id: str + object: Literal["model"] = "model" + created: int = Field(default_factory=lambda: int(time.time())) + owned_by: Literal["owner"] = "owner" + + +class ModelList(BaseModel): + object: Literal["list"] = "list" + data: list[ModelCard] = [] + + +class Function(BaseModel): + name: str + arguments: str + + +class FunctionDefinition(BaseModel): + name: str + description: str + parameters: dict[str, Any] + + +class FunctionAvailable(BaseModel): + type: Literal["function", "code_interpreter"] = "function" + function: FunctionDefinition | None = None + + +class FunctionCall(BaseModel): + id: str + type: Literal["function"] = "function" + function: Function + + +class URL(BaseModel): + url: str + detail: Literal["auto", "low", "high"] = "auto" + + +class MultimodalInputItem(BaseModel): + type: Literal["text", "image_url", "video_url", "audio_url"] + text: str | None = None + image_url: URL | None = None + video_url: URL | None = None + audio_url: URL | None = None + + +class ChatMessage(BaseModel): + role: Role + content: str | list[MultimodalInputItem] | None = None + tool_calls: list[FunctionCall] | None = None + + +class ChatCompletionMessage(BaseModel): + role: Role | None = None + content: str | None = None + tool_calls: list[FunctionCall] | None = None + + +class ChatCompletionRequest(BaseModel): + model: str + messages: list[ChatMessage] + tools: list[FunctionAvailable] | None = None + do_sample: bool | None = None + temperature: float | None = None + top_p: float | None = None + n: int = 1 + presence_penalty: float | None = None + max_tokens: int | None = None + stop: str | list[str] | None = None + stream: bool = False + + +class ChatCompletionResponseChoice(BaseModel): + index: int + message: ChatCompletionMessage + finish_reason: Finish + + +class ChatCompletionStreamResponseChoice(BaseModel): + index: int + delta: ChatCompletionMessage + finish_reason: Finish | None = None + + +class ChatCompletionResponseUsage(BaseModel): + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class ChatCompletionResponse(BaseModel): + id: str + object: Literal["chat.completion"] = "chat.completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: list[ChatCompletionResponseChoice] + usage: ChatCompletionResponseUsage + + +class ChatCompletionStreamResponse(BaseModel): + id: str + object: Literal["chat.completion.chunk"] = "chat.completion.chunk" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: list[ChatCompletionStreamResponseChoice] + + +class ScoreEvaluationRequest(BaseModel): + model: str + messages: list[str] + max_length: int | None = None + + +class ScoreEvaluationResponse(BaseModel): + id: str + object: Literal["score.evaluation"] = "score.evaluation" + model: str + scores: list[float] diff --git a/src/llamafactory/chat/__init__.py b/src/llamafactory/chat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..15d8b9ba2d77d6f300d59300da5a49abd3ed4e57 --- /dev/null +++ b/src/llamafactory/chat/__init__.py @@ -0,0 +1,19 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 .base_engine import BaseEngine +from .chat_model import ChatModel + + +__all__ = ["BaseEngine", "ChatModel"] diff --git a/src/llamafactory/chat/base_engine.py b/src/llamafactory/chat/base_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..6d497c1ae927f94f396c18833b18cdb894cbd59d --- /dev/null +++ b/src/llamafactory/chat/base_engine.py @@ -0,0 +1,98 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 collections.abc import AsyncGenerator +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, Optional, Union + + +if TYPE_CHECKING: + from transformers import PreTrainedModel, PreTrainedTokenizer + from vllm import AsyncLLMEngine + + from ..data import Template + from ..data.mm_plugin import AudioInput, ImageInput, VideoInput + from ..extras.constants import EngineName + from ..hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments + + +@dataclass +class Response: + response_text: str + response_length: int + prompt_length: int + finish_reason: Literal["stop", "length"] + + +class BaseEngine(ABC): + r"""Base class for inference engine of chat models. + + Must implements async methods: chat(), stream_chat() and get_scores(). + """ + + name: "EngineName" + model: Union["PreTrainedModel", "AsyncLLMEngine"] + tokenizer: "PreTrainedTokenizer" + can_generate: bool + template: "Template" + generating_args: dict[str, Any] + + @abstractmethod + def __init__( + self, + model_args: "ModelArguments", + data_args: "DataArguments", + finetuning_args: "FinetuningArguments", + generating_args: "GeneratingArguments", + ) -> None: + r"""Initialize an inference engine.""" + ... + + @abstractmethod + async def chat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> list["Response"]: + r"""Get a list of responses of the chat model.""" + ... + + @abstractmethod + async def stream_chat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> AsyncGenerator[str, None]: + r"""Get the response token-by-token of the chat model.""" + ... + + @abstractmethod + async def get_scores( + self, + batch_input: list[str], + **input_kwargs, + ) -> list[float]: + r"""Get a list of scores of the reward model.""" + ... diff --git a/src/llamafactory/chat/chat_model.py b/src/llamafactory/chat/chat_model.py new file mode 100644 index 0000000000000000000000000000000000000000..9ffd8647c7f5233957be9379f697a4818b3026e0 --- /dev/null +++ b/src/llamafactory/chat/chat_model.py @@ -0,0 +1,200 @@ +# Copyright 2025 THUDM and the LlamaFactory team. +# +# This code is inspired by the THUDM's ChatGLM implementation. +# https://github.com/THUDM/ChatGLM-6B/blob/main/cli_demo.py +# +# 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 asyncio +import os +from collections.abc import AsyncGenerator, Generator +from threading import Thread +from typing import TYPE_CHECKING, Any, Optional + +from ..extras.constants import EngineName +from ..extras.misc import torch_gc +from ..hparams import get_infer_args + + +if TYPE_CHECKING: + from ..data.mm_plugin import AudioInput, ImageInput, VideoInput + from .base_engine import BaseEngine, Response + + +def _start_background_loop(loop: "asyncio.AbstractEventLoop") -> None: + asyncio.set_event_loop(loop) + loop.run_forever() + + +class ChatModel: + r"""General class for chat models. Backed by huggingface or vllm engines. + + Supports both sync and async methods. + Sync methods: chat(), stream_chat() and get_scores(). + Async methods: achat(), astream_chat() and aget_scores(). + """ + + def __init__(self, args: Optional[dict[str, Any]] = None) -> None: + model_args, data_args, finetuning_args, generating_args = get_infer_args(args) + + if model_args.infer_backend == EngineName.HF: + from .hf_engine import HuggingfaceEngine + + self.engine: BaseEngine = HuggingfaceEngine(model_args, data_args, finetuning_args, generating_args) + elif model_args.infer_backend == EngineName.VLLM: + try: + from .vllm_engine import VllmEngine + + self.engine: BaseEngine = VllmEngine(model_args, data_args, finetuning_args, generating_args) + except ImportError as e: + raise ImportError( + "vLLM not install, you may need to run `pip install vllm`\n" + "or try to use HuggingFace backend: --infer_backend huggingface" + ) from e + elif model_args.infer_backend == EngineName.SGLANG: + try: + from .sglang_engine import SGLangEngine + + self.engine: BaseEngine = SGLangEngine(model_args, data_args, finetuning_args, generating_args) + except ImportError as e: + raise ImportError( + "SGLang not install, you may need to run `pip install sglang[all]`\n" + "or try to use HuggingFace backend: --infer_backend huggingface" + ) from e + else: + raise NotImplementedError(f"Unknown backend: {model_args.infer_backend}") + + self._loop = asyncio.new_event_loop() + self._thread = Thread(target=_start_background_loop, args=(self._loop,), daemon=True) + self._thread.start() + + def chat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> list["Response"]: + r"""Get a list of responses of the chat model.""" + task = asyncio.run_coroutine_threadsafe( + self.achat(messages, system, tools, images, videos, audios, **input_kwargs), self._loop + ) + return task.result() + + async def achat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> list["Response"]: + r"""Asynchronously get a list of responses of the chat model.""" + return await self.engine.chat(messages, system, tools, images, videos, audios, **input_kwargs) + + def stream_chat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> Generator[str, None, None]: + r"""Get the response token-by-token of the chat model.""" + generator = self.astream_chat(messages, system, tools, images, videos, audios, **input_kwargs) + while True: + try: + task = asyncio.run_coroutine_threadsafe(generator.__anext__(), self._loop) + yield task.result() + except StopAsyncIteration: + break + + async def astream_chat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> AsyncGenerator[str, None]: + r"""Asynchronously get the response token-by-token of the chat model.""" + async for new_token in self.engine.stream_chat( + messages, system, tools, images, videos, audios, **input_kwargs + ): + yield new_token + + def get_scores( + self, + batch_input: list[str], + **input_kwargs, + ) -> list[float]: + r"""Get a list of scores of the reward model.""" + task = asyncio.run_coroutine_threadsafe(self.aget_scores(batch_input, **input_kwargs), self._loop) + return task.result() + + async def aget_scores( + self, + batch_input: list[str], + **input_kwargs, + ) -> list[float]: + r"""Asynchronously get a list of scores of the reward model.""" + return await self.engine.get_scores(batch_input, **input_kwargs) + + +def run_chat() -> None: + if os.name != "nt": + try: + import readline # noqa: F401 + except ImportError: + print("Install `readline` for a better experience.") + + chat_model = ChatModel() + messages = [] + print("Welcome to the CLI application, use `clear` to remove the history, use `exit` to exit the application.") + + while True: + try: + query = input("\nUser: ") + except UnicodeDecodeError: + print("Detected decoding error at the inputs, please set the terminal encoding to utf-8.") + continue + except Exception: + raise + + if query.strip() == "exit": + break + + if query.strip() == "clear": + messages = [] + torch_gc() + print("History has been removed.") + continue + + messages.append({"role": "user", "content": query}) + print("Assistant: ", end="", flush=True) + + response = "" + for new_text in chat_model.stream_chat(messages): + print(new_text, end="", flush=True) + response += new_text + print() + messages.append({"role": "assistant", "content": response}) diff --git a/src/llamafactory/chat/hf_engine.py b/src/llamafactory/chat/hf_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..a7ef1e2cb6f5d40a0e3f7278660d9080434af3c9 --- /dev/null +++ b/src/llamafactory/chat/hf_engine.py @@ -0,0 +1,423 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 asyncio +import os +from collections.abc import AsyncGenerator, Callable +from threading import Thread +from typing import TYPE_CHECKING, Any, Optional, Union + +import torch +from transformers import GenerationConfig, TextIteratorStreamer, set_seed +from typing_extensions import override + +from ..data import get_template_and_fix_tokenizer +from ..extras import logging +from ..extras.constants import AUDIO_PLACEHOLDER, IMAGE_PLACEHOLDER, VIDEO_PLACEHOLDER, EngineName +from ..model import load_model, load_tokenizer +from .base_engine import BaseEngine, Response + + +if TYPE_CHECKING: + from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin + from trl import PreTrainedModelWrapper + + from ..data import Template + from ..data.mm_plugin import AudioInput, ImageInput, VideoInput + from ..hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments + + +logger = logging.get_logger(__name__) + + +class HuggingfaceEngine(BaseEngine): + def __init__( + self, + model_args: "ModelArguments", + data_args: "DataArguments", + finetuning_args: "FinetuningArguments", + generating_args: "GeneratingArguments", + ) -> None: + self.name = EngineName.HF + self.can_generate = finetuning_args.stage == "sft" + tokenizer_module = load_tokenizer(model_args) + self.tokenizer = tokenizer_module["tokenizer"] + self.processor = tokenizer_module["processor"] + self.tokenizer.padding_side = "left" if self.can_generate else "right" + self.template = get_template_and_fix_tokenizer(self.tokenizer, data_args) + self.model = load_model( + self.tokenizer, model_args, finetuning_args, is_trainable=False, add_valuehead=(not self.can_generate) + ) # must after fixing tokenizer to resize vocab + self.generating_args = generating_args.to_dict() + try: + asyncio.get_event_loop() + except RuntimeError: + logger.warning_rank0_once("There is no current event loop, creating a new one.") + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + self.semaphore = asyncio.Semaphore(int(os.getenv("MAX_CONCURRENT", "1"))) + + @staticmethod + def _process_args( + model: "PreTrainedModel", + tokenizer: "PreTrainedTokenizer", + processor: Optional["ProcessorMixin"], + template: "Template", + generating_args: dict[str, Any], + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + input_kwargs: Optional[dict[str, Any]] = {}, + ) -> tuple[dict[str, Any], int]: + mm_input_dict = {"images": [], "videos": [], "audios": [], "imglens": [0], "vidlens": [0], "audlens": [0]} + if images is not None: + mm_input_dict.update({"images": images, "imglens": [len(images)]}) + if not any(IMAGE_PLACEHOLDER in message["content"] for message in messages): + messages[0]["content"] = IMAGE_PLACEHOLDER * len(images) + messages[0]["content"] + + if videos is not None: + mm_input_dict.update({"videos": videos, "vidlens": [len(videos)]}) + if not any(VIDEO_PLACEHOLDER in message["content"] for message in messages): + messages[0]["content"] = VIDEO_PLACEHOLDER * len(videos) + messages[0]["content"] + + if audios is not None: + mm_input_dict.update({"audios": audios, "audlens": [len(audios)]}) + if not any(AUDIO_PLACEHOLDER in message["content"] for message in messages): + messages[0]["content"] = AUDIO_PLACEHOLDER * len(audios) + messages[0]["content"] + + messages = template.mm_plugin.process_messages( + messages, mm_input_dict["images"], mm_input_dict["videos"], mm_input_dict["audios"], processor + ) + paired_messages = messages + [{"role": "assistant", "content": ""}] + prompt_ids, _ = template.encode_oneturn(tokenizer, paired_messages, system, tools) + prompt_ids, _ = template.mm_plugin.process_token_ids( + prompt_ids, + None, + mm_input_dict["images"], + mm_input_dict["videos"], + mm_input_dict["audios"], + tokenizer, + processor, + ) + prompt_length = len(prompt_ids) + inputs = torch.tensor([prompt_ids], device=model.device) + attention_mask = torch.ones_like(inputs, dtype=torch.long) + + do_sample: Optional[bool] = input_kwargs.pop("do_sample", None) + temperature: Optional[float] = input_kwargs.pop("temperature", None) + top_p: Optional[float] = input_kwargs.pop("top_p", None) + top_k: Optional[float] = input_kwargs.pop("top_k", None) + num_return_sequences: int = input_kwargs.pop("num_return_sequences", 1) + repetition_penalty: Optional[float] = input_kwargs.pop("repetition_penalty", None) + length_penalty: Optional[float] = input_kwargs.pop("length_penalty", None) + skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None) + max_length: Optional[int] = input_kwargs.pop("max_length", None) + max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None) + seed: Optional[int] = input_kwargs.pop("seed", None) + stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None) + + if stop is not None: + logger.warning_rank0("Stop parameter is not supported by the huggingface engine yet.") + + generating_args = generating_args.copy() + generating_args.update( + dict( + do_sample=do_sample if do_sample is not None else generating_args["do_sample"], + temperature=temperature if temperature is not None else generating_args["temperature"], + top_p=top_p if top_p is not None else generating_args["top_p"], + top_k=top_k if top_k is not None else generating_args["top_k"], + num_return_sequences=num_return_sequences, + repetition_penalty=repetition_penalty + if repetition_penalty is not None + else generating_args["repetition_penalty"], + length_penalty=length_penalty if length_penalty is not None else generating_args["length_penalty"], + skip_special_tokens=skip_special_tokens + if skip_special_tokens is not None + else generating_args["skip_special_tokens"], + eos_token_id=template.get_stop_token_ids(tokenizer), + pad_token_id=tokenizer.pad_token_id, + ) + ) + + if isinstance(num_return_sequences, int) and num_return_sequences > 1: # do_sample needs temperature > 0 + generating_args["do_sample"] = True + generating_args["temperature"] = generating_args["temperature"] or 1.0 + + if not generating_args["temperature"]: + generating_args["do_sample"] = False + + if not generating_args["do_sample"]: + generating_args.pop("temperature", None) + generating_args.pop("top_p", None) + + if max_length: + generating_args.pop("max_new_tokens", None) + generating_args["max_length"] = max_length + + if max_new_tokens: + generating_args.pop("max_length", None) + generating_args["max_new_tokens"] = max_new_tokens + + gen_kwargs = dict( + inputs=inputs, + attention_mask=attention_mask, + generation_config=GenerationConfig(**generating_args), + ) + if seed is not None: + gen_kwargs["_seed"] = seed + + mm_inputs = template.mm_plugin.get_mm_inputs(**mm_input_dict, batch_ids=[prompt_ids], processor=processor) + for key, value in mm_inputs.items(): + if isinstance(value, list) and isinstance(value[0], torch.Tensor): # for pixtral inputs + value = torch.stack(value) # assume they have same sizes + elif ( + isinstance(value, list) and isinstance(value[0], list) and isinstance(value[0][0], torch.Tensor) + ): # for minicpmv inputs + value = torch.stack([torch.stack(v) for v in value]) + elif not isinstance(value, torch.Tensor): + value = torch.tensor(value) + + if torch.is_floating_point(value): # cast data dtype for paligemma + value = value.to(model.dtype) + + if key == "second_per_grid_ts": # qwen2.5vl special case + gen_kwargs[key] = value.tolist() + else: + gen_kwargs[key] = value.to(model.device) + + if getattr(model.config, "model_type", None) in ["minicpmv", "minicpmo"]: + gen_kwargs["input_ids"] = inputs + gen_kwargs["tokenizer"] = tokenizer + if "audio_feature_lens" in mm_inputs: + gen_kwargs["audio_feature_lens"] = mm_inputs["audio_feature_lens"] + + gen_kwargs.pop("image_sizes", None) + + return gen_kwargs, prompt_length + + @staticmethod + @torch.inference_mode() + def _chat( + model: "PreTrainedModel", + tokenizer: "PreTrainedTokenizer", + processor: Optional["ProcessorMixin"], + template: "Template", + generating_args: dict[str, Any], + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + input_kwargs: Optional[dict[str, Any]] = {}, + ) -> list["Response"]: + gen_kwargs, prompt_length = HuggingfaceEngine._process_args( + model, + tokenizer, + processor, + template, + generating_args, + messages, + system, + tools, + images, + videos, + audios, + input_kwargs, + ) + seed = gen_kwargs.pop("_seed", None) + if seed is not None: + set_seed(seed) + + generate_output = model.generate(**gen_kwargs) + if isinstance(generate_output, tuple): + generate_output = generate_output[1][0] # post-process the minicpm_o output + + response_ids = generate_output[:, prompt_length:] + response = tokenizer.batch_decode( + response_ids, + skip_special_tokens=getattr(gen_kwargs["generation_config"], "skip_special_tokens", True), + clean_up_tokenization_spaces=True, + ) + results = [] + for i in range(len(response)): + eos_index = (response_ids[i] == tokenizer.eos_token_id).nonzero() + response_length = (eos_index[0].item() + 1) if len(eos_index) else len(response_ids[i]) + results.append( + Response( + response_text=response[i], + response_length=response_length, + prompt_length=prompt_length, + finish_reason="stop" if len(eos_index) else "length", + ) + ) + + return results + + @staticmethod + @torch.inference_mode() + def _stream_chat( + model: "PreTrainedModel", + tokenizer: "PreTrainedTokenizer", + processor: Optional["ProcessorMixin"], + template: "Template", + generating_args: dict[str, Any], + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + input_kwargs: Optional[dict[str, Any]] = {}, + ) -> Callable[[], str]: + gen_kwargs, _ = HuggingfaceEngine._process_args( + model, + tokenizer, + processor, + template, + generating_args, + messages, + system, + tools, + images, + videos, + audios, + input_kwargs, + ) + seed = gen_kwargs.pop("_seed", None) + if seed is not None: + set_seed(seed) + + streamer = TextIteratorStreamer( + tokenizer, + skip_prompt=True, + skip_special_tokens=getattr(gen_kwargs["generation_config"], "skip_special_tokens", True), + ) + gen_kwargs["streamer"] = streamer + thread = Thread(target=model.generate, kwargs=gen_kwargs, daemon=True) + thread.start() + + def stream(): + try: + return streamer.__next__() + except StopIteration: + raise StopAsyncIteration() + + return stream + + @staticmethod + @torch.inference_mode() + def _get_scores( + model: "PreTrainedModelWrapper", + tokenizer: "PreTrainedTokenizer", + batch_input: list[str], + input_kwargs: Optional[dict[str, Any]] = {}, + ) -> list[float]: + max_length: Optional[int] = input_kwargs.pop("max_length", None) + device = getattr(model.pretrained_model, "device", "cuda") + inputs: dict[str, torch.Tensor] = tokenizer( + batch_input, + padding=True, + truncation=True, + max_length=max_length or getattr(model.config, "max_position_embeddings", 1024), + return_tensors="pt", + add_special_tokens=False, + ).to(device) + values: torch.Tensor = model(**inputs, return_dict=True, use_cache=False)[-1] + scores = values.gather(dim=-1, index=(inputs["attention_mask"].sum(dim=-1, keepdim=True) - 1)) + return scores + + @override + async def chat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> list["Response"]: + if not self.can_generate: + raise ValueError("The current model does not support `chat`.") + + input_args = ( + self.model, + self.tokenizer, + self.processor, + self.template, + self.generating_args, + messages, + system, + tools, + images, + videos, + audios, + input_kwargs, + ) + async with self.semaphore: + return await asyncio.to_thread(self._chat, *input_args) + + @override + async def stream_chat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> AsyncGenerator[str, None]: + if not self.can_generate: + raise ValueError("The current model does not support `stream_chat`.") + + input_args = ( + self.model, + self.tokenizer, + self.processor, + self.template, + self.generating_args, + messages, + system, + tools, + images, + videos, + audios, + input_kwargs, + ) + async with self.semaphore: + stream = self._stream_chat(*input_args) + while True: + try: + yield await asyncio.to_thread(stream) + except StopAsyncIteration: + break + + @override + async def get_scores( + self, + batch_input: list[str], + **input_kwargs, + ) -> list[float]: + if self.can_generate: + raise ValueError("Cannot get scores using an auto-regressive model.") + + input_args = (self.model, self.tokenizer, batch_input, input_kwargs) + async with self.semaphore: + return await asyncio.to_thread(self._get_scores, *input_args) diff --git a/src/llamafactory/chat/sglang_engine.py b/src/llamafactory/chat/sglang_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..342dea45544f0e6b259ceae0b41c27aec3dd8553 --- /dev/null +++ b/src/llamafactory/chat/sglang_engine.py @@ -0,0 +1,292 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 asyncio +import atexit +import json +from collections.abc import AsyncGenerator, AsyncIterator, Sequence +from typing import TYPE_CHECKING, Any, Optional, Union + +import requests +from typing_extensions import override + +from ..data import get_template_and_fix_tokenizer +from ..extras import logging +from ..extras.constants import AUDIO_PLACEHOLDER, IMAGE_PLACEHOLDER, VIDEO_PLACEHOLDER, EngineName +from ..extras.misc import get_device_count, torch_gc +from ..extras.packages import is_sglang_available +from ..hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments +from ..model import load_config, load_tokenizer +from ..model.model_utils.quantization import QuantizationMethod +from .base_engine import BaseEngine, Response + + +if is_sglang_available(): + from sglang.utils import launch_server_cmd, terminate_process, wait_for_server # type: ignore + + +if TYPE_CHECKING: + from ..data.mm_plugin import AudioInput, ImageInput, VideoInput + + +logger = logging.get_logger(__name__) + + +class SGLangEngine(BaseEngine): + """Inference engine for SGLang models. + + This class wraps the SGLang engine to provide a consistent interface for text generation + that matches LLaMA Factory's requirements. It uses the SGLang HTTP server approach for + better interaction and performance. The engine launches a server process and communicates + with it via HTTP requests. + + For more details on the SGLang HTTP server approach, see: + https://docs.sglang.ai/backend/send_request.html + """ + + def __init__( + self, + model_args: "ModelArguments", + data_args: "DataArguments", + finetuning_args: "FinetuningArguments", + generating_args: "GeneratingArguments", + ) -> None: + self.name = EngineName.SGLANG + self.model_args = model_args + config = load_config(model_args) # may download model from ms hub + if getattr(config, "quantization_config", None): # gptq models should use float16 + quantization_config: dict[str, Any] = getattr(config, "quantization_config", None) + quant_method = quantization_config.get("quant_method", "") + if quant_method == QuantizationMethod.GPTQ and model_args.infer_dtype == "auto": + model_args.infer_dtype = "float16" + + self.can_generate = finetuning_args.stage == "sft" + tokenizer_module = load_tokenizer(model_args) + self.tokenizer = tokenizer_module["tokenizer"] + self.processor = tokenizer_module["processor"] + self.tokenizer.padding_side = "left" + self.template = get_template_and_fix_tokenizer(self.tokenizer, data_args) + self.template.mm_plugin.expand_mm_tokens = False # for sglang generate + self.generating_args = generating_args.to_dict() + if model_args.adapter_name_or_path is not None: + self.lora_request = True + else: + self.lora_request = False + + launch_cmd = [ + "python3 -m sglang.launch_server", + f"--model-path {model_args.model_name_or_path}", + f"--dtype {model_args.infer_dtype}", + f"--context-length {model_args.sglang_maxlen}", + f"--mem-fraction-static {model_args.sglang_mem_fraction}", + f"--tp-size {model_args.sglang_tp_size if model_args.sglang_tp_size != -1 else get_device_count() or 1}", + f"--download-dir {model_args.cache_dir}", + "--log-level error", + ] + if self.lora_request: + launch_cmd.extend( + [ + "--max-loras-per-batch 1", + f"--lora-backend {model_args.sglang_lora_backend}", + f"--lora-paths lora0={model_args.adapter_name_or_path[0]}", + "--disable-radix-cache", + ] + ) + launch_cmd = " ".join(launch_cmd) + logger.info_rank0(f"Starting SGLang server with command: {launch_cmd}") + try: + torch_gc() + self.server_process, port = launch_server_cmd(launch_cmd) + self.base_url = f"http://localhost:{port}" + atexit.register(self._cleanup_server) + + logger.info_rank0(f"Waiting for SGLang server to be ready at {self.base_url}") + wait_for_server(self.base_url, timeout=300) + logger.info_rank0(f"SGLang server initialized successfully at {self.base_url}") + try: + response = requests.get(f"{self.base_url}/get_model_info", timeout=5) + if response.status_code == 200: + model_info = response.json() + logger.info(f"SGLang server model info: {model_info}") + except Exception as e: + logger.debug(f"Note: could not get model info: {str(e)}") + + except Exception as e: + logger.error(f"Failed to start SGLang server: {str(e)}") + self._cleanup_server() # make sure to clean up any started process + raise RuntimeError(f"SGLang server initialization failed: {str(e)}.") + + def _cleanup_server(self): + r"""Clean up the server process when the engine is destroyed.""" + if hasattr(self, "server_process") and self.server_process: + try: + logger.info("Terminating SGLang server process") + terminate_process(self.server_process) + logger.info("SGLang server process terminated") + except Exception as e: + logger.warning(f"Error terminating SGLang server: {str(e)}") + + async def _generate( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> AsyncIterator[dict[str, Any]]: + if images is not None and not any(IMAGE_PLACEHOLDER in message["content"] for message in messages): + messages[0]["content"] = IMAGE_PLACEHOLDER * len(images) + messages[0]["content"] + + if videos is not None and not any(VIDEO_PLACEHOLDER in message["content"] for message in messages): + messages[0]["content"] = VIDEO_PLACEHOLDER * len(videos) + messages[0]["content"] + + if audios is not None and not any(AUDIO_PLACEHOLDER in message["content"] for message in messages): + messages[0]["content"] = AUDIO_PLACEHOLDER * len(audios) + messages[0]["content"] + + messages = self.template.mm_plugin.process_messages( + messages, images or [], videos or [], audios or [], self.processor + ) + paired_messages = messages + [{"role": "assistant", "content": ""}] + prompt_ids, _ = self.template.encode_oneturn(self.tokenizer, paired_messages, system, tools) + prompt_length = len(prompt_ids) + + temperature: Optional[float] = input_kwargs.pop("temperature", None) + top_p: Optional[float] = input_kwargs.pop("top_p", None) + top_k: Optional[float] = input_kwargs.pop("top_k", None) + num_return_sequences: int = input_kwargs.pop("num_return_sequences", 1) + repetition_penalty: Optional[float] = input_kwargs.pop("repetition_penalty", None) + skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None) + max_length: Optional[int] = input_kwargs.pop("max_length", None) + max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None) + seed: Optional[int] = input_kwargs.pop("seed", None) + stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None) + + if num_return_sequences != 1: + raise NotImplementedError("SGLang only supports n=1.") + + if "max_new_tokens" in self.generating_args: + max_tokens = self.generating_args["max_new_tokens"] + elif "max_length" in self.generating_args: + if self.generating_args["max_length"] > prompt_length: + max_tokens = self.generating_args["max_length"] - prompt_length + else: + max_tokens = 1 + + if max_length: + max_tokens = max_length - prompt_length if max_length > prompt_length else 1 + + if max_new_tokens: + max_tokens = max_new_tokens + + sampling_params = { + "temperature": temperature if temperature is not None else self.generating_args["temperature"], + "top_p": (top_p if top_p is not None else self.generating_args["top_p"]) or 1.0, # top_p must > 0 + "top_k": (top_k if top_k is not None else self.generating_args["top_k"]) or -1, # top_k must > 0 + "stop": stop, + "stop_token_ids": self.template.get_stop_token_ids(self.tokenizer), + "max_new_tokens": max_tokens, + "repetition_penalty": ( + repetition_penalty if repetition_penalty is not None else self.generating_args["repetition_penalty"] + ) + or 1.0, # repetition_penalty must > 0 + "skip_special_tokens": skip_special_tokens + if skip_special_tokens is not None + else self.generating_args["skip_special_tokens"], + } + if seed is not None: + sampling_params["seed"] = seed + + def stream_request(): + json_data = { + "input_ids": prompt_ids, + "sampling_params": sampling_params, + "stream": True, + } + if self.lora_request: + json_data["lora_request"] = ["lora0"] + response = requests.post(f"{self.base_url}/generate", json=json_data, stream=True) + if response.status_code != 200: + raise RuntimeError(f"SGLang server error: {response.status_code}, {response.text}") + + for chunk in response.iter_lines(decode_unicode=False): + chunk = str(chunk.decode("utf-8")) + if chunk == "data: [DONE]": + break + + if chunk and chunk.startswith("data:"): + yield json.loads(chunk[5:].strip("\n")) + + return await asyncio.to_thread(stream_request) + + @override + async def chat( + self, + messages: Sequence[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[Sequence["ImageInput"]] = None, + videos: Optional[Sequence["VideoInput"]] = None, + audios: Optional[Sequence["AudioInput"]] = None, + **input_kwargs, + ) -> list["Response"]: + final_output = None + generator = await self._generate(messages, system, tools, images, videos, audios, **input_kwargs) + for request_output in generator: + final_output = request_output + + results = [ + Response( + response_text=final_output["text"], + response_length=final_output["meta_info"]["completion_tokens"], + prompt_length=final_output["meta_info"]["prompt_tokens"], + finish_reason="stop" if final_output["meta_info"]["finish_reason"] == "stop" else "length", + ) + ] + return results + + @override + async def stream_chat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> AsyncGenerator[str, None]: + generated_text = "" + generator = await self._generate(messages, system, tools, images, videos, audios, **input_kwargs) + for result in generator: + delta_text = result["text"][len(generated_text) :] + generated_text = result["text"] + yield delta_text + + @override + async def get_scores( + self, + batch_input: list[str], + **input_kwargs, + ) -> list[float]: + raise NotImplementedError("SGLang engine does not support `get_scores`.") + + def __del__(self): + r"""Ensure server is cleaned up when object is deleted.""" + self._cleanup_server() + try: + atexit.unregister(self._cleanup_server) + except Exception: + pass diff --git a/src/llamafactory/chat/vllm_engine.py b/src/llamafactory/chat/vllm_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..5e2b625a396160dc9c9bec1018157da716af912b --- /dev/null +++ b/src/llamafactory/chat/vllm_engine.py @@ -0,0 +1,273 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 uuid +from collections.abc import AsyncGenerator, AsyncIterator +from typing import TYPE_CHECKING, Any, Optional, Union + +from packaging import version +from typing_extensions import override + +from ..data import get_template_and_fix_tokenizer +from ..extras import logging +from ..extras.constants import AUDIO_PLACEHOLDER, IMAGE_PLACEHOLDER, VIDEO_PLACEHOLDER, EngineName +from ..extras.misc import get_device_count +from ..extras.packages import is_vllm_available +from ..model import load_config, load_tokenizer +from ..model.model_utils.quantization import QuantizationMethod +from ..model.model_utils.visual import LlavaMultiModalProjectorForYiVLForVLLM +from .base_engine import BaseEngine, Response + + +if is_vllm_available(): + from vllm import AsyncEngineArgs, AsyncLLMEngine, RequestOutput, SamplingParams + from vllm.lora.request import LoRARequest + + +if TYPE_CHECKING: + from ..data.mm_plugin import AudioInput, ImageInput, VideoInput + from ..hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments + + +logger = logging.get_logger(__name__) + + +class VllmEngine(BaseEngine): + def __init__( + self, + model_args: "ModelArguments", + data_args: "DataArguments", + finetuning_args: "FinetuningArguments", + generating_args: "GeneratingArguments", + ) -> None: + self.name = EngineName.VLLM + self.model_args = model_args + config = load_config(model_args) # may download model from ms hub + if getattr(config, "quantization_config", None): # gptq models should use float16 + quantization_config: dict[str, Any] = getattr(config, "quantization_config", None) + quant_method = quantization_config.get("quant_method", "") + if quant_method == QuantizationMethod.GPTQ and model_args.infer_dtype == "auto": + model_args.infer_dtype = "float16" + + self.can_generate = finetuning_args.stage == "sft" + tokenizer_module = load_tokenizer(model_args) + self.tokenizer = tokenizer_module["tokenizer"] + self.processor = tokenizer_module["processor"] + self.tokenizer.padding_side = "left" + self.template = get_template_and_fix_tokenizer(self.tokenizer, data_args) + self.template.mm_plugin.expand_mm_tokens = False # for vllm generate + self.generating_args = generating_args.to_dict() + + engine_args = { + "model": model_args.model_name_or_path, + "trust_remote_code": model_args.trust_remote_code, + "download_dir": model_args.cache_dir, + "dtype": model_args.infer_dtype, + "max_model_len": model_args.vllm_maxlen, + "tensor_parallel_size": get_device_count() or 1, + "gpu_memory_utilization": model_args.vllm_gpu_util, + "disable_log_stats": True, + "enforce_eager": model_args.vllm_enforce_eager, + "enable_lora": model_args.adapter_name_or_path is not None, + "max_lora_rank": model_args.vllm_max_lora_rank, + } + + import vllm + + if version.parse(vllm.__version__) <= version.parse("0.10.0"): + engine_args["disable_log_requests"] = True + else: + engine_args["enable_log_requests"] = False + + if self.template.mm_plugin.__class__.__name__ != "BasePlugin": + engine_args["limit_mm_per_prompt"] = {"image": 4, "video": 2, "audio": 2} + + if isinstance(model_args.vllm_config, dict): + engine_args.update(model_args.vllm_config) + + if getattr(config, "is_yi_vl_derived_model", None): + import vllm.model_executor.models.llava + + logger.info_rank0("Detected Yi-VL model, applying projector patch.") + vllm.model_executor.models.llava.LlavaMultiModalProjector = LlavaMultiModalProjectorForYiVLForVLLM + + self.model = AsyncLLMEngine.from_engine_args(AsyncEngineArgs(**engine_args)) + if model_args.adapter_name_or_path is not None: + self.lora_request = LoRARequest("default", 1, model_args.adapter_name_or_path[0]) + else: + self.lora_request = None + + async def _generate( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> AsyncIterator["RequestOutput"]: + request_id = f"chatcmpl-{uuid.uuid4().hex}" + if images is not None and not any(IMAGE_PLACEHOLDER in message["content"] for message in messages): + messages[0]["content"] = IMAGE_PLACEHOLDER * len(images) + messages[0]["content"] + + if videos is not None and not any(VIDEO_PLACEHOLDER in message["content"] for message in messages): + messages[0]["content"] = VIDEO_PLACEHOLDER * len(videos) + messages[0]["content"] + + if audios is not None and not any(AUDIO_PLACEHOLDER in message["content"] for message in messages): + messages[0]["content"] = AUDIO_PLACEHOLDER * len(audios) + messages[0]["content"] + + messages = self.template.mm_plugin.process_messages( + messages, images or [], videos or [], audios or [], self.processor + ) + paired_messages = messages + [{"role": "assistant", "content": ""}] + prompt_ids, _ = self.template.encode_oneturn(self.tokenizer, paired_messages, system, tools) + prompt_length = len(prompt_ids) + + temperature: Optional[float] = input_kwargs.pop("temperature", None) + top_p: Optional[float] = input_kwargs.pop("top_p", None) + top_k: Optional[float] = input_kwargs.pop("top_k", None) + num_return_sequences: int = input_kwargs.pop("num_return_sequences", 1) + repetition_penalty: Optional[float] = input_kwargs.pop("repetition_penalty", None) + length_penalty: Optional[float] = input_kwargs.pop("length_penalty", None) + skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None) + max_length: Optional[int] = input_kwargs.pop("max_length", None) + max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None) + seed: Optional[int] = input_kwargs.pop("seed", None) + stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None) + + if length_penalty is not None: + logger.warning_rank0("Length penalty is not supported by the vllm engine yet.") + + if "max_new_tokens" in self.generating_args: + max_tokens = self.generating_args["max_new_tokens"] + elif "max_length" in self.generating_args: + if self.generating_args["max_length"] > prompt_length: + max_tokens = self.generating_args["max_length"] - prompt_length + else: + max_tokens = 1 + + if max_length: + max_tokens = max_length - prompt_length if max_length > prompt_length else 1 + + if max_new_tokens: + max_tokens = max_new_tokens + + sampling_kwargs = dict( + n=num_return_sequences, + repetition_penalty=( + repetition_penalty if repetition_penalty is not None else self.generating_args["repetition_penalty"] + ) + or 1.0, # repetition_penalty must > 0 + temperature=temperature if temperature is not None else self.generating_args["temperature"], + top_p=(top_p if top_p is not None else self.generating_args["top_p"]) or 1.0, # top_p must > 0 + top_k=(top_k if top_k is not None else self.generating_args["top_k"]) or -1, # top_k must > 0 + stop=stop, + stop_token_ids=self.template.get_stop_token_ids(self.tokenizer), + max_tokens=max_tokens, + skip_special_tokens=skip_special_tokens + if skip_special_tokens is not None + else self.generating_args["skip_special_tokens"], + ) + if seed is not None: + sampling_kwargs["seed"] = seed + + sampling_params = SamplingParams(**sampling_kwargs) + + multi_modal_data = {} + if images is not None: # add image features + multi_modal_data["image"] = self.template.mm_plugin._regularize_images( + images, + image_max_pixels=self.model_args.image_max_pixels, + image_min_pixels=self.model_args.image_min_pixels, + )["images"] + + if videos is not None: + multi_modal_data["video"] = self.template.mm_plugin._regularize_videos( + videos, + image_max_pixels=self.model_args.video_max_pixels, + image_min_pixels=self.model_args.video_min_pixels, + video_fps=self.model_args.video_fps, + video_maxlen=self.model_args.video_maxlen, + )["videos"] + + if audios is not None: + audio_data = self.template.mm_plugin._regularize_audios( + audios, + sampling_rate=self.model_args.audio_sampling_rate, + ) + multi_modal_data["audio"] = zip(audio_data["audios"], audio_data["sampling_rates"]) + + result_generator = self.model.generate( + {"prompt_token_ids": prompt_ids, "multi_modal_data": multi_modal_data or None}, + sampling_params=sampling_params, + request_id=request_id, + lora_request=self.lora_request, + ) + return result_generator + + @override + async def chat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> list["Response"]: + final_output = None + generator = await self._generate(messages, system, tools, images, videos, audios, **input_kwargs) + async for request_output in generator: + final_output = request_output + + results = [] + for output in final_output.outputs: + results.append( + Response( + response_text=output.text, + response_length=len(output.token_ids), + prompt_length=len(final_output.prompt_token_ids), + finish_reason=output.finish_reason, + ) + ) + + return results + + @override + async def stream_chat( + self, + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + images: Optional[list["ImageInput"]] = None, + videos: Optional[list["VideoInput"]] = None, + audios: Optional[list["AudioInput"]] = None, + **input_kwargs, + ) -> AsyncGenerator[str, None]: + generated_text = "" + generator = await self._generate(messages, system, tools, images, videos, audios, **input_kwargs) + async for result in generator: + delta_text = result.outputs[0].text[len(generated_text) :] + generated_text = result.outputs[0].text + yield delta_text + + @override + async def get_scores( + self, + batch_input: list[str], + **input_kwargs, + ) -> list[float]: + raise NotImplementedError("vLLM engine does not support `get_scores`.") diff --git a/src/llamafactory/cli.py b/src/llamafactory/cli.py new file mode 100644 index 0000000000000000000000000000000000000000..d574bf1db543f5379f074e276898826234708037 --- /dev/null +++ b/src/llamafactory/cli.py @@ -0,0 +1,31 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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. + + +def main(): + from .extras.misc import is_env_enabled + + if is_env_enabled("USE_V1"): + from .v1 import launcher + else: + from . import launcher + + launcher.launch() + + +if __name__ == "__main__": + from multiprocessing import freeze_support + + freeze_support() + main() diff --git a/src/llamafactory/data/__init__.py b/src/llamafactory/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..11c8c9fcecd10e736e240196fde98f833c9df3dc --- /dev/null +++ b/src/llamafactory/data/__init__.py @@ -0,0 +1,37 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 .collator import ( + KTODataCollatorWithPadding, + MultiModalDataCollatorForSeq2Seq, + PairwiseDataCollatorWithPadding, + SFTDataCollatorWith4DAttentionMask, +) +from .data_utils import Role, split_dataset +from .loader import get_dataset +from .template import TEMPLATES, Template, get_template_and_fix_tokenizer + + +__all__ = [ + "TEMPLATES", + "KTODataCollatorWithPadding", + "MultiModalDataCollatorForSeq2Seq", + "PairwiseDataCollatorWithPadding", + "Role", + "SFTDataCollatorWith4DAttentionMask", + "Template", + "get_dataset", + "get_template_and_fix_tokenizer", + "split_dataset", +] diff --git a/src/llamafactory/data/__pycache__/__init__.cpython-312.pyc b/src/llamafactory/data/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2adfe8d9154abda81ad8f2fb50a5f20dbb8f3f1e Binary files /dev/null and b/src/llamafactory/data/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/llamafactory/data/__pycache__/collator.cpython-312.pyc b/src/llamafactory/data/__pycache__/collator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a770a1f8a5f07496073182e6bacd964c34335d5d Binary files /dev/null and b/src/llamafactory/data/__pycache__/collator.cpython-312.pyc differ diff --git a/src/llamafactory/data/__pycache__/converter.cpython-312.pyc b/src/llamafactory/data/__pycache__/converter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4d25513f59fb49caeaa517c8d90e51778a91a7d8 Binary files /dev/null and b/src/llamafactory/data/__pycache__/converter.cpython-312.pyc differ diff --git a/src/llamafactory/data/__pycache__/data_utils.cpython-312.pyc b/src/llamafactory/data/__pycache__/data_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..724b3cf0752033cb682f4965487fe507bf8fbfae Binary files /dev/null and b/src/llamafactory/data/__pycache__/data_utils.cpython-312.pyc differ diff --git a/src/llamafactory/data/__pycache__/formatter.cpython-312.pyc b/src/llamafactory/data/__pycache__/formatter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..137741245ed49c2be6bb839e3091d317c3dc17b1 Binary files /dev/null and b/src/llamafactory/data/__pycache__/formatter.cpython-312.pyc differ diff --git a/src/llamafactory/data/__pycache__/loader.cpython-312.pyc b/src/llamafactory/data/__pycache__/loader.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad262c644949309a4fb3257306a8a81f4be6806f Binary files /dev/null and b/src/llamafactory/data/__pycache__/loader.cpython-312.pyc differ diff --git a/src/llamafactory/data/__pycache__/mm_plugin.cpython-312.pyc b/src/llamafactory/data/__pycache__/mm_plugin.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..930b1478026f60a72fce644b2040baf4e989de2a --- /dev/null +++ b/src/llamafactory/data/__pycache__/mm_plugin.cpython-312.pyc @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c41a6a5d07168e7ea8e94f99c087bbaf48c1ea00e8ad95c8e343884be0e4d3a5 +size 125555 diff --git a/src/llamafactory/data/__pycache__/parser.cpython-312.pyc b/src/llamafactory/data/__pycache__/parser.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..437fd584a785d22bbd5caada2b8f90ffc7292c3d Binary files /dev/null and b/src/llamafactory/data/__pycache__/parser.cpython-312.pyc differ diff --git a/src/llamafactory/data/__pycache__/template.cpython-312.pyc b/src/llamafactory/data/__pycache__/template.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c4246c93e178437f246d0e7d1cf4fec7d14a8b7c Binary files /dev/null and b/src/llamafactory/data/__pycache__/template.cpython-312.pyc differ diff --git a/src/llamafactory/data/__pycache__/tool_utils.cpython-312.pyc b/src/llamafactory/data/__pycache__/tool_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ef9d728eb1b0e0b3c42dc2eca46c9715dee4ba5 Binary files /dev/null and b/src/llamafactory/data/__pycache__/tool_utils.cpython-312.pyc differ diff --git a/src/llamafactory/data/collator.py b/src/llamafactory/data/collator.py new file mode 100644 index 0000000000000000000000000000000000000000..af234d99bfe6917b53f6b995b1097fa182c726dc --- /dev/null +++ b/src/llamafactory/data/collator.py @@ -0,0 +1,608 @@ +# Copyright 2025 OpenAccess AI Collective and the LlamaFactory team. +# +# This code is inspired by the OpenAccess AI Collective's axolotl library. +# https://github.com/OpenAccess-AI-Collective/axolotl/blob/main/src/axolotl/monkeypatch/utils.py +# +# 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 copy +import inspect +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, Optional + +import numpy as np +import torch +import torch.nn.functional as F +from peft import PeftModel +from transformers import DataCollatorForSeq2Seq + +from ..extras.constants import AUDIO_PLACEHOLDER, IGNORE_INDEX, IMAGE_PLACEHOLDER, MROPE_MODELS +from ..extras.packages import is_pillow_available + + +if is_pillow_available(): + from PIL import Image + + +if TYPE_CHECKING: + from transformers import ProcessorMixin + + from .template import Template + + +def _slice_mm_inputs_for_sample( + mm_inputs: dict[str, Any], + batch_imglens: list[int], + batch_vidlens: list[int], + batch_idx: int, + images_per_subseq: Optional[list[int]] = None, + videos_per_subseq: Optional[list[int]] = None, + subseq_idx: Optional[int] = None, +) -> dict[str, Any]: + r"""Slice mm_inputs for one batch sample, optionally for a single sub-sequence when packing. + + image_grid_thw / video_grid_thw have shape [num_items, 3]. Indices for sample batch_idx + are batch_imglens[batch_idx] images and batch_vidlens[batch_idx] videos. When subseq_idx + is given, further restrict to that sub-seq's counts via packed_*_counts. + has_dummy_image=True means only batch[0] will be concated with fake image and no multimodal data. + """ + image_start_idx = sum(batch_imglens[:batch_idx]) + image_end_idx = sum(batch_imglens[: batch_idx + 1]) + video_start_idx = sum(batch_vidlens[:batch_idx]) + video_end_idx = sum(batch_vidlens[: batch_idx + 1]) + + if subseq_idx is not None and images_per_subseq is not None: + image_start_idx += sum(images_per_subseq[:subseq_idx]) + image_end_idx = image_start_idx + images_per_subseq[subseq_idx] + + if subseq_idx is not None and videos_per_subseq is not None: + video_start_idx += sum(videos_per_subseq[:subseq_idx]) + video_end_idx = video_start_idx + videos_per_subseq[subseq_idx] + + sliced_mm_inputs: dict[str, Any] = {} + key_to_slice_meta = { + "image_grid_thw": (image_start_idx, image_end_idx, True), + "video_grid_thw": (video_start_idx, video_end_idx, True), + "second_per_grid_ts": (video_start_idx, video_end_idx, False), # qwen2.5vl + "video_second_per_grid": (video_start_idx, video_end_idx, False), # qwen omni + } + + for key, (start_idx, end_idx, assign_none_when_empty) in key_to_slice_meta.items(): + if key not in mm_inputs: + continue + + mm_value = mm_inputs[key] + if mm_value is not None and end_idx > start_idx: + sliced_mm_inputs[key] = mm_value[start_idx:end_idx] + elif assign_none_when_empty: + sliced_mm_inputs[key] = None + + return sliced_mm_inputs + + +def prepare_4d_attention_mask(attention_mask_with_indices: "torch.Tensor", dtype: "torch.dtype") -> "torch.Tensor": + r"""Expand 2d attention mask to 4d attention mask. + + Expand the attention mask with indices from (batch_size, seq_len) to (batch_size, 1, seq_len, seq_len), + handle packed sequences and transforms the mask to lower triangular form to prevent future peeking. + + e.g. + ```python + # input + [[1, 1, 2, 2, 2, 0]] + # output + [ + [ + [ + [o, x, x, x, x, x], + [o, o, x, x, x, x], + [x, x, o, x, x, x], + [x, x, o, o, x, x], + [x, x, o, o, o, x], + [x, x, x, x, x, x], + ] + ] + ] + ``` + where `o` equals to `0.0`, `x` equals to `min_dtype`. + """ + _, seq_len = attention_mask_with_indices.size() + min_dtype = torch.finfo(dtype).min + zero_tensor = torch.tensor(0, dtype=dtype) + + # Create a non-padding mask. + non_padding_mask = (attention_mask_with_indices != 0).unsqueeze(1).unsqueeze(2) + # Create indices for comparison. + indices = attention_mask_with_indices.unsqueeze(1).unsqueeze(2) # [bsz, 1, 1, seq_len] + indices_t = attention_mask_with_indices.unsqueeze(1).unsqueeze(3) # [bsz, 1, seq_len, 1] + # Create a lower triangular mask. + tril_mask = torch.tril(torch.ones((seq_len, seq_len), dtype=torch.bool)) + attention_mask_4d = (indices == indices_t) & non_padding_mask & tril_mask + # Invert the attention mask. + attention_mask_4d = torch.where(attention_mask_4d, zero_tensor, min_dtype) + return attention_mask_4d + + +@dataclass +class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq): + r"""Data collator that supports VLMs. + + Features should contain input_ids, attention_mask, labels, and optionally contain images, videos and audios. + """ + + template: Optional["Template"] = None + processor: Optional["ProcessorMixin"] = None + + def __post_init__(self): + if self.template is None: + raise ValueError("Template is required for MultiModalDataCollator.") + + if isinstance(self.model, PeftModel): + self.model = self.model.base_model.model + + if self.model is not None and hasattr(self.model, "get_rope_index"): # for qwen2vl mrope + self.get_rope_func = self.model.get_rope_index # transformers < 4.52.0 or qwen2.5 omni + elif self.model is not None and hasattr(self.model, "model") and hasattr(self.model.model, "get_rope_index"): + self.get_rope_func = self.model.model.get_rope_index # transformers >= 4.52.0 + else: + self.get_rope_func = None + + def _compute_rope_position_ids(self, features: dict[str, "torch.Tensor"], mm_inputs: dict[str, Any]) -> None: + r"""Compute position_ids and rope_deltas via get_rope_func for VLMs.""" + rope_index_kwargs = { + "input_ids": features["input_ids"], + "image_grid_thw": mm_inputs.get("image_grid_thw"), + "video_grid_thw": mm_inputs.get("video_grid_thw"), + "attention_mask": (features["attention_mask"] >= 1).float(), + } + if features["attention_mask"].sum() == 0: # for pad tokens + seq_len = features["input_ids"].shape[-1] + features["position_ids"] = ( + torch.arange(seq_len).view(1, 1, seq_len).expand(3, *features["input_ids"].shape).contiguous() + ) + features["rope_deltas"] = torch.zeros(features["input_ids"].shape[0]) + return + + if "mm_token_type_ids" in inspect.signature(self.get_rope_func).parameters: + image_token_id = getattr(self.model.config, "image_token_id", None) + video_token_id = getattr(self.model.config, "video_token_id", None) + if image_token_id is not None or video_token_id is not None: + mm_token_type_ids = torch.zeros_like(features["input_ids"]) + if image_token_id is not None: + mm_token_type_ids[features["input_ids"] == image_token_id] = 1 + if video_token_id is not None: + mm_token_type_ids[features["input_ids"] == video_token_id] = 2 + rope_index_kwargs["mm_token_type_ids"] = mm_token_type_ids + + if "second_per_grid_ts" in mm_inputs: # for qwen2vl + rope_index_kwargs["second_per_grid_ts"] = mm_inputs.get("second_per_grid_ts") + elif "video_second_per_grid" in mm_inputs: # for qwen2.5 omni + rope_index_kwargs["second_per_grids"] = mm_inputs.get("video_second_per_grid") + + if getattr(self.model.config, "model_type", None) in ["qwen2_5_omni_thinker", "qwen3_omni_moe_thinker"]: + rope_index_kwargs["use_audio_in_video"] = getattr(self.processor, "use_audio_in_video", False) + feature_attention_mask = mm_inputs.get("feature_attention_mask", None) + if feature_attention_mask is not None: # FIXME: need to get video image lengths + audio_feature_lengths = torch.sum(feature_attention_mask, dim=1) + rope_index_kwargs["audio_seqlens"] = audio_feature_lengths # prepare for input + + features["position_ids"], rope_deltas = self.get_rope_func(**rope_index_kwargs) + features["rope_deltas"] = rope_deltas - (1 - rope_index_kwargs["attention_mask"]).sum(dim=-1).unsqueeze(-1) + else: # for qwen vl + features["position_ids"], features["rope_deltas"] = self.get_rope_func(**rope_index_kwargs) + + def _compute_rope_position_ids_with_packing( + self, + features: dict[str, "torch.Tensor"], + mm_inputs: dict[str, Any], + packing_params_list: list[dict[str, Any] | None], + batch_imglens: list[int], + batch_vidlens: list[int], + batch_audlens: list[int], + has_dummy_image: bool, + ) -> None: + r"""Compute position_ids and rope_deltas per sample (or per sub-sequence when packed), then merge and validate.""" + bsz = features["input_ids"].size(0) + seq_len = features["input_ids"].size(1) + all_position_ids: list[torch.Tensor] = [] + all_rope_deltas: list[torch.Tensor] = [] + + if has_dummy_image: + # for [0, seq_len] = [0, unpadded_length + right_padding_length + fake_input_ids_len + collator_padding_length] + # FIXME: maybe right_padding_length is large, with improper max_cutoff_len + unpadded_length = int(features["attention_mask"][0].bool().sum().item()) + right_padding_length = int((packing_params_list[0] or {}).get("right_padding_length") or 0) + fake_input_padding_length = max(0, seq_len - unpadded_length - right_padding_length) + # avoid continual cuseqlens breaking varlen attention @kuangdd + # https://github.com/hiyouga/LlamaFactory/issues/10452 + dummy_image_right_padding_mrope = ( + torch.arange(fake_input_padding_length) + .view(1, 1, fake_input_padding_length) + .expand(3, bsz, fake_input_padding_length) + ) + dummy_image_right_padding_attention_mask = torch.zeros((bsz, fake_input_padding_length)) + assert self.tokenizer.padding_side == "right", "padding_side should be right when fake image is injected" + dummy_mm_inputs = copy.deepcopy(mm_inputs) + + for sample_idx in range(bsz): + sample_packing = (packing_params_list[sample_idx] or {}) if sample_idx < len(packing_params_list) else {} + sequence_boundaries = sample_packing.get("sequence_boundaries") + num_sub_seqs = ( + (len(sequence_boundaries) - 1) if sequence_boundaries and len(sequence_boundaries) > 1 else 1 + ) + image_subseq_ids = sample_packing.get("image_subseq_ids") or [] + video_subseq_ids = sample_packing.get("video_subseq_ids") or [] + images_per_subseq = ( + [image_subseq_ids.count(i) for i in range(num_sub_seqs)] + if image_subseq_ids and num_sub_seqs > 1 + else None + ) + videos_per_subseq = ( + [video_subseq_ids.count(i) for i in range(num_sub_seqs)] + if video_subseq_ids and num_sub_seqs > 1 + else None + ) + if has_dummy_image: + mm_inputs = {} + + if num_sub_seqs <= 1: + sample_features = { + "input_ids": features["input_ids"], + "attention_mask": features["attention_mask"][sample_idx : sample_idx + 1], + } + mm_inputs_for_sample = _slice_mm_inputs_for_sample( + mm_inputs, batch_imglens, batch_vidlens, sample_idx=sample_idx + ) + self._compute_rope_position_ids(sample_features, mm_inputs_for_sample) + all_position_ids.append(sample_features["position_ids"]) + all_rope_deltas.append(sample_features["rope_deltas"]) + else: + # when we do packing, don't need rope_deltas when training. + sample_position_ids: list[torch.Tensor] = [] + for subseq_idx in range(num_sub_seqs): + subseq_start = sequence_boundaries[subseq_idx] + subseq_end = sequence_boundaries[subseq_idx + 1] + subseq_features = { + "input_ids": features["input_ids"][sample_idx : sample_idx + 1, subseq_start:subseq_end], + "attention_mask": features["attention_mask"][ + sample_idx : sample_idx + 1, subseq_start:subseq_end + ], + } + mm_inputs_for_subseq = _slice_mm_inputs_for_sample( + mm_inputs, + batch_imglens, + batch_vidlens, + sample_idx, + images_per_subseq, + videos_per_subseq, + subseq_idx, + ) + self._compute_rope_position_ids(subseq_features, mm_inputs_for_subseq) + sample_position_ids.append(subseq_features["position_ids"]) + + all_position_ids.append(torch.cat(sample_position_ids, dim=-1)) + + batch_dim_for_position_ids = 1 if all_position_ids[0].dim() == 3 else 0 + + features["position_ids"] = torch.cat(all_position_ids, dim=batch_dim_for_position_ids) + if has_dummy_image: + mm_inputs = dummy_mm_inputs + + expected_position_ids_shape = ( + (bsz, seq_len) + if all_position_ids[0].dim() == 2 + else ( + all_position_ids[0].size(0), + bsz, + seq_len, + ) + ) + # Check if position_ids shape matches expected shape. + # for further usage, we should padding to the right when some padding token on the right. + if has_dummy_image: + features["position_ids"] = torch.cat([features["position_ids"], dummy_image_right_padding_mrope], dim=-1) + features["attention_mask"] = torch.cat( + [features["attention_mask"], dummy_image_right_padding_attention_mask], dim=-1 + ) + + if features["position_ids"].shape != expected_position_ids_shape: + raise ValueError( + "Merged position_ids shape mismatch: " + f"got {features['position_ids'].shape}, expected {expected_position_ids_shape}." + ) + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]: + batch_images, batch_videos, batch_audios = [], [], [] + batch_imglens, batch_vidlens, batch_audlens, batch_input_ids = [], [], [], [] + packing_params_list: list[dict[str, Any] | None] = [] + for feature in features: + images = feature.pop("images", None) or [] + videos = feature.pop("videos", None) or [] + audios = feature.pop("audios", None) or [] + batch_images.extend(images) + batch_videos.extend(videos) + batch_audios.extend(audios) + batch_imglens.append(len(images)) + batch_vidlens.append(len(videos)) + batch_audlens.append(len(audios)) + batch_input_ids.append(feature["input_ids"]) + packing_params_list.append(feature.pop("packing_params", None)) + + fake_input_ids = [] + has_dummy_image = False + if ( + self.template.mm_plugin.image_token is not None and sum(batch_imglens) == 0 and sum(batch_vidlens) == 0 + ): # avoid process hanging in zero3/fsdp case + fake_messages = [{"role": "user", "content": IMAGE_PLACEHOLDER}] + fake_images = [Image.new("RGB", (64, 64), (255, 255, 255))] + fake_messages = self.template.mm_plugin.process_messages( + fake_messages, fake_images, [], [], self.processor + ) + _fake_input_ids = self.tokenizer.encode(fake_messages[0]["content"], add_special_tokens=False) + _fake_input_ids, _ = self.template.mm_plugin.process_token_ids( + _fake_input_ids, None, fake_images, [], [], self.tokenizer, self.processor + ) + fake_input_ids.extend(_fake_input_ids) + batch_images = fake_images + batch_imglens[0] = 1 + has_dummy_image = True + + if ( + self.template.mm_plugin.audio_token is not None and sum(batch_audlens) == 0 + ): # avoid process hanging in zero3/fsdp case + fake_messages = [{"role": "user", "content": AUDIO_PLACEHOLDER}] + fake_audios = [np.zeros(1600)] + fake_messages = self.template.mm_plugin.process_messages( + fake_messages, [], [], fake_audios, self.processor + ) + _fake_input_ids = self.tokenizer.encode(fake_messages[0]["content"], add_special_tokens=False) + _fake_input_ids, _ = self.template.mm_plugin.process_token_ids( + _fake_input_ids, None, [], [], fake_audios, self.tokenizer, self.processor + ) + fake_input_ids.extend(_fake_input_ids) + batch_audios = fake_audios + batch_audlens[0] = 1 + + if len(fake_input_ids) != 0: + if self.tokenizer.padding_side == "right": + features[0]["input_ids"] = features[0]["input_ids"] + fake_input_ids + features[0]["attention_mask"] = features[0]["attention_mask"] + [0] * len(fake_input_ids) + features[0]["labels"] = features[0]["labels"] + [IGNORE_INDEX] * len(fake_input_ids) + else: + features[0]["input_ids"] = fake_input_ids + features[0]["input_ids"] + features[0]["attention_mask"] = [0] * len(fake_input_ids) + features[0]["attention_mask"] + features[0]["labels"] = [IGNORE_INDEX] * len(fake_input_ids) + features[0]["labels"] + + batch_input_ids[0] = features[0]["input_ids"] + + mm_inputs = self.template.mm_plugin.get_mm_inputs( + batch_images, + batch_videos, + batch_audios, + batch_imglens, + batch_vidlens, + batch_audlens, + batch_input_ids, + self.processor, + ) + if "token_type_ids" in mm_inputs: + token_type_ids = mm_inputs.pop("token_type_ids") + for i, feature in enumerate(features): + feature["token_type_ids"] = token_type_ids[i] + + if "mm_token_type_ids" in mm_inputs: # need tensor-like for gemma4 + mm_token_type_ids = mm_inputs.pop("mm_token_type_ids") + max_len = max(len(ids) for ids in mm_token_type_ids) + padded = [] + for ids in mm_token_type_ids: + pad_len = max_len - len(ids) + if self.tokenizer.padding_side == "right": + padded.append(ids + [0] * pad_len) + else: + padded.append([0] * pad_len + ids) + + mm_inputs["mm_token_type_ids"] = torch.tensor(padded, dtype=torch.long) + + features: dict[str, torch.Tensor] = super().__call__(features) + + bsz, seq_len = features["input_ids"].shape[:2] + model_type = getattr(self.model.config, "model_type", None) if self.model is not None else None + is_omni = model_type in [ + "qwen2_5_omni_thinker", + "qwen3_omni_moe_thinker", + ] + + if self.get_rope_func is not None: + # for mmrope situation, we should calculate position_ids and rope_deltas per sample. + # When neat_packing is on, each sample has packing_params; None means no packing for that sample. + boundaries_list = [p.get("sequence_boundaries") if p is not None else None for p in packing_params_list] + has_packing = any(b is not None and len(b) > 2 for b in boundaries_list) + if has_dummy_image and has_packing: + # FIXME: too tricky, need to be refactored @kuangdd + features["has_dummy_image"] = True + + # When fake image/audio was injected, sequence_boundaries no longer match the tensor; use non-packing path. + if not has_packing: + self._compute_rope_position_ids(features, mm_inputs) + else: + if is_omni: # TODO: support omni models for packed sequences @kuangdd + raise RuntimeError("Omni models are not supported for packed sequences for now.") + + self._compute_rope_position_ids_with_packing( + features, + mm_inputs, + packing_params_list, + batch_imglens, + batch_vidlens, + batch_audlens, + has_dummy_image, + ) + + # For transformers compatibility, after https://github.com/huggingface/transformers/issues/39400 + if features["position_ids"].dim() == 3: + features["position_ids"] = torch.cat( + [features["position_ids"][0].unsqueeze(0), features["position_ids"]], dim=0 + ) + + if ( + self.model is not None + and getattr(self.model.config, "model_type", None) in MROPE_MODELS + and ("position_ids" not in features or features["position_ids"].dim() != 3) + ): + raise ValueError(f"{self.model.config.model_type} requires 3D position ids for mrope.") + + if "cross_attention_mask" in mm_inputs: # for mllama inputs when pad_to_multiple_of is enabled + cross_attention_mask = mm_inputs.pop("cross_attention_mask") + seq_len = features["input_ids"].size(1) + orig_len = cross_attention_mask.size(1) + mm_inputs["cross_attention_mask"] = F.pad(cross_attention_mask, (0, 0, 0, 0, 0, seq_len - orig_len)) + + features.update(mm_inputs) + + if "image_bound" in features: # for minicpmv inputs + bsz, seq_length = features["input_ids"].shape + features["position_ids"] = torch.arange(seq_length).long().repeat(bsz, 1) + return {"data": features, "input_ids": features["input_ids"], "labels": features["labels"]} + + return features + + +@dataclass +class SFTDataCollatorWith4DAttentionMask(MultiModalDataCollatorForSeq2Seq): + r"""Data collator for 4d attention mask.""" + + block_diag_attn: bool = False + attn_implementation: Literal["eager", "sdpa", "flash_attention_2"] = "eager" + compute_dtype: "torch.dtype" = torch.float32 + neat_packing: bool = False + + def __post_init__(self): + super().__post_init__() + if self.neat_packing and self.attn_implementation == "flash_attention_2": + if self.model is not None and getattr(self.model.config, "model_type", None) in ["gemma4", "gpt_oss"]: + raise ValueError("Neat packing is not supported for gemma4, gpt_oss models for now.") + + @staticmethod + def _unpad_packed_features(features: dict[str, Any]) -> None: + r"""Trim padded positions for packed FA2 batches.""" + attention_mask = features.get("attention_mask") + if not torch.is_tensor(attention_mask) or attention_mask.dim() != 2 or attention_mask.size(0) != 1: + return + + seq_len = attention_mask.size(1) + non_padding_indices = torch.nonzero(attention_mask[0] != 0, as_tuple=False).flatten() + if non_padding_indices.numel() == seq_len: + return + + keys_on_seq_dim_1 = {"input_ids", "labels", "attention_mask", "token_type_ids"} + for key, value in list(features.items()): + if not torch.is_tensor(value): + continue + + if key == "position_ids" and value.size(-1) == seq_len: + features[key] = value.index_select(-1, non_padding_indices) + elif ( + key == "cross_attention_mask" and value.dim() >= 2 and value.size(0) == 1 and value.size(1) == seq_len + ): + features[key] = value.index_select(1, non_padding_indices) + elif key in keys_on_seq_dim_1 and value.dim() == 2 and value.size(0) == 1 and value.size(1) == seq_len: + features[key] = value.index_select(1, non_padding_indices) + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]: + features = super().__call__(features) + has_dummy_image = features.pop("has_dummy_image", False) + if self.block_diag_attn and self.attn_implementation != "flash_attention_2": + features["attention_mask"] = prepare_4d_attention_mask(features["attention_mask"], self.compute_dtype) + + if self.neat_packing and self.attn_implementation == "flash_attention_2": # FIXME compatibility fa3/fa4 + assert features["input_ids"].shape[0] == 1, "bsz should be 1 for neat packing" + if not has_dummy_image: + self._unpad_packed_features(features) + + features["attention_mask"] = None # let transformers handle causal packed mask. + + for key, value in features.items(): # cast data dtype for paligemma + if torch.is_tensor(value) and torch.is_floating_point(value): + features[key] = value.to(self.compute_dtype) + + return features + + +@dataclass +class PairwiseDataCollatorWithPadding(MultiModalDataCollatorForSeq2Seq): + r"""Data collator for pairwise data.""" + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]: + r"""Pad batched data to the longest sequence in the batch. + + We generate 2 * n examples where the first n examples represent chosen examples and + the last n examples represent rejected examples. + """ + concatenated_features = [] + for key in ("chosen", "rejected"): + for feature in features: + target_feature = { + "input_ids": feature[f"{key}_input_ids"], + "attention_mask": feature[f"{key}_attention_mask"], + "labels": feature[f"{key}_labels"], + "images": feature["images"], + "videos": feature["videos"], + "audios": feature["audios"], + } + concatenated_features.append(target_feature) + + return super().__call__(concatenated_features) + + +@dataclass +class KTODataCollatorWithPadding(MultiModalDataCollatorForSeq2Seq): + r"""Data collator for KTO data.""" + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]: + target_features = [] + kl_features = [] + kto_tags = [] + for feature in features: + target_feature = { + "input_ids": feature["input_ids"], + "attention_mask": feature["attention_mask"], + "labels": feature["labels"], + "images": feature["images"], + "videos": feature["videos"], + "audios": feature["audios"], + } + kl_feature = { + "input_ids": feature["kl_input_ids"], + "attention_mask": feature["kl_attention_mask"], + "labels": feature["kl_labels"], + "images": feature["images"], + "videos": feature["videos"], + "audios": feature["audios"], + } + target_features.append(target_feature) + kl_features.append(kl_feature) + kto_tags.append(feature["kto_tags"]) + + batch = super().__call__(target_features) + kl_batch = super().__call__(kl_features) + batch["kl_input_ids"] = kl_batch["input_ids"] + batch["kl_attention_mask"] = kl_batch["attention_mask"] + batch["kl_labels"] = kl_batch["labels"] + if "cross_attention_mask" in kl_batch: # for mllama inputs + batch["kl_cross_attention_mask"] = kl_batch["cross_attention_mask"] + + if "token_type_ids" in kl_batch: + batch["kl_token_type_ids"] = kl_batch["token_type_ids"] + + batch["kto_tags"] = torch.tensor(kto_tags) + return batch diff --git a/src/llamafactory/data/converter.py b/src/llamafactory/data/converter.py new file mode 100644 index 0000000000000000000000000000000000000000..ad49deded441318e722cd07bfe331feff6c498a5 --- /dev/null +++ b/src/llamafactory/data/converter.py @@ -0,0 +1,425 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 abc import abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Union + +from ..extras import logging +from .data_utils import Role + + +if TYPE_CHECKING: + from datasets import Dataset, IterableDataset + from transformers import Seq2SeqTrainingArguments + + from ..hparams import DataArguments + from .mm_plugin import AudioInput, ImageInput, VideoInput + from .parser import DatasetAttr + + MediaType = Union[ImageInput, VideoInput, AudioInput] + + +logger = logging.get_logger(__name__) + + +@dataclass +class DatasetConverter: + dataset_attr: "DatasetAttr" + data_args: "DataArguments" + + def _find_medias(self, medias: Union["MediaType", list["MediaType"], None]) -> list["MediaType"] | None: + r"""Optionally concatenate media path to media dir when loading from local disk.""" + if medias is None: + return None + elif not isinstance(medias, list): + medias = [medias] + elif len(medias) == 0: + return None + else: + medias = medias[:] + + if self.dataset_attr.load_from in ["script", "file"]: + if isinstance(medias[0], str): + for i in range(len(medias)): + media_path = os.path.join(self.data_args.media_dir, medias[i]) + if os.path.isfile(media_path): + medias[i] = media_path + else: + logger.warning_rank0_once( + f"Media {medias[i]} does not exist in `media_dir`. Use original path." + ) + elif isinstance(medias[0], list): # for processed video frames + # medias is a list of lists, e.g., [[frame1.jpg, frame2.jpg], [frame3.jpg, frame4.jpg]] + for i in range(len(medias)): + for j in range(len(medias[i])): + media_path = os.path.join(self.data_args.media_dir, medias[i][j]) + if os.path.isfile(media_path): + medias[i][j] = media_path + else: + logger.warning_rank0_once( + f"Media {medias[i][j]} does not exist in `media_dir`. Use original path." + ) + + return medias + + @abstractmethod + def __call__(self, example: dict[str, Any]) -> dict[str, Any]: + r"""Convert a single example in the dataset to the standard format.""" + ... + + +@dataclass +class AlpacaDatasetConverter(DatasetConverter): + def __call__(self, example: dict[str, Any]) -> dict[str, Any]: + prompt = [] + if self.dataset_attr.history and isinstance(example[self.dataset_attr.history], list): + for old_prompt, old_response in example[self.dataset_attr.history]: + prompt.append({"role": Role.USER.value, "content": old_prompt}) + prompt.append({"role": Role.ASSISTANT.value, "content": old_response}) + + query = [] + if self.dataset_attr.prompt and example[self.dataset_attr.prompt]: + query.append(example[self.dataset_attr.prompt]) + + if self.dataset_attr.query and example[self.dataset_attr.query]: + query.append(example[self.dataset_attr.query]) + + prompt.append({"role": Role.USER.value, "content": "\n".join(query)}) # "prompt\nquery" + + if self.dataset_attr.kto_tag and isinstance(example[self.dataset_attr.kto_tag], bool): # kto example + response = [{"role": Role.ASSISTANT.value, "content": example[self.dataset_attr.response]}] + if example[self.dataset_attr.kto_tag]: + response = response + [{"role": Role.ASSISTANT.value, "content": ""}] + else: + response = [{"role": Role.ASSISTANT.value, "content": ""}] + response + elif ( + self.dataset_attr.ranking + and isinstance(example[self.dataset_attr.chosen], str) + and isinstance(example[self.dataset_attr.rejected], str) + ): # pairwise example + response = [ + {"role": Role.ASSISTANT.value, "content": example[self.dataset_attr.chosen]}, + {"role": Role.ASSISTANT.value, "content": example[self.dataset_attr.rejected]}, + ] + elif self.dataset_attr.response and isinstance(example[self.dataset_attr.response], str): # normal example + response = [{"role": Role.ASSISTANT.value, "content": example[self.dataset_attr.response]}] + else: # unsupervised + response = [] + + output = { + "_prompt": prompt, + "_response": response, + "_system": example[self.dataset_attr.system] if self.dataset_attr.system else "", + "_tools": example[self.dataset_attr.tools] if self.dataset_attr.tools else "", + "_images": self._find_medias(example[self.dataset_attr.images]) if self.dataset_attr.images else None, + "_videos": self._find_medias(example[self.dataset_attr.videos]) if self.dataset_attr.videos else None, + "_audios": self._find_medias(example[self.dataset_attr.audios]) if self.dataset_attr.audios else None, + } + return output + + +@dataclass +class SharegptDatasetConverter(DatasetConverter): + def __call__(self, example: dict[str, Any]) -> dict[str, Any]: + tag_mapping = { + self.dataset_attr.user_tag: Role.USER.value, + self.dataset_attr.assistant_tag: Role.ASSISTANT.value, + self.dataset_attr.observation_tag: Role.OBSERVATION.value, + self.dataset_attr.function_tag: Role.FUNCTION.value, + self.dataset_attr.system_tag: Role.SYSTEM.value, + } + odd_tags = (self.dataset_attr.user_tag, self.dataset_attr.observation_tag) + even_tags = (self.dataset_attr.assistant_tag, self.dataset_attr.function_tag) + accept_tags = (odd_tags, even_tags) + messages = example[self.dataset_attr.messages] + if ( + self.dataset_attr.system_tag + and len(messages) != 0 + and messages[0][self.dataset_attr.role_tag] == self.dataset_attr.system_tag + ): + system = messages[0][self.dataset_attr.content_tag] + messages = messages[1:] + else: + system = example[self.dataset_attr.system] if self.dataset_attr.system else "" + + aligned_messages = [] + broken_data = False + for turn_idx, message in enumerate(messages): + if message[self.dataset_attr.role_tag] not in accept_tags[turn_idx % 2]: + logger.warning_rank0(f"Invalid role tag in {messages}.") + broken_data = True + break + + aligned_messages.append( + { + "role": tag_mapping[message[self.dataset_attr.role_tag]], + "content": message[self.dataset_attr.content_tag], + } + ) + + if (not self.dataset_attr.ranking and len(aligned_messages) % 2 != 0) or ( + self.dataset_attr.ranking and len(aligned_messages) % 2 == 0 + ): + logger.warning_rank0(f"Invalid message count in {messages}.") + broken_data = True + + if broken_data: + logger.warning_rank0("Skipping this abnormal example.") + prompt, response = [], [] + elif self.dataset_attr.kto_tag and isinstance(example[self.dataset_attr.kto_tag], bool): # kto example + prompt = aligned_messages[:-1] + response = aligned_messages[-1:] + if example[self.dataset_attr.kto_tag]: + response = response + [{"role": Role.ASSISTANT.value, "content": ""}] + else: + response = [{"role": Role.ASSISTANT.value, "content": ""}] + response + elif ( + self.dataset_attr.ranking + and isinstance(example[self.dataset_attr.chosen], dict) + and isinstance(example[self.dataset_attr.rejected], dict) + ): # pairwise example + chosen = example[self.dataset_attr.chosen] + rejected = example[self.dataset_attr.rejected] + if ( + chosen[self.dataset_attr.role_tag] not in accept_tags[-1] + or rejected[self.dataset_attr.role_tag] not in accept_tags[-1] + ): + logger.warning_rank0(f"Invalid role tag in {[chosen, rejected]}.") + broken_data = True + + prompt = aligned_messages + response = [ + { + "role": tag_mapping[chosen[self.dataset_attr.role_tag]], + "content": chosen[self.dataset_attr.content_tag], + }, + { + "role": tag_mapping[rejected[self.dataset_attr.role_tag]], + "content": rejected[self.dataset_attr.content_tag], + }, + ] + else: # normal example + prompt = aligned_messages[:-1] + response = aligned_messages[-1:] + + output = { + "_prompt": prompt, + "_response": response, + "_system": system, + "_tools": example[self.dataset_attr.tools] if self.dataset_attr.tools else "", + "_images": self._find_medias(example[self.dataset_attr.images]) if self.dataset_attr.images else None, + "_videos": self._find_medias(example[self.dataset_attr.videos]) if self.dataset_attr.videos else None, + "_audios": self._find_medias(example[self.dataset_attr.audios]) if self.dataset_attr.audios else None, + } + return output + + +@dataclass +class OpenAIDatasetConverter(DatasetConverter): + def __call__(self, example: dict[str, Any]) -> dict[str, Any]: + tag_mapping = { + self.dataset_attr.user_tag: Role.USER.value, + self.dataset_attr.assistant_tag: Role.ASSISTANT.value, + self.dataset_attr.observation_tag: Role.OBSERVATION.value, + self.dataset_attr.function_tag: Role.FUNCTION.value, + self.dataset_attr.system_tag: Role.SYSTEM.value, + } + + messages = example[self.dataset_attr.messages] + if ( + self.dataset_attr.system_tag + and len(messages) != 0 + and messages[0][self.dataset_attr.role_tag] == self.dataset_attr.system_tag + ): + system = messages[0][self.dataset_attr.content_tag] + messages = messages[1:] + else: + system = example.get(self.dataset_attr.system, "") if self.dataset_attr.system else "" + + aligned_messages = [] + tool_responses = [] + broken_data = False + for turn_idx, message in enumerate(messages): + role = message[self.dataset_attr.role_tag] + content = message[self.dataset_attr.content_tag] + + if role in [self.dataset_attr.assistant_tag, self.dataset_attr.function_tag]: + if tool_calls := message.get("tool_calls"): + tool_calls_list = [tool["function"] for tool in tool_calls] + content = json.dumps(tool_calls_list, ensure_ascii=False) + role = self.dataset_attr.function_tag + + if role == self.dataset_attr.observation_tag: + tool_responses.append(content) + continue + elif len(tool_responses) > 0: + _content = "\n\n\n".join(tool_responses) + aligned_messages.append( + { + "role": Role.OBSERVATION.value, + "content": _content, + } + ) + tool_responses = [] + + aligned_messages.append( + { + "role": tag_mapping[role], + "content": content, + } + ) + + odd_tags = (Role.USER.value, Role.OBSERVATION.value) + even_tags = (Role.ASSISTANT.value, Role.FUNCTION.value) + accept_tags = (odd_tags, even_tags) + for turn_idx, message in enumerate(aligned_messages): + if message["role"] not in accept_tags[turn_idx % 2]: + logger.warning_rank0(f"Invalid role tag in {messages}.") + broken_data = True + break + + if (not self.dataset_attr.ranking and len(aligned_messages) % 2 != 0) or ( + self.dataset_attr.ranking and len(aligned_messages) % 2 == 0 + ): + logger.warning_rank0(f"Invalid message count in {messages}.") + broken_data = True + + if broken_data: + logger.warning_rank0("Skipping this abnormal example.") + prompt, response = [], [] + elif self.dataset_attr.kto_tag and isinstance(example[self.dataset_attr.kto_tag], bool): # kto example + prompt = aligned_messages[:-1] + response = aligned_messages[-1:] + if example[self.dataset_attr.kto_tag]: + response = response + [{"role": Role.ASSISTANT.value, "content": ""}] + else: + response = [{"role": Role.ASSISTANT.value, "content": ""}] + response + elif ( + self.dataset_attr.ranking + and isinstance(example[self.dataset_attr.chosen], dict) + and isinstance(example[self.dataset_attr.rejected], dict) + ): # pairwise example + chosen = example[self.dataset_attr.chosen] + rejected = example[self.dataset_attr.rejected] + if ( + chosen[self.dataset_attr.role_tag] not in accept_tags[-1] + or rejected[self.dataset_attr.role_tag] not in accept_tags[-1] + ): + logger.warning_rank0(f"Invalid role tag in {[chosen, rejected]}.") + broken_data = True + + prompt = aligned_messages + response = [ + { + "role": tag_mapping[chosen[self.dataset_attr.role_tag]], + "content": chosen[self.dataset_attr.content_tag], + }, + { + "role": tag_mapping[rejected[self.dataset_attr.role_tag]], + "content": rejected[self.dataset_attr.content_tag], + }, + ] + else: # normal example + prompt = aligned_messages[:-1] + response = aligned_messages[-1:] + + tools = example.get(self.dataset_attr.tools, "") if self.dataset_attr.tools else "" + if isinstance(tools, dict) or isinstance(tools, list): + tools = json.dumps(tools, ensure_ascii=False) + + short_system_prompt = "detailed thinking off" + if not system: + if not tools: + system = short_system_prompt + else: + pass + else: + if not tools: + if "detailed thinking on" in system or "detailed thinking off" in system: + pass + else: + system += "\n" + short_system_prompt + else: + system += "\n" + + output = { + "_prompt": prompt, + "_response": response, + "_system": system, + "_tools": tools, + "_images": self._find_medias(example[self.dataset_attr.images]) if self.dataset_attr.images else None, + "_videos": self._find_medias(example[self.dataset_attr.videos]) if self.dataset_attr.videos else None, + "_audios": self._find_medias(example[self.dataset_attr.audios]) if self.dataset_attr.audios else None, + } + return output + + +DATASET_CONVERTERS = { + "alpaca": AlpacaDatasetConverter, + "sharegpt": SharegptDatasetConverter, + "openai": OpenAIDatasetConverter, +} + + +def register_dataset_converter(name: str, dataset_converter: type["DatasetConverter"]) -> None: + r"""Register a new dataset converter.""" + if name in DATASET_CONVERTERS: + raise ValueError(f"Dataset converter {name} already exists.") + + DATASET_CONVERTERS[name] = dataset_converter + + +def get_dataset_converter(name: str, dataset_attr: "DatasetAttr", data_args: "DataArguments") -> "DatasetConverter": + r"""Get a dataset converter.""" + if name not in DATASET_CONVERTERS: + raise ValueError(f"Dataset converter {name} not found.") + + return DATASET_CONVERTERS[name](dataset_attr, data_args) + + +def align_dataset( + dataset: Union["Dataset", "IterableDataset"], + dataset_attr: "DatasetAttr", + data_args: "DataArguments", + training_args: "Seq2SeqTrainingArguments", +) -> Union["Dataset", "IterableDataset"]: + r"""Align the dataset to a specific format. + + Aligned dataset: + _prompt: [{"role": "user", "content": "..."}] * (2T - 1) + _response: [{"role": "assistant", "content": "..."}] * N (N > 1 for ranking dataset) + _system: "..." + _tools: "..." + _images: [] + _videos: [] + _audios: [] + """ + column_names = list(next(iter(dataset)).keys()) + kwargs = {} + if not data_args.streaming: + kwargs = dict( + num_proc=data_args.preprocessing_num_workers, + load_from_cache_file=(not data_args.overwrite_cache) or (training_args.local_process_index != 0), + desc="Converting format of dataset", + ) + + dataset_converter = get_dataset_converter(dataset_attr.formatting, dataset_attr, data_args) + return dataset.map( + dataset_converter, + batched=False, + remove_columns=column_names, + **kwargs, + ) diff --git a/src/llamafactory/data/data_utils.py b/src/llamafactory/data/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..144dbde1d8efa55fe1e9ff350110242678aaa223 --- /dev/null +++ b/src/llamafactory/data/data_utils.py @@ -0,0 +1,203 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 enum import StrEnum, unique +from typing import TYPE_CHECKING, Any, Optional, TypedDict, Union + +import fsspec +from datasets import DatasetDict, concatenate_datasets, interleave_datasets + +from ..extras import logging + + +if TYPE_CHECKING: + from datasets import Dataset, IterableDataset + + from ..hparams import DataArguments + + +logger = logging.get_logger(__name__) + + +SLOTS = list[Union[str, set[str], dict[str, str]]] + + +@unique +class Role(StrEnum): + USER = "user" + ASSISTANT = "assistant" + SYSTEM = "system" + FUNCTION = "function" + OBSERVATION = "observation" + + +class DatasetModule(TypedDict): + train_dataset: Optional[Union["Dataset", "IterableDataset"]] + eval_dataset: Optional[Union["Dataset", "IterableDataset", dict[str, "Dataset"]]] + + +def merge_dataset( + all_datasets: list[Union["Dataset", "IterableDataset"]], data_args: "DataArguments", seed: int +) -> Union["Dataset", "IterableDataset"]: + r"""Merge multiple datasets to a unified dataset.""" + if len(all_datasets) == 1: + return all_datasets[0] + + elif data_args.mix_strategy == "concat": + if data_args.streaming: + logger.warning_rank0_once("The samples between different datasets will not be mixed in streaming mode.") + + return concatenate_datasets(all_datasets) + + elif data_args.mix_strategy.startswith("interleave"): + if not data_args.streaming: + logger.warning_rank0_once("We recommend using `mix_strategy=concat` in non-streaming mode.") + + strategy_map: str = { + "interleave_under": "first_exhausted", + "interleave_over": "all_exhausted", + "interleave_once": "all_exhausted_without_replacement", + }[data_args.mix_strategy] + + return interleave_datasets( + datasets=all_datasets, + probabilities=data_args.interleave_probs, + seed=seed, + stopping_strategy=strategy_map, # type: ignore + ) + + else: + raise ValueError(f"Unknown mixing strategy: {data_args.mix_strategy}.") + + +def split_dataset( + dataset: Optional[Union["Dataset", "IterableDataset"]], + eval_dataset: Optional[Union["Dataset", "IterableDataset", dict[str, "Dataset"]]], + data_args: "DataArguments", + seed: int, +) -> tuple[dict, dict]: + r"""Split the dataset and returns two dicts containing train set and validation set. + + Support both map dataset and iterable dataset. + + Returns: + train_dict: Dictionary containing training data with key "train" + eval_dict: Dictionary containing evaluation data with keys "validation" or "validation_{name}" + """ + if eval_dataset is not None and data_args.val_size > 1e-6: + raise ValueError("Cannot specify `val_size` if `eval_dataset` is not None.") + + # the train and eval better to in dict dtype and separately return for cpode clearly and good handle outside + train_dict, eval_dict = {}, {} + + if dataset is not None: + if data_args.streaming: + dataset = dataset.shuffle(buffer_size=data_args.buffer_size, seed=seed) + + if data_args.val_size > 1e-6: + if data_args.streaming: + eval_dict["validation"] = dataset.take(int(data_args.val_size)) + train_dict["train"] = dataset.skip(int(data_args.val_size)) + else: + val_size = int(data_args.val_size) if data_args.val_size > 1 else data_args.val_size + split_result = dataset.train_test_split(test_size=val_size, seed=seed) + train_dict["train"] = split_result["train"] + eval_dict["validation"] = split_result["test"] + else: + train_dict["train"] = dataset + + if eval_dataset is not None: + if isinstance(eval_dataset, dict): + for name, data in eval_dataset.items(): + eval_dict[f"validation_{name}"] = data + else: + if data_args.streaming: + eval_dataset = eval_dataset.shuffle(buffer_size=data_args.buffer_size, seed=seed) + + eval_dict["validation"] = eval_dataset + + return train_dict, eval_dict + + +def get_dataset_module(dataset: Union["Dataset", "DatasetDict"]) -> "DatasetModule": + r"""Convert dataset or dataset dict to dataset module.""" + dataset_module: DatasetModule = {} + if isinstance(dataset, DatasetDict): # dataset dict + if "train" in dataset: + dataset_module["train_dataset"] = dataset["train"] + + if "validation" in dataset: + dataset_module["eval_dataset"] = dataset["validation"] + else: + eval_dataset = {} + for key in dataset.keys(): + if key.startswith("validation_"): + eval_dataset[key[len("validation_") :]] = dataset[key] + + if len(eval_dataset): + dataset_module["eval_dataset"] = eval_dataset + + else: # single dataset + dataset_module["train_dataset"] = dataset + + return dataset_module + + +def setup_fs(path: str, anon: bool = False) -> "fsspec.AbstractFileSystem": + r"""Set up a filesystem object based on the path protocol.""" + storage_options = {"anon": anon} if anon else {} + if path.startswith("s3://"): + fs = fsspec.filesystem("s3", **storage_options) + elif path.startswith(("gs://", "gcs://")): + fs = fsspec.filesystem("gcs", **storage_options) + else: + raise ValueError(f"Unsupported protocol in path: {path}. Use 's3://' or 'gs://'.") + + if not fs.exists(path): + raise ValueError(f"Path does not exist: {path}.") + + return fs + + +def _read_json_with_fs(fs: "fsspec.AbstractFileSystem", path: str) -> list[Any]: + r"""Helper function to read JSON/JSONL files using fsspec.""" + with fs.open(path, "r") as f: + if path.endswith(".jsonl"): + return [json.loads(line) for line in f if line.strip()] + else: + return json.load(f) + + +def read_cloud_json(cloud_path: str) -> list[Any]: + r"""Read a JSON/JSONL file from cloud storage (S3 or GCS). + + Args: + cloud_path: str + Cloud path in the format: + - 's3://bucket-name/file.json' for AWS S3 + - 'gs://bucket-name/file.jsonl' or 'gcs://bucket-name/file.jsonl' for Google Cloud Storage + """ + try: + fs = setup_fs(cloud_path, anon=True) # try with anonymous access first + except Exception: + fs = setup_fs(cloud_path) # try again with credentials + + # filter out non-JSON files + files = [x["Key"] for x in fs.listdir(cloud_path)] if fs.isdir(cloud_path) else [cloud_path] + files = list(filter(lambda file: file.endswith(".json") or file.endswith(".jsonl"), files)) + if not files: + raise ValueError(f"No JSON/JSONL files found in the specified path: {cloud_path}.") + + return sum([_read_json_with_fs(fs, file) for file in files], []) diff --git a/src/llamafactory/data/formatter.py b/src/llamafactory/data/formatter.py new file mode 100644 index 0000000000000000000000000000000000000000..1c080f8812d0d9feb320e949d39bfcb1a0e1e582 --- /dev/null +++ b/src/llamafactory/data/formatter.py @@ -0,0 +1,159 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 re +from abc import ABC, abstractmethod +from dataclasses import dataclass, field + +from typing_extensions import override + +from .data_utils import SLOTS +from .tool_utils import FunctionCall, get_tool_utils + + +@dataclass +class Formatter(ABC): + slots: SLOTS = field(default_factory=list) + tool_format: str | None = None + + @abstractmethod + def apply(self, **kwargs) -> SLOTS: + r"""Forms a list of slots according to the inputs to encode.""" + ... + + def extract(self, content: str) -> str | list["FunctionCall"]: + r"""Extract a list of tuples from the response message if using tools. + + Each tuple consists of function name and function arguments. + """ + raise NotImplementedError + + +@dataclass +class EmptyFormatter(Formatter): + def __post_init__(self): + has_placeholder = False + for slot in filter(lambda s: isinstance(s, str), self.slots): + if re.search(r"\{\{[a-zA-Z_][a-zA-Z0-9_]*\}\}", slot): + has_placeholder = True + + if has_placeholder: + raise ValueError("Empty formatter should not contain any placeholder.") + + @override + def apply(self, **kwargs) -> SLOTS: + return self.slots + + +@dataclass +class StringFormatter(Formatter): + def __post_init__(self): + has_placeholder = False + for slot in filter(lambda s: isinstance(s, str), self.slots): + if re.search(r"\{\{[a-zA-Z_][a-zA-Z0-9_]*\}\}", slot): + has_placeholder = True + + if not has_placeholder: + raise ValueError("A placeholder is required in the string formatter.") + + @override + def apply(self, **kwargs) -> SLOTS: + elements = [] + for slot in self.slots: + if isinstance(slot, str): + for name, value in kwargs.items(): + if not isinstance(value, str): + raise RuntimeError(f"Expected a string, got {value}") + + slot = slot.replace("{{" + name + "}}", value, 1) + elements.append(slot) + elif isinstance(slot, (dict, set)): + elements.append(slot) + else: + raise RuntimeError(f"Input must be string, set[str] or dict[str, str], got {type(slot)}.") + + return elements + + +@dataclass +class FunctionFormatter(StringFormatter): + def __post_init__(self): + super().__post_init__() + self.tool_utils = get_tool_utils(self.tool_format) + + @override + def apply(self, **kwargs) -> SLOTS: + content: str = kwargs.pop("content") + thought_words = kwargs.pop("thought_words", None) + tool_call_words = kwargs.pop("tool_call_words", None) + + def _parse_functions(json_content: str) -> list["FunctionCall"]: + try: + tool_calls = json.loads(json_content) + if not isinstance(tool_calls, list): # parallel function call + tool_calls = [tool_calls] + + return [FunctionCall(tc["name"], json.dumps(tc["arguments"], ensure_ascii=False)) for tc in tool_calls] + except json.JSONDecodeError: + raise RuntimeError(f"Invalid JSON format in function message: {str([content])}.") + + tool_call_match = None + if tool_call_words and len(tool_call_words) == 2: + tool_call_regex = re.compile( + rf"{re.escape(tool_call_words[0])}(.*?){re.escape(tool_call_words[1])}", re.DOTALL + ) + tool_call_match = re.search(tool_call_regex, content) + + if tool_call_match is None: + thought_match = None + if thought_words and len(thought_words) == 2: + regex = re.compile(rf"{re.escape(thought_words[0])}(.*?){re.escape(thought_words[1])}", re.DOTALL) + thought_match = re.search(regex, content) + + if thought_match: + json_part = content.replace(thought_match.group(0), "") + else: + json_part = content + + functions = _parse_functions(json_part) + function_str = self.tool_utils.function_formatter(functions) + if thought_match: + function_str = thought_match.group(0) + function_str + else: + thought_content = content.replace(tool_call_match.group(0), "") + functions = _parse_functions(tool_call_match.group(1)) + function_str = self.tool_utils.function_formatter(functions) + function_str = thought_content + function_str + + return super().apply(content=function_str) + + +@dataclass +class ToolFormatter(Formatter): + def __post_init__(self): + self.tool_utils = get_tool_utils(self.tool_format) + + @override + def apply(self, **kwargs) -> SLOTS: + content = kwargs.pop("content") + try: + tools = json.loads(content) + return [self.tool_utils.tool_formatter(tools) if len(tools) != 0 else ""] + except json.JSONDecodeError: + raise RuntimeError(f"Invalid JSON format in tool description: {str([content])}.") # flat string + + @override + def extract(self, content: str) -> str | list["FunctionCall"]: + return self.tool_utils.tool_extractor(content) diff --git a/src/llamafactory/data/loader.py b/src/llamafactory/data/loader.py new file mode 100644 index 0000000000000000000000000000000000000000..d3d44e6f3c25de0fa46c333be6a3a39d5f1dcfee --- /dev/null +++ b/src/llamafactory/data/loader.py @@ -0,0 +1,336 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 TYPE_CHECKING, Literal, Optional, Union + +import numpy as np +from datasets import Dataset, DatasetDict, load_dataset, load_from_disk + +from ..extras import logging +from ..extras.constants import FILEEXT2TYPE +from ..extras.misc import check_version, has_tokenized_data +from .converter import align_dataset +from .data_utils import get_dataset_module, merge_dataset, read_cloud_json, split_dataset +from .parser import get_dataset_list +from .processor import ( + FeedbackDatasetProcessor, + PackedSupervisedDatasetProcessor, + PairwiseDatasetProcessor, + PretrainDatasetProcessor, + SupervisedDatasetProcessor, + UnsupervisedDatasetProcessor, +) + + +if TYPE_CHECKING: + from datasets import Dataset, IterableDataset + from transformers import PreTrainedTokenizer, ProcessorMixin, Seq2SeqTrainingArguments + + from ..hparams import DataArguments, ModelArguments + from .data_utils import DatasetModule + from .parser import DatasetAttr + from .processor import DatasetProcessor + from .template import Template + + +logger = logging.get_logger(__name__) + + +def _load_single_dataset( + dataset_attr: "DatasetAttr", + model_args: "ModelArguments", + data_args: "DataArguments", + training_args: "Seq2SeqTrainingArguments", +) -> Union["Dataset", "IterableDataset"]: + r"""Load a single dataset and aligns it to the standard format.""" + logger.info_rank0(f"Loading dataset {dataset_attr}...") + data_path, data_name, data_dir, data_files = None, None, None, None + if dataset_attr.load_from in ["hf_hub", "ms_hub", "om_hub"]: + data_path = dataset_attr.dataset_name + data_name = dataset_attr.subset + data_dir = dataset_attr.folder + + elif dataset_attr.load_from == "script": + data_path = os.path.join(data_args.dataset_dir, dataset_attr.dataset_name) + data_name = dataset_attr.subset + data_dir = dataset_attr.folder + + elif dataset_attr.load_from == "cloud_file": + data_path = dataset_attr.dataset_name + + elif dataset_attr.load_from == "file": + data_files = [] + local_path = os.path.join(data_args.dataset_dir, dataset_attr.dataset_name) + if os.path.isdir(local_path): # is directory + for file_name in os.listdir(local_path): + data_files.append(os.path.join(local_path, file_name)) + elif os.path.isfile(local_path): # is file + data_files.append(local_path) + else: + raise ValueError(f"File {local_path} not found.") + + data_path = FILEEXT2TYPE.get(os.path.splitext(data_files[0])[-1][1:], None) + if data_path is None: + raise ValueError("Allowed file types: {}.".format(",".join(FILEEXT2TYPE.keys()))) + + if any(data_path != FILEEXT2TYPE.get(os.path.splitext(data_file)[-1][1:], None) for data_file in data_files): + raise ValueError("File types should be identical.") + else: + raise NotImplementedError(f"Unknown load type: {dataset_attr.load_from}.") + + if dataset_attr.load_from == "ms_hub": + check_version("modelscope>=1.14.0", mandatory=True) + from modelscope import MsDataset # type: ignore + from modelscope.utils.config_ds import MS_DATASETS_CACHE # type: ignore + + cache_dir = model_args.cache_dir or MS_DATASETS_CACHE + dataset = MsDataset.load( + dataset_name=data_path, + subset_name=data_name, + data_dir=data_dir, + data_files=data_files, + split=dataset_attr.split, + cache_dir=cache_dir, + token=model_args.ms_hub_token, + use_streaming=data_args.streaming, + ) + if isinstance(dataset, MsDataset): + dataset = dataset.to_hf_dataset() + + elif dataset_attr.load_from == "om_hub": + check_version("openmind>=0.8.0", mandatory=True) + from openmind import OmDataset # type: ignore + from openmind.utils.hub import OM_DATASETS_CACHE # type: ignore + + cache_dir = model_args.cache_dir or OM_DATASETS_CACHE + dataset = OmDataset.load_dataset( + path=data_path, + name=data_name, + data_dir=data_dir, + data_files=data_files, + split=dataset_attr.split, + cache_dir=cache_dir, + token=model_args.om_hub_token, + streaming=data_args.streaming, + ) + elif dataset_attr.load_from == "cloud_file": + dataset = Dataset.from_list(read_cloud_json(data_path), split=dataset_attr.split) + else: + dataset = load_dataset( + path=data_path, + name=data_name, + data_dir=data_dir, + data_files=data_files, + split=dataset_attr.split, + cache_dir=model_args.cache_dir, + token=model_args.hf_hub_token, + num_proc=data_args.preprocessing_num_workers, + streaming=data_args.streaming and dataset_attr.load_from != "file", + ) + if data_args.streaming and dataset_attr.load_from == "file": + dataset = dataset.to_iterable_dataset(num_shards=training_args.dataloader_num_workers) + + if dataset_attr.num_samples is not None and not data_args.streaming: + target_num = dataset_attr.num_samples + indexes = np.random.permutation(len(dataset))[:target_num] # all samples should be included + target_num -= len(indexes) + if target_num > 0: + expand_indexes = np.random.choice(len(dataset), target_num) + indexes = np.concatenate((indexes, expand_indexes), axis=0) + + assert len(indexes) == dataset_attr.num_samples, "Sample num mismatched." + dataset = dataset.select(indexes) + logger.info_rank0(f"Sampled {dataset_attr.num_samples} examples from dataset {dataset_attr}.") + + if data_args.max_samples is not None: # truncate dataset + max_samples = min(data_args.max_samples, len(dataset)) + dataset = dataset.select(range(max_samples)) + + return align_dataset(dataset, dataset_attr, data_args, training_args) + + +def _get_merged_dataset( + dataset_names: list[str] | None, + model_args: "ModelArguments", + data_args: "DataArguments", + training_args: "Seq2SeqTrainingArguments", + stage: Literal["pt", "sft", "rm", "ppo", "kto"], + return_dict: bool = False, +) -> Union["Dataset", "IterableDataset", dict[str, "Dataset"]] | None: + r"""Return the merged datasets in the standard format.""" + if dataset_names is None: + return None + + datasets = {} + for dataset_name, dataset_attr in zip(dataset_names, get_dataset_list(dataset_names, data_args.dataset_dir)): + if (stage == "rm" and dataset_attr.ranking is False) or (stage != "rm" and dataset_attr.ranking is True): + raise ValueError("The dataset is not applicable in the current training stage.") + + datasets[dataset_name] = _load_single_dataset(dataset_attr, model_args, data_args, training_args) + + if return_dict: + return datasets + else: + return merge_dataset(list(datasets.values()), data_args, seed=training_args.seed) + + +def _get_dataset_processor( + data_args: "DataArguments", + stage: Literal["pt", "sft", "rm", "ppo", "kto"], + template: "Template", + tokenizer: "PreTrainedTokenizer", + processor: Optional["ProcessorMixin"], + do_generate: bool = False, +) -> "DatasetProcessor": + r"""Return the corresponding dataset processor.""" + if stage == "pt": + dataset_processor_class = PretrainDatasetProcessor + elif stage == "sft" and not do_generate: + if data_args.packing: + if data_args.neat_packing: # hack datasets to have int32 attention mask + from datasets.arrow_writer import OptimizedTypedSequence, TypedSequence + + def __init__(self, data, **kwargs): + return TypedSequence.__init__( + self, + data, + type=kwargs.pop("type", None), + try_type=kwargs.pop("try_type", None), + optimized_int_type=kwargs.pop("optimized_int_type", None), + ) + + OptimizedTypedSequence.__init__ = __init__ + dataset_processor_class = PackedSupervisedDatasetProcessor + else: + dataset_processor_class = SupervisedDatasetProcessor + + elif stage == "rm": + dataset_processor_class = PairwiseDatasetProcessor + elif stage == "kto": + dataset_processor_class = FeedbackDatasetProcessor + else: + dataset_processor_class = UnsupervisedDatasetProcessor + + return dataset_processor_class(template=template, tokenizer=tokenizer, processor=processor, data_args=data_args) + + +def _get_preprocessed_dataset( + dataset: Union["Dataset", "IterableDataset"] | None, + data_args: "DataArguments", + training_args: "Seq2SeqTrainingArguments", + stage: Literal["pt", "sft", "rm", "ppo", "kto"], + template: "Template", + tokenizer: "PreTrainedTokenizer", + processor: Optional["ProcessorMixin"] = None, + is_eval: bool = False, +) -> Union["Dataset", "IterableDataset"] | None: + r"""Preprocesses the dataset, including format checking and tokenization.""" + if dataset is None: + return None + + dataset_processor = _get_dataset_processor( + data_args, stage, template, tokenizer, processor, do_generate=(training_args.predict_with_generate and is_eval) + ) + column_names = list(next(iter(dataset)).keys()) + kwargs = {} + if not data_args.streaming: + kwargs = dict( + num_proc=data_args.preprocessing_num_workers, + load_from_cache_file=(not data_args.overwrite_cache) or (training_args.local_process_index != 0), + desc="Running tokenizer on dataset", + ) + + dataset = dataset.map( + dataset_processor.preprocess_dataset, + batched=True, + batch_size=data_args.preprocessing_batch_size, + remove_columns=column_names, + **kwargs, + ) + + if training_args.should_log: + try: + print("eval example:" if is_eval else "training example:") + dataset_processor.print_data_example(next(iter(dataset))) + except StopIteration: + if stage == "pt": + raise RuntimeError("Cannot find sufficient samples, consider increasing dataset size.") + else: + raise RuntimeError("Cannot find valid samples, check `data/README.md` for the data format.") + + return dataset + + +def get_dataset( + template: "Template", + model_args: "ModelArguments", + data_args: "DataArguments", + training_args: "Seq2SeqTrainingArguments", + stage: Literal["pt", "sft", "rm", "ppo", "kto"], + tokenizer: "PreTrainedTokenizer", + processor: Optional["ProcessorMixin"] = None, +) -> "DatasetModule": + r"""Get the train dataset and optionally gets the evaluation dataset.""" + # Load tokenized dataset if path exists + if data_args.tokenized_path is not None: + if has_tokenized_data(data_args.tokenized_path): + logger.warning_rank0("Loading dataset from disk will ignore other data arguments.") + tokenized_data = load_from_disk(data_args.tokenized_path) + dataset_module = get_dataset_module(tokenized_data) + if data_args.streaming: + dataset_module["train_dataset"] = dataset_module["train_dataset"].to_iterable_dataset() + + logger.info_rank0(f"Loaded tokenized dataset from {data_args.tokenized_path}.") + return dataset_module + + if data_args.streaming: + raise ValueError("Turn off `streaming` when saving dataset to disk.") + + # Load and preprocess dataset + with training_args.main_process_first(desc="load dataset", local=(not data_args.data_shared_file_system)): + dataset = _get_merged_dataset(data_args.dataset, model_args, data_args, training_args, stage) + eval_dataset = _get_merged_dataset( + data_args.eval_dataset, + model_args, + data_args, + training_args, + stage, + return_dict=data_args.eval_on_each_dataset, + ) + + with training_args.main_process_first(desc="pre-process dataset", local=(not data_args.data_shared_file_system)): + # move front to make sure eval_dataset(if contain or split) can preprocessed appropriately + train_dict, eval_dict = split_dataset(dataset, eval_dataset, data_args, seed=training_args.seed) + + if "train" in train_dict: + train_dict["train"] = _get_preprocessed_dataset( + train_dict["train"], data_args, training_args, stage, template, tokenizer, processor, is_eval=False + ) + + for key in eval_dict: + eval_dict[key] = _get_preprocessed_dataset( + eval_dict[key], data_args, training_args, stage, template, tokenizer, processor, is_eval=True + ) + + # Combine train and eval dictionaries + dataset_dict = DatasetDict({**train_dict, **eval_dict}) + + if data_args.tokenized_path is not None: # save tokenized dataset to disk + if training_args.should_save: + dataset_dict.save_to_disk(data_args.tokenized_path) + logger.info_rank0(f"Tokenized dataset is saved at {data_args.tokenized_path}.") + logger.info_rank0(f"Please launch the training with `tokenized_path: {data_args.tokenized_path}`.") + + return get_dataset_module(dataset_dict) diff --git a/src/llamafactory/data/mm_plugin.py b/src/llamafactory/data/mm_plugin.py new file mode 100644 index 0000000000000000000000000000000000000000..8c2d66c2f5b769a571eb67b7371158a84d5eb380 --- /dev/null +++ b/src/llamafactory/data/mm_plugin.py @@ -0,0 +1,2944 @@ +# Copyright 2025 HuggingFace Inc. and the LlamaFactory team. +# +# This code is inspired by the HuggingFace's Transformers library. +# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/models/llava/processing_llava.py +# +# 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 inspect +import math +import os +import re +from copy import deepcopy +from dataclasses import dataclass +from io import BytesIO +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, BinaryIO, Literal, NotRequired, Optional, TypedDict, Union + +import numpy as np +import torch +import torchaudio +from transformers.image_utils import get_image_size, is_valid_image, make_flat_list_of_images, to_numpy_array +from transformers.models.mllama.processing_mllama import ( + convert_sparse_cross_attention_mask_to_dense, + get_cross_attention_token_mask, +) +from transformers.video_utils import make_batched_videos +from typing_extensions import override + +from ..extras.constants import AUDIO_PLACEHOLDER, IGNORE_INDEX, IMAGE_PLACEHOLDER, VIDEO_PLACEHOLDER +from ..extras.packages import is_pillow_available, is_pyav_available, is_transformers_version_greater_than + + +if is_pillow_available(): + from PIL import Image + from PIL.Image import Image as ImageObject + + +if is_pyav_available(): + import av + + +if TYPE_CHECKING: + from av.stream import Stream + from numpy.typing import NDArray + from transformers import PreTrainedTokenizer, ProcessorMixin + from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor + from transformers.image_processing_utils import BaseImageProcessor + from transformers.video_processing_utils import BaseVideoProcessor + + class EncodedImage(TypedDict): + path: str | None + bytes: bytes | None + + ImageInput = Union[str, bytes, EncodedImage, BinaryIO, ImageObject] + VideoInput = Union[str, BinaryIO, list[list[ImageInput]]] + AudioInput = Union[str, BinaryIO, NDArray] + + class RegularizedImageOutput(TypedDict): + images: list[ImageObject] + + class RegularizedVideoOutput(TypedDict): + videos: list[list[ImageObject]] + durations: list[float] + fps_per_video: NotRequired[list[float]] + + class RegularizedAudioOutput(TypedDict): + audios: list[NDArray] + sampling_rates: list[float] + + class MMProcessor(ProcessorMixin): + patch_size: int + image_seq_length: int + num_additional_image_tokens: int + vision_feature_select_strategy: Literal["default", "full"] + + def _get_number_of_features(self, orig_height: int, orig_width: int, height: int, width: int) -> int: + pass + + +def _get_paligemma_token_type_ids(imglens: list[int], seqlens: list[int], processor: "MMProcessor") -> list[list[int]]: + r"""Get paligemma token type ids for computing loss. + + It is slightly different with the original token type ids where the prompt part is 0. + + Returns: + batch_token_type_ids: shape (batch_size, seq_length) + + """ + batch_token_type_ids = [] + for imglen, seqlen in zip(imglens, seqlens): + image_seqlen = imglen * processor.image_seq_length + batch_token_type_ids.append([0] * image_seqlen + [1] * (seqlen - image_seqlen)) + + return batch_token_type_ids + + +def _get_gemma3_token_type_ids(batch_ids: list[list[int]], processor: "MMProcessor"): + r"""Get gemma3 token type ids for computing loss. + + Returns: + batch_token_type_ids: shape (batch_size, seq_length) + + """ + image_token_id: int = getattr(processor, "image_token_id") + batch_token_type_ids = [] + for token_ids in batch_ids: + token_ids = np.array(token_ids) + token_type_ids = np.zeros_like(token_ids) + token_type_ids[token_ids == image_token_id] = 1 + batch_token_type_ids.append(token_type_ids.tolist()) + + return batch_token_type_ids + + +def _make_batched_images(images: list["ImageObject"], imglens: list[int]) -> list[list["ImageObject"]]: + r"""Make nested list of images.""" + batch_images = [] + for imglen in imglens: + batch_images.append(images[:imglen]) + images = images[imglen:] + + return batch_images + + +def _check_video_is_nested_images(video: "VideoInput") -> bool: + r"""Check if the video is nested images.""" + return isinstance(video, list) and all(isinstance(frame, (str, BinaryIO, dict, ImageObject)) for frame in video) + + +@dataclass +class MMPluginMixin: + image_token: str | None + video_token: str | None + audio_token: str | None + expand_mm_tokens: bool = True + + def _validate_input( + self, + processor: Optional["MMProcessor"], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + ) -> None: + r"""Validate if this model accepts the input modalities.""" + image_processor: BaseImageProcessor = getattr(processor, "image_processor", None) + video_processor: BaseImageProcessor = getattr( + processor, "video_processor", getattr(processor, "image_processor", None) + ) + feature_extractor: SequenceFeatureExtractor = getattr(processor, "feature_extractor", None) or getattr( + processor, "audio_processor", None + ) + if len(images) != 0 and self.image_token is None: + raise ValueError( + "This model does not support image input. Please check whether the correct `template` is used." + ) + + if len(videos) != 0 and self.video_token is None: + raise ValueError( + "This model does not support video input. Please check whether the correct `template` is used." + ) + + if len(audios) != 0 and self.audio_token is None: + raise ValueError( + "This model does not support audio input. Please check whether the correct `template` is used." + ) + + if self.image_token is not None and processor is None: + raise ValueError("Processor was not found, please check and update your model file.") + + if self.image_token is not None and image_processor is None: + raise ValueError("Image processor was not found, please check and update your model file.") + + if self.video_token is not None and video_processor is None: + raise ValueError("Video processor was not found, please check and update your model file.") + + if self.audio_token is not None and feature_extractor is None: + raise ValueError("Audio feature extractor was not found, please check and update your model file.") + + def _validate_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + ): + r"""Validate if the number of images, videos and audios match the number of placeholders in messages.""" + num_image_tokens, num_video_tokens, num_audio_tokens = 0, 0, 0 + for message in messages: + num_image_tokens += message["content"].count(IMAGE_PLACEHOLDER) + num_video_tokens += message["content"].count(VIDEO_PLACEHOLDER) + num_audio_tokens += message["content"].count(AUDIO_PLACEHOLDER) + + if len(images) != num_image_tokens: + raise ValueError( + f"The number of images does not match the number of {IMAGE_PLACEHOLDER} tokens in {messages}." + ) + + if len(videos) != num_video_tokens: + raise ValueError( + f"The number of videos does not match the number of {VIDEO_PLACEHOLDER} tokens in {messages}." + ) + + if len(audios) != num_audio_tokens: + raise ValueError( + f"The number of audios does not match the number of {AUDIO_PLACEHOLDER} tokens in {messages}." + ) + + def _preprocess_image( + self, image: "ImageObject", image_max_pixels: int, image_min_pixels: int, **kwargs + ) -> "ImageObject": + r"""Pre-process a single image.""" + if (image.width * image.height) > image_max_pixels: + resize_factor = math.sqrt(image_max_pixels / (image.width * image.height)) + width, height = int(image.width * resize_factor), int(image.height * resize_factor) + image = image.resize((width, height)) + + if (image.width * image.height) < image_min_pixels: + resize_factor = math.sqrt(image_min_pixels / (image.width * image.height)) + width, height = int(image.width * resize_factor), int(image.height * resize_factor) + image = image.resize((width, height)) + + if image.mode != "RGB": + image = image.convert("RGB") + + return image + + def _get_video_sample_indices( + self, video_stream: "Stream", video_fps: float, video_maxlen: int, **kwargs + ) -> list[int]: + r"""Compute video sample indices according to fps.""" + total_frames = video_stream.frames + if total_frames == 0: # infinite video + return np.linspace(0, video_maxlen - 1, video_maxlen).astype(np.int32) + + sample_frames = max(1, math.floor(float(video_stream.duration * video_stream.time_base) * video_fps)) + sample_frames = min(total_frames, video_maxlen, sample_frames) + return np.linspace(0, total_frames - 1, sample_frames).astype(np.int32) + + def _get_video_token_metadata( + self, + videos: list["VideoInput"], + processor: "MMProcessor", + ) -> Optional[dict[str, Any]]: + r"""Build metadata used to expand video tokens without decoding frames.""" + return None + + def _regularize_images(self, images: list["ImageInput"], **kwargs) -> "RegularizedImageOutput": + r"""Regularize images to avoid error. Including reading and pre-processing.""" + results = [] + for image in images: + if isinstance(image, (str, BinaryIO)): + image = Image.open(image) + elif isinstance(image, bytes): + image = Image.open(BytesIO(image)) + elif isinstance(image, dict): + if image["bytes"] is not None: + image = Image.open(BytesIO(image["bytes"])) + else: + image = Image.open(image["path"]) + + if not isinstance(image, ImageObject): + raise ValueError(f"Expect input is a list of images, but got {type(image)}.") + + results.append(self._preprocess_image(image, **kwargs)) + + return {"images": results} + + def _regularize_videos(self, videos: list["VideoInput"], **kwargs) -> "RegularizedVideoOutput": + r"""Regularizes videos to avoid error. Including reading, resizing and converting.""" + results = [] + durations = [] + for video in videos: + frames: list[ImageObject] = [] + if _check_video_is_nested_images(video): + for frame in video: + if not is_valid_image(frame) and not isinstance(frame, dict) and not os.path.exists(frame): + raise ValueError("Invalid image found in video frames.") + frames = video + durations.append(len(frames) / kwargs.get("video_fps", 2.0)) + else: + container = av.open(video, "r") + video_stream = next(stream for stream in container.streams if stream.type == "video") + sample_indices = self._get_video_sample_indices(video_stream, **kwargs) + container.seek(0) + for frame_idx, frame in enumerate(container.decode(video_stream)): + if frame_idx in sample_indices: + frames.append(frame.to_image()) + + if video_stream.duration is None: + durations.append(len(frames) / kwargs.get("video_fps", 2.0)) + else: + durations.append(float(video_stream.duration * video_stream.time_base)) + + frames = self._regularize_images(frames, **kwargs)["images"] + results.append(frames) + + return {"videos": results, "durations": durations} + + def _regularize_audios( + self, audios: list["AudioInput"], sampling_rate: float, **kwargs + ) -> "RegularizedAudioOutput": + r"""Regularizes audios to avoid error. Including reading and resampling.""" + results, sampling_rates = [], [] + for audio in audios: + if not isinstance(audio, np.ndarray): + audio, sr = torchaudio.load(audio) + if audio.shape[0] > 1: + audio = audio.mean(dim=0, keepdim=True) + + if sr != sampling_rate: + audio = torchaudio.functional.resample(audio, sr, sampling_rate) + + audio = audio.squeeze(0).numpy() + + results.append(audio) + sampling_rates.append(sampling_rate) + + return {"audios": results, "sampling_rates": sampling_rates} + + def _get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "MMProcessor", + imglens: list[int] | None = None, + ) -> dict[str, "torch.Tensor"]: + r"""Process visual inputs. + + Returns: (llava and paligemma) + pixel_values: tensor with shape (B, C, H, W) + + Returns: (qwen2-vl) + pixel_values: tensor with shape (num_patches, patch_dim) + image_grid_thw: tensor with shape (num_images, 3), where the three numbers are time, width, height + where num_patches == torch.prod(image_grid_thw) + + Returns: (mllama) + pixel_values: tensor with shape + (batch_size, max_num_images, max_image_tiles, channels, tile_height, tile_width) + For example, (2, 1, 4, 3, 560, 560). + aspect_ratio_ids: tensor with shape (batch_size, max_num_images). For example, (2, 1). + aspect_ratio_mask: tensor with shape (batch_size, max_num_images, max_image_tiles). For example, (2, 1, 4). + num_tiles: List[List[int]] with shape (batch_size, num_images_in_batch). For example, (2, 1). + + """ + mm_inputs = {} + if len(images) != 0: + image_processor: BaseImageProcessor = getattr(processor, "image_processor", None) + images = self._regularize_images( + images, + image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768), + image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32), + )["images"] + if imglens is not None: # if imglens are provided, make batched images + images = _make_batched_images(images, imglens) + + image_processor_kwargs = {} + if getattr(processor, "image_do_pan_and_scan", False): # gemma3 image processor + image_processor_kwargs.update( + { + "do_pan_and_scan": True, + "pan_and_scan_min_crop_size": 256, + "pan_and_scan_max_num_crops": 4, + "pan_and_scan_min_ratio_to_activate": 1.2, + } + ) + + mm_inputs.update(image_processor(images, return_tensors="pt", **image_processor_kwargs)) + + if len(videos) != 0: + video_processor: BaseImageProcessor = getattr( + processor, "video_processor", getattr(processor, "image_processor", None) + ) + videos = self._regularize_videos( + videos, + image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256), + image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16), + video_fps=getattr(processor, "video_fps", 2.0), + video_maxlen=getattr(processor, "video_maxlen", 128), + )["videos"] + if "videos" in inspect.signature(video_processor.preprocess).parameters: # for qwen2_vl and video_llava + mm_inputs.update(video_processor(images=None, videos=videos, return_tensors="pt")) + else: # for llava_next_video + mm_inputs.update(video_processor(videos, return_tensors="pt")) + + if len(audios) != 0: + feature_extractor: SequenceFeatureExtractor = getattr(processor, "feature_extractor", None) or getattr( + processor, "audio_processor", None + ) + audios = self._regularize_audios( + audios, + sampling_rate=getattr(processor, "audio_sampling_rate", 16000), + )["audios"] + mm_inputs.update( + feature_extractor( + audios, + sampling_rate=getattr(processor, "audio_sampling_rate", 16000), + return_attention_mask=True, + padding="max_length", + return_tensors="pt", + ) + ) + mm_inputs["feature_attention_mask"] = mm_inputs.pop("attention_mask", None) # prevent conflicts + + return mm_inputs + + +@dataclass +class BasePlugin(MMPluginMixin): + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + r"""Pre-process input messages before tokenization for VLMs.""" + self._validate_input(processor, images, videos, audios) + return messages + + def process_token_ids( + self, + input_ids: list[int], + labels: list[int] | None, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + tokenizer: "PreTrainedTokenizer", + processor: Optional["MMProcessor"], + ) -> tuple[list[int], list[int] | None]: + r"""Pre-process token ids after tokenization for VLMs.""" + self._validate_input(processor, images, videos, audios) + return input_ids, labels + + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["MMProcessor"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + r"""Build batched multimodal inputs for VLMs. + + Arguments: + images: a list of image inputs, shape (num_images,) + videos: a list of video inputs, shape (num_videos,) + audios: a list of audio inputs, shape (num_audios,) + imglens: number of images in each sample, shape (batch_size,) + vidlens: number of videos in each sample, shape (batch_size,) + audlens: number of audios in each sample, shape (batch_size,) + batch_ids: token ids of input samples, shape (batch_size, seq_len) + processor: a processor for pre-processing images and videos + + """ + self._validate_input(processor, images, videos, audios) + return self._get_mm_inputs(images, videos, audios, processor) + + +@dataclass +class ErnieVLPlugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + messages = deepcopy(messages) + + image_processor: BaseImageProcessor = getattr(processor, "image_processor") + + merge_length: int = getattr(image_processor, "merge_size") ** 2 + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + image_grid_thw = mm_inputs.get("image_grid_thw", []) + video_grid_thw = mm_inputs.get("video_grid_thw", []) + else: + image_grid_thw = [None] * len(images) + video_grid_thw = [None] * len(videos) + + image_idx, video_idx = 0, 0 + for message in messages: + content = message["content"] + image_token = self.image_token or "<|IMAGE_PLACEHOLDER|>" + video_token = self.video_token or "<|VIDEO_PLACEHOLDER|>" + while IMAGE_PLACEHOLDER in content: + image_seqlen = image_grid_thw[image_idx].prod() // merge_length if self.expand_mm_tokens else 1 + content = content.replace( + IMAGE_PLACEHOLDER, + f"Picture {image_idx + 1}:<|IMAGE_START|>{image_token * image_seqlen}<|IMAGE_END|>", + 1, + ) + image_idx += 1 + while VIDEO_PLACEHOLDER in content: + video_seqlen = video_grid_thw[video_idx].prod() // merge_length if self.expand_mm_tokens else 1 + content = content.replace( + VIDEO_PLACEHOLDER, + f"Video {video_idx + 1}:<|VIDEO_START|>{video_token * video_seqlen}<|VIDEO_END|>", + 1, + ) + video_idx += 1 + message["content"] = content + return messages + + +@dataclass +class Gemma3Plugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens = 0 + messages = deepcopy(messages) + boi_token: str = getattr(processor, "boi_token") + full_image_sequence: str = getattr(processor, "full_image_sequence") + image_str = full_image_sequence if self.expand_mm_tokens else boi_token + + do_pan_and_scan: bool = getattr(processor, "image_do_pan_and_scan", False) + if do_pan_and_scan: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + if do_pan_and_scan: + image_placeholder_str = ( + "Here is the original image {{image}} and here are some crops to help you see better " + + " ".join(["{{image}}"] * mm_inputs["num_crops"][0][num_image_tokens]) + ) + else: + image_placeholder_str = "{{image}}" + + content = content.replace(IMAGE_PLACEHOLDER, image_placeholder_str, 1) + num_image_tokens += 1 + + message["content"] = content.replace("{{image}}", image_str) + + return messages + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["MMProcessor"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + mm_inputs.pop("num_crops", None) + mm_inputs["token_type_ids"] = _get_gemma3_token_type_ids(batch_ids, processor) + return mm_inputs + + +class Gemma3nPlugin(Gemma3Plugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + messages = deepcopy(messages) + boi_token: str = getattr(processor, "boi_token") + boa_token: str = getattr(processor, "boa_token") + full_image_sequence: str = getattr(processor, "full_image_sequence") + full_audio_sequence: str = getattr(processor, "full_audio_sequence") + image_str = full_image_sequence if self.expand_mm_tokens else boi_token + audio_str = full_audio_sequence if self.expand_mm_tokens else boa_token + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + content = content.replace(IMAGE_PLACEHOLDER, image_str, 1) + + while AUDIO_PLACEHOLDER in content: + content = content.replace(AUDIO_PLACEHOLDER, audio_str, 1) + + message["content"] = content + + return messages + + +@dataclass +class Gemma4Plugin(BasePlugin): + r"""Plugin for the Gemma4 multimodal model.""" + + @override + def _regularize_videos(self, videos: list["VideoInput"], **kwargs) -> "RegularizedVideoOutput": + r"""Regularize videos, also tracking per-video FPS and frame indices for timestamp generation.""" + results, fps_per_video, durations, frames_indices = [], [], [], [] + for video in videos: + frames: list[ImageObject] = [] + if _check_video_is_nested_images(video): + frames = video + fps_per_video.append(kwargs.get("video_fps", 2.0)) + durations.append(len(frames) / kwargs.get("video_fps", 2.0)) + frames_indices.append(list(range(len(frames)))) + else: + container = av.open(video, "r") + video_stream = next(stream for stream in container.streams if stream.type == "video") + sample_indices = self._get_video_sample_indices(video_stream, **kwargs) + original_fps = float(video_stream.average_rate) + # for correctly calculate timestamps + frames_indices.append([idx / original_fps * kwargs.get("video_fps", 2.0) for idx in sample_indices]) + container.seek(0) + for frame_idx, frame in enumerate(container.decode(video_stream)): + if frame_idx in sample_indices: + frames.append(frame.to_image()) + + if video_stream.duration is None: + durations.append(len(frames) / kwargs.get("video_fps", 2.0)) + else: + durations.append(float(video_stream.duration * video_stream.time_base)) + + frames = self._regularize_images(frames, **kwargs)["images"] + results.append(frames) + + return { + "videos": results, + "fps_per_video": fps_per_video, + "durations": durations, + "frames_indices": frames_indices, + } + + @override + def _get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "MMProcessor", + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + image_processor = getattr(processor, "image_processor", None) + video_processor = getattr(processor, "video_processor", None) + feature_extractor = getattr(processor, "feature_extractor", None) + mm_inputs = {} + + if len(images) != 0: + regularized = self._regularize_images( + images, + image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768), + image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32), + )["images"] + mm_inputs.update(image_processor(regularized, return_tensors="pt")) + + if len(videos) != 0: + video_data = self._regularize_videos( + videos, + image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256), + image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16), + video_fps=getattr(processor, "video_fps", 2.0), + video_maxlen=getattr(processor, "video_maxlen", 128), + ) + video_metadata = [ + { + "fps": getattr(processor, "video_fps", 2.0), + "duration": duration, + "total_num_frames": len(video), + "frames_indices": sample_indices, + } + for video, duration, sample_indices in zip( + video_data["videos"], video_data["durations"], video_data["frames_indices"] + ) + ] + mm_inputs.update( + video_processor( + videos=video_data["videos"], + video_metadata=video_metadata, + return_tensors="pt", + return_metadata=True, + do_sample_frames=False, + ) + ) + + if len(audios) != 0: # only for gemma4n + audios = self._regularize_audios( + audios, + sampling_rate=getattr(processor, "audio_sampling_rate", 16000), + )["audios"] + + mm_inputs.update( + feature_extractor( + audios, + padding="max_length", + return_tensors="pt", + ) + ) + + return mm_inputs + + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + messages = deepcopy(messages) + + boi_token: str = getattr(processor, "boi_token") + eoi_token: str = getattr(processor, "eoi_token") + boa_token: str = getattr(processor, "boa_token") + eoa_token: str = getattr(processor, "eoa_token") + image_token: str = getattr(processor, "image_token") + video_token: str = getattr(processor, "video_token") + audio_token: str = getattr(processor, "audio_token") + + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + num_image_soft_tokens: list[int] = list( + mm_inputs.get("num_soft_tokens_per_image", [getattr(processor, "image_seq_length", 256)] * len(images)) + ) + num_video_soft_tokens: list[int] = list(mm_inputs.get("num_soft_tokens_per_video", [1] * len(videos))) + video_metadata = mm_inputs.get("video_metadata", []) + else: + num_image_soft_tokens = [1] * len(images) + num_video_soft_tokens = [1] * len(videos) + video_metadata = [None] * len(videos) + + audio_iter = iter(audios) + image_iter = iter(num_image_soft_tokens) + video_iter = iter(zip(num_video_soft_tokens, video_metadata)) + + for message in messages: + content = message["content"] + + while IMAGE_PLACEHOLDER in content: + n = next(image_iter) + content = content.replace(IMAGE_PLACEHOLDER, f"{boi_token}{image_token * n}{eoi_token}", 1) + + while VIDEO_PLACEHOLDER in content: + num_soft_tokens_per_frame, metadata = next(video_iter) + if self.expand_mm_tokens: + timestamp_strs = [f"{int(t // 60):02d}:{int(t % 60):02d}" for t in metadata.timestamps] + frame_strs = [ + f"{ts} {boi_token}{video_token * num_soft_tokens_per_frame}{eoi_token}" + for ts in timestamp_strs + ] + video_str = " ".join(frame_strs) + else: + video_str = f"{boi_token}{video_token * num_soft_tokens_per_frame}{eoi_token}" + content = content.replace(VIDEO_PLACEHOLDER, video_str, 1) + + while AUDIO_PLACEHOLDER in content: + current_audio = next(audio_iter) + if self.expand_mm_tokens: + num_audio_tokens = processor._compute_audio_num_tokens( + current_audio, processor.feature_extractor.sampling_rate + ) + audio_str = f"{boa_token}{audio_token * num_audio_tokens}{eoa_token}" + else: + audio_str = f"{boa_token}{audio_token}{eoa_token}" + + content = content.replace(AUDIO_PLACEHOLDER, audio_str, 1) + + message["content"] = content + + return messages + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["MMProcessor"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + # Pop metadata keys that must not be passed to the model. + for key in ( + "num_soft_tokens_per_image", + "num_soft_tokens_per_video", + "video_metadata", + "_gemma4_fps_per_video", + "_gemma4_frames_indices", + "_gemma4_num_audio_soft_tokens", + ): + mm_inputs.pop(key, None) + + mm_inputs["mm_token_type_ids"] = processor.create_mm_token_type_ids(batch_ids) + + return mm_inputs + + +@dataclass +class InternVLPlugin(BasePlugin): + @override + def _get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "ProcessorMixin", + **kwargs, + ) -> dict[str, "torch.Tensor"]: + image_processor: BaseImageProcessor = getattr(processor, "image_processor") + image_processor_kwargs = {} + if getattr(processor, "crop_to_patches", False): + image_processor_kwargs.update( + { + "crop_to_patches": True, + "max_patches": 12, + "min_patches": 1, + } + ) + + mm_inputs = {} + image_video_patches = [] + + if len(images) != 0: + images = self._regularize_images( + images, + image_max_pixels=getattr(processor, "image_max_pixels", 1024 * 1024), + image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32), + )["images"] + + if len(videos) != 0: + videos = self._regularize_videos( + videos, + image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256), + image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16), + video_fps=getattr(processor, "video_fps", 2.0), + video_maxlen=getattr(processor, "video_maxlen", 128), + )["videos"] + + if len(images) != 0: + images = make_flat_list_of_images(images) + image_inputs = image_processor(images=images, return_tensors="pt", **image_processor_kwargs) + image_num_patches = image_inputs.pop("num_patches") + image_pixel_values = image_inputs.pop("pixel_values") + image_num_patches_indices = np.cumsum(image_num_patches) + + if len(videos) != 0: + videos = make_batched_videos(videos) + num_frames_per_video = [len(video) for video in videos] + patch_indices = np.cumsum(num_frames_per_video) + image_processor_kwargs["crop_to_patches"] = False + video_inputs = image_processor(images=videos, return_tensors="pt", **image_processor_kwargs) + video_num_patches = video_inputs.pop("num_patches") + video_pixel_values = video_inputs.pop("pixel_values") + video_num_patches_indices = np.cumsum(video_num_patches) + + # NOT SUPPORT IMAGE VIDEO INTERLEAVED + if len(images) != 0 and image_pixel_values is not None: + for i in range(len(images)): + start_index = image_num_patches_indices[i - 1] if i > 0 else 0 + end_index = image_num_patches_indices[i] + image_video_patches.append(image_pixel_values[start_index:end_index]) + + if len(videos) != 0 and video_pixel_values is not None: + patch_indices_with_prefix = [0] + list(patch_indices) + for i in range(len(videos)): + current_patch_index = patch_indices_with_prefix[i] + end_patch_index = patch_indices_with_prefix[i + 1] + start_index = video_num_patches_indices[current_patch_index - 1] if i > 0 else 0 + end_index = video_num_patches_indices[end_patch_index - 1] + image_video_patches.append(video_pixel_values[start_index:end_index]) + + if len(images) != 0 or len(videos) != 0: + mm_inputs["pixel_values"] = torch.cat(image_video_patches, dim=0) + + if len(images) != 0: + mm_inputs.update({"image_num_patches": image_num_patches}) + + if len(videos) != 0: + mm_inputs.update({"video_patch_indices": patch_indices}) + mm_inputs.update({"video_num_patches": video_num_patches}) + + return mm_inputs + + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["ProcessorMixin"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens, num_video_tokens = 0, 0 + image_seqlen = getattr(processor, "image_seq_length") if self.expand_mm_tokens else 1 + messages = deepcopy(messages) + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + + image_pixel_patch_list = mm_inputs.get("image_num_patches") # patches of images + video_num_patches = mm_inputs.get("video_num_patches") # all patches for frames of videos + video_patch_indices = mm_inputs.get("video_patch_indices") # num frames of per video + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + content = content.replace( + IMAGE_PLACEHOLDER, + f"{'' * image_seqlen * image_pixel_patch_list[num_image_tokens]}", + 1, + ) + num_image_tokens += 1 + + while VIDEO_PLACEHOLDER in content: + current_patch_index = video_patch_indices[num_video_tokens - 1] if num_video_tokens > 0 else 0 + end_patch_index = video_patch_indices[num_video_tokens] + num_patches = list(video_num_patches[current_patch_index:end_patch_index]) + video_replaced_prompt = "\n".join( + f"Frame{i + 1}: {'' * image_seqlen * num_patches[i]}" + for i in range(len(num_patches)) + ) + content = content.replace(VIDEO_PLACEHOLDER, video_replaced_prompt, 1) + num_video_tokens += 1 + + message["content"] = content + + return messages + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["ProcessorMixin"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + mm_inputs.pop("image_num_patches", None) + mm_inputs.pop("video_patch_indices", None) + mm_inputs.pop("video_num_patches", None) + return mm_inputs + + +class KimiVLPlugin(BasePlugin): + @override + def process_messages(self, messages, images, videos, audios, processor): + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + image_grid_hws = mm_inputs.get("image_grid_hws", []) + else: + image_grid_hws = [None] * len(images) + + num_image_tokens = 0 + image_processor: BaseImageProcessor = getattr(processor, "image_processor") + merge_length = math.prod(image_processor.merge_kernel_size) + messages = deepcopy(messages) + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + image_seqlen = image_grid_hws[num_image_tokens].prod() // merge_length if self.expand_mm_tokens else 1 + content = content.replace( + IMAGE_PLACEHOLDER, + f"<|media_start|>image<|media_content|>{self.image_token * image_seqlen}<|media_end|>", + 1, + ) + num_image_tokens += 1 + + message["content"] = content + + return messages + + +@dataclass +class Llama4Plugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + if "pixel_values" in mm_inputs: + image_height, image_width = mm_inputs["pixel_values"][0].shape[-2:] + num_patches_per_chunk = int( + (image_height // processor.patch_size) + * (image_width // processor.patch_size) + // processor.downsample_ratio + ) + aspect_ratios = mm_inputs.pop("aspect_ratios") + + num_image_tokens = 0 + messages = deepcopy(messages) + for message in messages: + content = message["content"] + if self.expand_mm_tokens: + placeholder_count = content.count(IMAGE_PLACEHOLDER) + prompt_splits = content.split(IMAGE_PLACEHOLDER) + new_content = [] + for local_image_index, split_part in enumerate(prompt_splits): + new_content.append(split_part) + if local_image_index < placeholder_count: + tokens_for_this_image = processor._prompt_split_image( + aspect_ratios[num_image_tokens], num_patches_per_chunk + ) + num_image_tokens += 1 + new_content.append(tokens_for_this_image) + + content = "".join(new_content) + else: + content = content.replace(IMAGE_PLACEHOLDER, self.image_token) + + message["content"] = content + + return messages + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["MMProcessor"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + mm_inputs.pop("aspect_ratios", None) + return mm_inputs + + +@dataclass +class LlavaPlugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + messages = deepcopy(messages) + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + if "pixel_values" in mm_inputs: + height, width = get_image_size(to_numpy_array(mm_inputs["pixel_values"][0])) + image_seqlen = (height // processor.patch_size) * ( + width // processor.patch_size + ) + processor.num_additional_image_tokens + if processor.vision_feature_select_strategy == "default": + image_seqlen -= 1 + else: + image_seqlen = 1 + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + content = content.replace(IMAGE_PLACEHOLDER, "{{image}}" * image_seqlen, 1) + + message["content"] = content.replace("{{image}}", self.image_token) + + return messages + + +@dataclass +class LlavaNextPlugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens = 0 + messages = deepcopy(messages) + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + if "pixel_values" in mm_inputs: + image_sizes = iter(mm_inputs["image_sizes"].tolist()) + height, width = get_image_size(to_numpy_array(mm_inputs["pixel_values"][0][0])) + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + if self.expand_mm_tokens: + orig_height, orig_width = next(image_sizes) + image_seqlen = processor._get_number_of_features(orig_height, orig_width, height, width) + if processor.vision_feature_select_strategy == "default": + image_seqlen -= 1 + else: + image_seqlen = 1 + + content = content.replace(IMAGE_PLACEHOLDER, "{{image}}" * image_seqlen, 1) + num_image_tokens += 1 + + message["content"] = content.replace("{{image}}", self.image_token) + + return messages + + +@dataclass +class LlavaNextVideoPlugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + messages = deepcopy(messages) + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + if "pixel_values" in mm_inputs: + image_sizes = iter(mm_inputs["image_sizes"].tolist()) + height, width = get_image_size(to_numpy_array(mm_inputs["pixel_values"][0][0])) + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + if self.expand_mm_tokens: + orig_height, orig_width = next(image_sizes) + image_seqlen = processor._get_number_of_features(orig_height, orig_width, height, width) + if processor.vision_feature_select_strategy == "default": + image_seqlen -= 1 + else: + image_seqlen = 1 + + content = content.replace(IMAGE_PLACEHOLDER, "{{image}}" * image_seqlen, 1) + + message["content"] = content.replace("{{image}}", self.image_token) + + if self.expand_mm_tokens: + if "pixel_values_videos" in mm_inputs: + one_video = to_numpy_array(mm_inputs.get("pixel_values_videos")[0]) + height, width = get_image_size(one_video[0]) + num_frames = one_video.shape[0] # frame dim is always after batch dim + image_seqlen = (height // processor.patch_size) * (width // processor.patch_size) + video_seqlen = image_seqlen // 4 * num_frames # divide by 4 needed for avg pooling layer + else: + video_seqlen = 1 + + for message in messages: + content = message["content"] + while VIDEO_PLACEHOLDER in content: + content = content.replace(VIDEO_PLACEHOLDER, "{{video}}" * video_seqlen, 1) + + message["content"] = content.replace("{{video}}", self.video_token) + + return messages + + +@dataclass +class MiniCPMVPlugin(BasePlugin): + @override + def _get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "MMProcessor", + **kwargs, + ) -> dict[str, "torch.Tensor"]: + image_processor: BaseImageProcessor = getattr(processor, "image_processor") + mm_inputs = {} + if len(images) != 0: + images = self._regularize_images( + images, + image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768), + image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32), + )["images"] + if "valid_image_nums_ls" in kwargs: + valid_image_nums_ls = kwargs["valid_image_nums_ls"] + new_images = [] + idx = 0 + for valid_image_nums in valid_image_nums_ls: + new_images.append(images[idx : idx + valid_image_nums]) + idx += valid_image_nums + + images = new_images + + image_inputs = image_processor( + images, do_pad=True, max_slice_nums=image_processor.max_slice_nums, return_tensors="pt" + ) + mm_inputs.update(image_inputs) + + if len(videos) != 0: + videos = self._regularize_videos( + videos, + image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256), + image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16), + video_fps=getattr(processor, "video_fps", 2.0), + video_maxlen=getattr(processor, "video_maxlen", 128), + )["videos"] + video_inputs = image_processor(videos, do_pad=True, max_slice_nums=2, return_tensors="pt") + mm_inputs.update(video_inputs) + + if len(audios) != 0: + audios = self._regularize_audios( + audios, + sampling_rate=getattr(processor, "audio_sampling_rate", 16000), + )["audios"] + if "valid_audio_nums_ls" in kwargs: + valid_audio_nums_ls = kwargs["valid_audio_nums_ls"] + audios_ls = [] + idx = 0 + for valid_audio_nums in valid_audio_nums_ls: + audios_ls.append(audios[idx : idx + valid_audio_nums]) + idx += valid_audio_nums + else: + audios_ls = [audios] + + audio_features, audio_feature_lens, audio_phs = processor.audio_feature_extract( + audios_ls, + chunk_input=True, + sampling_rate=getattr(processor, "audio_sampling_rate", 16000), + ) + audio_feature_lens = [ + x.clone().detach() if isinstance(x, torch.Tensor) else torch.tensor(x) for x in audio_feature_lens + ] + mm_inputs.update({"audio_features": audio_features, "audio_feature_lens": audio_feature_lens}) + if kwargs.get("ret_phs", False): + mm_inputs.update({"audio_phs": audio_phs}) + + return mm_inputs + + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens, num_video_tokens, num_audio_tokens = 0, 0, 0 + messages = deepcopy(messages) + image_processor: BaseImageProcessor = getattr(processor, "image_processor") + mm_inputs, audio_inputs = {}, {} + if len(images) != 0 and len(videos) != 0: + raise ValueError("MiniCPM-V model does not support input images and videos at the same time.") + + if len(videos) != 0: + max_slice_nums = 2 + use_image_id = False + mm_inputs = self._get_mm_inputs([], videos, [], processor) + else: + max_slice_nums = image_processor.max_slice_nums + use_image_id = image_processor.use_image_id + + for i, message in enumerate(messages): + content = message["content"] + while IMAGE_PLACEHOLDER in content: + content = content.replace(IMAGE_PLACEHOLDER, "{{image}}", 1) + num_image_tokens += 1 + + while VIDEO_PLACEHOLDER in content: + video_seqlen = len(mm_inputs["image_sizes"][num_video_tokens]) if self.expand_mm_tokens else 1 + content = content.replace(VIDEO_PLACEHOLDER, "{{image}}" * video_seqlen, 1) + num_video_tokens += 1 + + while AUDIO_PLACEHOLDER in content: + content = content.replace(AUDIO_PLACEHOLDER, "{{audio}}", 1) + num_audio_tokens += 1 + + message["content"] = content.replace("{{image}}", "(./)").replace( + "{{audio}}", "()" + ) + + if len(images): + mm_inputs = self._get_mm_inputs(images, [], [], processor) + + if len(audios): + audio_inputs = self._get_mm_inputs([], [], audios, processor, ret_phs=True) + + if self.expand_mm_tokens and mm_inputs: + pattern = "(./)" + image_sizes = mm_inputs["image_sizes"] + idx = 0 + for index, message in enumerate(messages): + text = message["content"] + image_tags = re.findall(pattern, text) + text_chunks = text.split(pattern) + final_text = "" + for i in range(len(image_tags)): + final_text = ( + final_text + + text_chunks[i] + + image_processor.get_slice_image_placeholder( + image_sizes[0][idx], idx, max_slice_nums, use_image_id + ) + ) + idx += 1 + + final_text += text_chunks[-1] + messages[index]["content"] = final_text + + if self.expand_mm_tokens and audio_inputs: + pattern = "()" + idx = 0 + for index, message in enumerate(messages): + text = message["content"] + audio_tags = re.findall(pattern, text) + text_chunks = text.split(pattern) + final_text = "" + for i in range(len(audio_tags)): + audio_placeholder = audio_inputs["audio_phs"][0][idx] + final_text = final_text + text_chunks[i] + audio_placeholder + idx += 1 + + final_text += text_chunks[-1] + messages[index]["content"] = final_text + + return messages + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["MMProcessor"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + # image bound + image_bounds_list = [] + valid_image_nums_ls = [] + for i, input_ids in enumerate(batch_ids): + input_ids_ = torch.tensor(input_ids) + start_cond = (input_ids_ == processor.tokenizer.im_start_id) | ( + input_ids_ == processor.tokenizer.slice_start_id + ) + end_cond = (input_ids_ == processor.tokenizer.im_end_id) | (input_ids_ == processor.tokenizer.slice_end_id) + image_start_tokens = torch.where(start_cond)[0] + image_start_tokens += 1 + image_end_tokens = torch.where(end_cond)[0] + valid_image_nums_ls.append(imglens[i]) + image_bounds = torch.hstack( + [ + image_start_tokens.unsqueeze(-1), + image_end_tokens.unsqueeze(-1), + ] + ) + image_bounds_list.append(image_bounds) + + mm_inputs = self._get_mm_inputs(images, videos, [], processor, valid_image_nums_ls=valid_image_nums_ls) + if "tgt_sizes" not in mm_inputs: + dummy_data = [torch.empty(0) for _ in range(len(batch_ids))] + mm_inputs.update({"tgt_sizes": dummy_data, "pixel_values": dummy_data, "image_sizes": dummy_data}) + + mm_inputs.update({"image_bound": image_bounds_list}) + + if len(audios) > 0: + # audio bound + audio_bounds_ls = [] + spk_bounds_ls = [] + valid_audio_nums_ls = [] + + for input_ids, audiolen in zip(batch_ids, audlens): + input_ids_ = torch.tensor(input_ids) + audio_start_idx = torch.where(input_ids_ == processor.tokenizer.audio_start_id)[0] + audio_end_idx = torch.where(input_ids_ == processor.tokenizer.audio_end_id)[0] + assert len(audio_start_idx) == len(audio_end_idx) + audio_bounds = torch.hstack([(audio_start_idx + 1).unsqueeze(-1), audio_end_idx.unsqueeze(-1)]) + audio_bounds_ls.append(audio_bounds) + valid_audio_nums_ls.append(audiolen) + + spk_start_idx = torch.where(input_ids_ == processor.tokenizer.spk_start_id)[0] + spk_end_idx = torch.where(input_ids_ == processor.tokenizer.spk_end_id)[0] + assert len(spk_start_idx) == len(spk_end_idx) + spk_bounds = torch.hstack([(spk_start_idx + 1).unsqueeze(-1), spk_end_idx.unsqueeze(-1)]) + spk_bounds_ls.append(spk_bounds) + + audio_inputs = self._get_mm_inputs([], [], audios, processor, valid_audio_nums_ls=valid_audio_nums_ls) + mm_inputs.update(audio_inputs) + mm_inputs.update({"audio_bounds": audio_bounds_ls, "spk_bounds": spk_bounds_ls}) + + return mm_inputs + + +@dataclass +class MiniCPMV4_6Plugin(BasePlugin): + """Plugin for MiniCPM-V-4.6 with new transformers (NaViT vision + get_placeholder_mask API).""" + + def _get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "MMProcessor", + **kwargs, + ) -> dict[str, "torch.Tensor"]: + image_processor = getattr(processor, "image_processor") + video_processor = getattr(processor, "video_processor", None) + mm_inputs = {} + + if len(images) != 0: + # The image_processor ignores downsample_mode; target_sizes are always based on patch_size. + # downsample_mode only affects the token divisor in _build_v4_6_placeholder and model forward. + mm_inputs.update(image_processor(images, return_tensors="pt")) + + if len(videos) != 0: + if video_processor is not None: + video_inputs = video_processor(videos, return_tensors="pt") + mm_inputs["pixel_values_videos"] = video_inputs["pixel_values_videos"] + mm_inputs["target_sizes_videos"] = video_inputs["target_sizes_videos"] + else: + video_inputs = image_processor(videos, return_tensors="pt") + mm_inputs["pixel_values_videos"] = video_inputs["pixel_values"] + mm_inputs["target_sizes_videos"] = video_inputs["target_sizes"] + + if len(audios) != 0: + audio_features, audio_feature_lens, audio_phs = processor.audio_feature_extract( + [audios], + chunk_input=True, + sampling_rate=getattr(processor, "audio_sampling_rate", 16000), + ) + audio_feature_lens = [ + x.clone().detach() if isinstance(x, torch.Tensor) else torch.tensor(x) for x in audio_feature_lens + ] + mm_inputs.update({"audio_features": audio_features, "audio_feature_lens": audio_feature_lens}) + if kwargs.get("ret_phs", False): + mm_inputs.update({"audio_phs": audio_phs}) + + return mm_inputs + + def _build_v4_6_placeholder( + self, + image_inputs: dict[str, Any], + image_idx: int, + use_image_id: bool, + processor: "MMProcessor", + ) -> str: + """Build image placeholder for MiniCPM-V-4.6 using NaViT token count computation.""" + grids = image_inputs.get("grids", [[0, 0]]) + num_patches_per_image = image_inputs.get("num_patches_per_image", [1]) + target_sizes = image_inputs.get("target_sizes") + + downsample_mode = os.getenv("DOWNSAMPLE_MODE") + if downsample_mode is None: + image_processor = getattr(processor, "image_processor") + downsample_mode = getattr(image_processor, "downsample_mode", "16x") + token_divisor = 4 if downsample_mode == "4x" else 16 + + flat_index = 0 + for idx in range(image_idx): + flat_index += num_patches_per_image[idx] + n_patches = num_patches_per_image[image_idx] + + img_target_sizes = target_sizes[flat_index : flat_index + n_patches] + num_tokens_per_patch = img_target_sizes.prod(-1) // token_divisor + num_rows, num_cols = grids[image_idx] + + image_start = getattr(processor, "image_start_token", "") + image_end = getattr(processor, "image_end_token", "") + slice_start = getattr(processor, "slice_start_token", "") + slice_end = getattr(processor, "slice_end_token", "") + image_id_start = getattr(processor, "image_id_start_token", "") + image_id_end = getattr(processor, "image_id_end_token", "") + image_token = ( + getattr(processor, "image_token", None) + or getattr(getattr(processor, "tokenizer", None), "image_token", None) + or "" + ) + + image_placeholder = image_start + "<|ph|>" * int(num_tokens_per_patch[0]) + image_end + if use_image_id: + image_placeholder = f"{image_id_start}{image_idx}{image_id_end}" + image_placeholder + + slice_mode = getattr(processor, "slice_mode", True) + if slice_mode and num_rows > 0 and num_cols > 0: + per_slice_tokens = int(num_tokens_per_patch[1]) if len(num_tokens_per_patch) > 1 else 0 + slice_placeholder = slice_start + "<|ph|>" * per_slice_tokens + slice_end + slices = [slice_placeholder * num_cols for _ in range(num_rows)] + image_placeholder += "\n".join(slices) + + return image_placeholder.replace("<|ph|>", image_token) + + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens, num_video_tokens, num_audio_tokens = 0, 0, 0 + messages = deepcopy(messages) + mm_inputs, audio_inputs = {}, {} + if len(images) != 0 and len(videos) != 0: + raise ValueError("MiniCPM-V model does not support input images and videos at the same time.") + + use_image_id = getattr(processor, "default_use_image_id", True) + + if len(videos) != 0: + use_image_id = False + mm_inputs = self._get_mm_inputs([], videos, [], processor) + + for i, message in enumerate(messages): + content = message["content"] + while IMAGE_PLACEHOLDER in content: + content = content.replace(IMAGE_PLACEHOLDER, "{{image}}", 1) + num_image_tokens += 1 + + while VIDEO_PLACEHOLDER in content: + num_frames = 1 + if "num_frames_per_video" in mm_inputs: + num_frames = sum(mm_inputs["num_frames_per_video"]) + content = content.replace(VIDEO_PLACEHOLDER, "{{image}}" * num_frames, 1) + num_video_tokens += 1 + + while AUDIO_PLACEHOLDER in content: + content = content.replace(AUDIO_PLACEHOLDER, "{{audio}}", 1) + num_audio_tokens += 1 + + message["content"] = content.replace("{{image}}", "(./)").replace( + "{{audio}}", "()" + ) + + if len(images): + mm_inputs = self._get_mm_inputs(images, [], [], processor) + + if len(audios): + audio_inputs = self._get_mm_inputs([], [], audios, processor, ret_phs=True) + + if self.expand_mm_tokens and mm_inputs: + pattern = "(./)" + idx = 0 + for index, message in enumerate(messages): + text = message["content"] + image_tags = re.findall(pattern, text) + text_chunks = text.split(pattern) + final_text = "" + for i in range(len(image_tags)): + image_placeholder = self._build_v4_6_placeholder(mm_inputs, idx, use_image_id, processor) + final_text = final_text + text_chunks[i] + image_placeholder + idx += 1 + final_text += text_chunks[-1] + messages[index]["content"] = final_text + + if self.expand_mm_tokens and audio_inputs: + pattern = "()" + idx = 0 + for index, message in enumerate(messages): + text = message["content"] + audio_tags = re.findall(pattern, text) + text_chunks = text.split(pattern) + final_text = "" + for i in range(len(audio_tags)): + audio_placeholder = audio_inputs["audio_phs"][0][idx] + final_text = final_text + text_chunks[i] + audio_placeholder + idx += 1 + final_text += text_chunks[-1] + messages[index]["content"] = final_text + + return messages + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["MMProcessor"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + + # v4.6 does NOT use image_bound — the model finds image tokens via get_placeholder_mask + # Ensure target_sizes key name matches the model's expected input + if "target_sizes" not in mm_inputs and "tgt_sizes" in mm_inputs: + mm_inputs["target_sizes"] = mm_inputs.pop("tgt_sizes") + + if "target_sizes" not in mm_inputs: + mm_inputs["target_sizes"] = torch.empty(0, 2, dtype=torch.int32) + + if "pixel_values" not in mm_inputs: + mm_inputs["pixel_values"] = torch.empty(1, 3, 14, 0) + + # Pass downsample_mode to model forward so it matches the placeholder divisor + _ds = os.getenv("DOWNSAMPLE_MODE") + if _ds is None: + _ds = getattr(getattr(processor, "image_processor", None), "downsample_mode", "16x") + mm_inputs["downsample_mode"] = _ds + + if len(audios) > 0: + audio_inputs = self._get_mm_inputs([], [], audios, processor) + mm_inputs.update(audio_inputs) + + return mm_inputs + + +@dataclass +class MllamaPlugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens = 0 + messages = deepcopy(messages) + for message in messages: + content = message["content"] + num_image_tokens += content.count(IMAGE_PLACEHOLDER) + message["content"] = content.replace(IMAGE_PLACEHOLDER, self.image_token) + + return messages + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["MMProcessor"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + mm_inputs = self._get_mm_inputs(images, videos, audios, processor, imglens) + if mm_inputs: + num_tiles = mm_inputs.pop("num_tiles") + image_token_id: int = getattr(processor, "image_token_id") + max_image_tiles: int = getattr(processor.image_processor, "max_image_tiles") + cross_attention_token_mask = [ + get_cross_attention_token_mask(input_ids, image_token_id) for input_ids in batch_ids + ] + mm_inputs["cross_attention_mask"] = torch.from_numpy( + convert_sparse_cross_attention_mask_to_dense( + cross_attention_token_mask, + num_tiles=num_tiles, + max_num_tiles=max_image_tiles, + length=max(len(input_ids) for input_ids in batch_ids), + ) + ) # shape: (batch_size, length, max_num_images, max_num_tiles) + + return mm_inputs + + +@dataclass +class PaliGemmaPlugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens = 0 + messages = deepcopy(messages) + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + content = content.replace(IMAGE_PLACEHOLDER, "", 1) + num_image_tokens += 1 + + message["content"] = content + + return messages + + @override + def process_token_ids( + self, + input_ids: list[int], + labels: list[int] | None, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + tokenizer: "PreTrainedTokenizer", + processor: Optional["MMProcessor"], + ) -> tuple[list[int], list[int] | None]: + self._validate_input(processor, images, videos, audios) + num_images = len(images) + image_seqlen = processor.image_seq_length if self.expand_mm_tokens else 0 # skip mm token + image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) + input_ids = [image_token_id] * num_images * image_seqlen + input_ids + if labels is not None: + labels = [IGNORE_INDEX] * num_images * image_seqlen + labels + + return input_ids, labels + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["MMProcessor"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + seqlens = [len(input_ids) for input_ids in batch_ids] + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + mm_inputs["token_type_ids"] = _get_paligemma_token_type_ids(imglens, seqlens, processor) + return mm_inputs + + +@dataclass +class PixtralPlugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + messages = deepcopy(messages) + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + if "pixel_values" in mm_inputs: + # BC for transformers < 4.49.0 + if isinstance(mm_inputs["image_sizes"], list): + image_sizes = iter(mm_inputs["image_sizes"][0]) + else: + image_sizes = iter(mm_inputs["image_sizes"].tolist()) + + image_break_token: str = getattr(processor, "image_break_token") + image_end_token: str = getattr(processor, "image_end_token") + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + if self.expand_mm_tokens: + patch_size = processor.patch_size * getattr(processor, "spatial_merge_size", 1) + height, width = next(image_sizes) + num_height_tokens = height // patch_size + num_width_tokens = width // patch_size + replace_tokens = [[self.image_token] * num_width_tokens + [image_break_token]] * num_height_tokens + replace_tokens = [item for sublist in replace_tokens for item in sublist] # flatten list + replace_tokens[-1] = image_end_token + replace_str = "".join(replace_tokens) + else: + replace_str = self.image_token + + content = content.replace(IMAGE_PLACEHOLDER, replace_str, 1) + + message["content"] = content + + return messages + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["MMProcessor"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + # ref to this commit https://github.com/huggingface/transformers/pull/35122 + # after transformers 4.49.0, the `image_sizes` is mandatory as an input parameter for Pixtral VisionEncoder forwarding. + # it can be passed into `LlavaConditionalGeneration` as a parameter. + if not is_transformers_version_greater_than("4.49.0"): + mm_inputs.pop("image_sizes", None) + return mm_inputs + + +@dataclass +class Qwen2AudioPlugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + bos_token: str = getattr(processor, "audio_bos_token") + eos_token: str = getattr(processor, "audio_eos_token") + messages = deepcopy(messages) + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs([], [], audios, processor) + if "feature_attention_mask" in mm_inputs: + audio_lengths = mm_inputs["feature_attention_mask"].sum(-1).tolist() + + for message in messages: + content = message["content"] + while AUDIO_PLACEHOLDER in content: + if self.expand_mm_tokens: + audio_length = audio_lengths.pop(0) + input_length = (audio_length - 1) // 2 + 1 + audio_seqlen = (input_length - 2) // 2 + 1 + else: + audio_seqlen = 1 + + content = content.replace( + AUDIO_PLACEHOLDER, f"{bos_token}{self.audio_token * audio_seqlen}{eos_token}", 1 + ) + + message["content"] = content + + return messages + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["MMProcessor"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + return self._get_mm_inputs(images, videos, audios, processor) + + +@dataclass +class Qwen2VLPlugin(BasePlugin): + vision_bos_token: str = "<|vision_start|>" + vision_eos_token: str = "<|vision_end|>" + + @override + def _preprocess_image(self, image: "ImageObject", **kwargs) -> "ImageObject": + image = super()._preprocess_image(image, **kwargs) + if min(image.width, image.height) < 28: + width, height = max(image.width, 28), max(image.height, 28) + image = image.resize((width, height)) + + if image.width / image.height > 200: + width, height = image.height * 180, image.height + image = image.resize((width, height)) + + if image.height / image.width > 200: + width, height = image.width, image.width * 180 + image = image.resize((width, height)) + + return image + + @override + def _regularize_videos(self, videos: list["VideoInput"], **kwargs) -> "RegularizedVideoOutput": + results, fps_per_video, durations, frames_indices = [], [], [], [] + for video in videos: + frames: list[ImageObject] = [] + if _check_video_is_nested_images(video): + # we assume already sample frames from videos + for frame in video: + if not is_valid_image(frame) and not isinstance(frame, dict) and not os.path.exists(frame): + raise ValueError("Invalid image found in video frames.") + + frames = video + fps_per_video.append(kwargs.get("video_fps", 2.0)) + durations.append(len(frames) / kwargs.get("video_fps", 2.0)) + frames_indices.append(list(range(len(frames)))) + else: + container = av.open(video, "r") + video_stream = next(stream for stream in container.streams if stream.type == "video") + sample_indices = self._get_video_sample_indices(video_stream, **kwargs) + original_fps = float(video_stream.average_rate) + # for qwen3vl video timestamp calculation + frames_indices.append( + [idx / original_fps * kwargs.get("video_fps", 2.0) for idx in sample_indices] + ) # hack usage when do_sample_frames=False + container.seek(0) + for frame_idx, frame in enumerate(container.decode(video_stream)): + if frame_idx in sample_indices: + frames.append(frame.to_image()) + + if video_stream.duration is None: + fps_per_video.append(kwargs.get("video_fps", 2.0)) + durations.append(len(frames) / kwargs.get("video_fps", 2.0)) + else: + fps_per_video.append(len(sample_indices) / float(video_stream.duration * video_stream.time_base)) + durations.append(float(video_stream.duration * video_stream.time_base)) + + if len(frames) % 2 != 0: + frames.append(frames[-1]) + + frames = self._regularize_images(frames, **kwargs)["images"] + results.append(frames) + + return { + "videos": results, + "fps_per_video": fps_per_video, + "durations": durations, + "frames_indices": frames_indices, + } + + def _get_qwen_video_size_after_regularization( + self, width: int, height: int, image_max_pixels: int, image_min_pixels: int + ) -> tuple[int, int]: + r"""Compute the frame size produced by Qwen-VL image regularization.""" + if (width * height) > image_max_pixels: + resize_factor = math.sqrt(image_max_pixels / (width * height)) + width, height = int(width * resize_factor), int(height * resize_factor) + + if (width * height) < image_min_pixels: + resize_factor = math.sqrt(image_min_pixels / (width * height)) + width, height = int(width * resize_factor), int(height * resize_factor) + + if min(width, height) < 28: + width, height = max(width, 28), max(height, 28) + + if width / height > 200: + width, height = height * 180, height + + if height / width > 200: + width, height = width, width * 180 + + return width, height + + def _get_qwen_video_stream_metadata( + self, + video: "VideoInput", + video_fps: float, + video_maxlen: int, + ) -> Optional[dict[str, Any]]: + if not is_pyav_available() or not isinstance(video, (str, os.PathLike)): + return None + + try: + container = av.open(video, "r") + except (av.FFmpegError, OSError): + return None + + try: + video_stream = next((stream for stream in container.streams if stream.type == "video"), None) + if video_stream is None: + return None + + if video_stream.duration is None or video_stream.average_rate is None: + return None + + average_fps = float(video_stream.average_rate) + if average_fps <= 0: + return None + + sample_indices = self._get_video_sample_indices( + video_stream, video_fps=video_fps, video_maxlen=video_maxlen + ) + return { + "width": video_stream.width, + "height": video_stream.height, + "average_fps": average_fps, + "sample_indices": sample_indices, + } + finally: + container.close() + + def _get_qwen_video_resize( + self, + num_frames: int, + height: int, + width: int, + patch_size: int, + temporal_patch_size: int, + merge_size: int, + min_pixels: int, + max_pixels: int, + ) -> tuple[int, int]: + from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize + + return smart_resize( + height=height, + width=width, + factor=patch_size * merge_size, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + + def _get_qwen_video_grid_metadata( + self, + videos: list["VideoInput"], + processor: "MMProcessor", + ) -> Optional[dict[str, Any]]: + if len(videos) == 0: + return {"video_grid_thw": torch.empty((0, 3), dtype=torch.long), "frames_indices": [], "fps": 2.0} + + image_processor: BaseImageProcessor = getattr(processor, "image_processor", None) + video_processor: BaseVideoProcessor = getattr(processor, "video_processor", None) or image_processor + if image_processor is None or video_processor is None: + return None + + patch_size = getattr(video_processor, "patch_size", None) + temporal_patch_size = getattr(video_processor, "temporal_patch_size", None) + merge_size = getattr(video_processor, "merge_size", None) + size = getattr(video_processor, "size", None) + if patch_size is None or temporal_patch_size is None or merge_size is None or size is None: + return None + + if isinstance(size, dict): + min_pixels = size.get("shortest_edge") + max_pixels = size.get("longest_edge") + else: + min_pixels = getattr(size, "shortest_edge", None) + max_pixels = getattr(size, "longest_edge", None) + + if min_pixels is None or max_pixels is None: + return None + + video_fps = getattr(processor, "video_fps", 2.0) + video_maxlen = getattr(processor, "video_maxlen", 128) + image_max_pixels = getattr(processor, "video_max_pixels", 256 * 256) + image_min_pixels = getattr(processor, "video_min_pixels", 16 * 16) + + video_grid_thw = [] + frames_indices = [] + for video in videos: + metadata = self._get_qwen_video_stream_metadata(video, video_fps, video_maxlen) + if metadata is None: + return None + + width, height = self._get_qwen_video_size_after_regularization( + metadata["width"], metadata["height"], image_max_pixels, image_min_pixels + ) + num_frames = len(metadata["sample_indices"]) + if num_frames % 2 != 0: + num_frames += 1 + + resized_size = self._get_qwen_video_resize( + num_frames, + height, + width, + patch_size, + temporal_patch_size, + merge_size, + min_pixels, + max_pixels, + ) + + resized_height, resized_width = resized_size + video_grid_thw.append( + [ + math.ceil(num_frames / temporal_patch_size), + resized_height // patch_size, + resized_width // patch_size, + ] + ) + frames_indices.append([idx / metadata["average_fps"] * video_fps for idx in metadata["sample_indices"]]) + + return { + "video_grid_thw": torch.tensor(video_grid_thw, dtype=torch.long), + "frames_indices": frames_indices, + "fps": video_fps, + } + + @override + def _get_video_token_metadata( + self, + videos: list["VideoInput"], + processor: "MMProcessor", + ) -> Optional[dict[str, Any]]: + video_metadata = self._get_qwen_video_grid_metadata(videos, processor) + if video_metadata is None: + return None + + return {"video_grid_thw": video_metadata["video_grid_thw"]} + + def _get_mm_token_metadata( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "MMProcessor", + ) -> Optional[dict[str, Any]]: + if len(audios) != 0: + return None + + mm_inputs = {} + if len(images) != 0: + mm_inputs.update(self._get_mm_inputs(images, [], [], processor)) + + if len(videos) != 0: + video_inputs = self._get_video_token_metadata(videos, processor) + if video_inputs is None: + return None + + mm_inputs.update(video_inputs) + + return mm_inputs + + @override + def _get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "MMProcessor", + ) -> dict[str, "torch.Tensor"]: + image_processor: BaseImageProcessor = getattr(processor, "image_processor", None) + video_processor: BaseVideoProcessor = getattr(processor, "video_processor", None) + mm_inputs = {} + if len(images) != 0: + images = self._regularize_images( + images, + image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768), + image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32), + )["images"] + mm_inputs.update(image_processor(images, return_tensors="pt")) + + if len(videos) != 0: + video_data = self._regularize_videos( + videos, + image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256), + image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16), + video_fps=getattr(processor, "video_fps", 2.0), + video_maxlen=getattr(processor, "video_maxlen", 128), + ) + mm_inputs.update(video_processor(videos=video_data["videos"], return_tensors="pt")) + temporal_patch_size: int = getattr(image_processor, "temporal_patch_size", 2) + if "second_per_grid_ts" in processor.model_input_names: + mm_inputs["second_per_grid_ts"] = [temporal_patch_size / fps for fps in video_data["fps_per_video"]] + + return mm_inputs + + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens, num_video_tokens = 0, 0 + messages = deepcopy(messages) + image_processor: BaseImageProcessor = getattr(processor, "image_processor") + + merge_length: int = getattr(image_processor, "merge_size") ** 2 + if self.expand_mm_tokens: + mm_inputs = self._get_mm_token_metadata(images, videos, audios, processor) + if mm_inputs is None: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + + image_grid_thw = mm_inputs.get("image_grid_thw", []) + video_grid_thw = mm_inputs.get("video_grid_thw", []) + else: + image_grid_thw = [None] * len(images) + video_grid_thw = [None] * len(videos) + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + image_seqlen = image_grid_thw[num_image_tokens].prod() // merge_length if self.expand_mm_tokens else 1 + content = content.replace( + IMAGE_PLACEHOLDER, + f"{self.vision_bos_token}{self.image_token * image_seqlen}{self.vision_eos_token}", + 1, + ) + num_image_tokens += 1 + + while VIDEO_PLACEHOLDER in content: + video_seqlen = video_grid_thw[num_video_tokens].prod() // merge_length if self.expand_mm_tokens else 1 + content = content.replace( + VIDEO_PLACEHOLDER, + f"{self.vision_bos_token}{self.video_token * video_seqlen}{self.vision_eos_token}", + 1, + ) + num_video_tokens += 1 + + message["content"] = content + + return messages + + +@dataclass +class Qwen3VLPlugin(Qwen2VLPlugin): + @override + def _get_qwen_video_resize( + self, + num_frames: int, + height: int, + width: int, + patch_size: int, + temporal_patch_size: int, + merge_size: int, + min_pixels: int, + max_pixels: int, + ) -> tuple[int, int]: + from transformers.models.qwen3_vl.video_processing_qwen3_vl import smart_resize + + return smart_resize( + num_frames=num_frames, + height=height, + width=width, + temporal_factor=temporal_patch_size, + factor=patch_size * merge_size, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + + @override + def _get_video_token_metadata( + self, + videos: list["VideoInput"], + processor: "MMProcessor", + ) -> Optional[dict[str, Any]]: + video_metadata = self._get_qwen_video_grid_metadata(videos, processor) + if video_metadata is None: + return None + + return { + "video_grid_thw": video_metadata["video_grid_thw"], + "video_metadata": [ + SimpleNamespace( + frames_indices=frames_indices, + fps=video_metadata["fps"], + ) + for frames_indices in video_metadata["frames_indices"] + ], + } + + @override + def _get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "MMProcessor", + ) -> dict[str, "torch.Tensor"]: + image_processor: BaseImageProcessor = getattr(processor, "image_processor", None) + video_processor: BaseImageProcessor = getattr(processor, "video_processor", None) + mm_inputs = {} + if len(images) != 0: + images = self._regularize_images( + images, + image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768), + image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32), + )["images"] + mm_inputs.update(image_processor(images, return_tensors="pt")) + + if len(videos) != 0: + videos = self._regularize_videos( + videos, + image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256), + image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16), + video_fps=getattr(processor, "video_fps", 2.0), + video_maxlen=getattr(processor, "video_maxlen", 128), + ) + video_metadata = [ + { + "fps": getattr(processor, "video_fps", 2.0), + "duration": duration, + "total_num_frames": len(video), + "frames_indices": sample_indices, + } + for video, duration, sample_indices in zip( + videos["videos"], videos["durations"], videos["frames_indices"] + ) + ] + mm_inputs.update( + video_processor( + videos=videos["videos"], + video_metadata=video_metadata, + fps=getattr(processor, "video_fps", 2.0), + return_metadata=True, + do_sample_frames=False, # avoid changing frames_indices + ) + ) + temporal_patch_size: int = getattr(image_processor, "temporal_patch_size", 2) + if "second_per_grid_ts" in processor.model_input_names: + mm_inputs["second_per_grid_ts"] = [temporal_patch_size / fps for fps in videos["fps_per_video"]] + + return mm_inputs + + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens, num_video_tokens = 0, 0 + messages = deepcopy(messages) + image_processor: BaseImageProcessor = getattr(processor, "image_processor") + video_processor: BaseImageProcessor = getattr(processor, "video_processor") + + image_merge_length: int = getattr(image_processor, "merge_size") ** 2 + video_merge_length: int = getattr(video_processor, "merge_size") ** 2 + if self.expand_mm_tokens: + mm_inputs = self._get_mm_token_metadata(images, videos, audios, processor) + if mm_inputs is None: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + + image_grid_thw = mm_inputs.get("image_grid_thw", []) + video_grid_thw = mm_inputs.get("video_grid_thw", []) + video_metadata = mm_inputs.get("video_metadata", []) + + else: + image_grid_thw = [None] * len(images) + video_grid_thw = [None] * len(videos) + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + image_seqlen = ( + image_grid_thw[num_image_tokens].prod() // image_merge_length if self.expand_mm_tokens else 1 + ) + content = content.replace( + IMAGE_PLACEHOLDER, + f"{self.vision_bos_token}{self.image_token * image_seqlen}{self.vision_eos_token}", + 1, + ) + num_image_tokens += 1 + + while VIDEO_PLACEHOLDER in content: + if self.expand_mm_tokens: + video_grid = video_grid_thw[num_video_tokens] + num_frames = int(video_grid[0].item()) + metadata = video_metadata[num_video_tokens] + timestamps = processor._calculate_timestamps( + metadata.frames_indices, + metadata.fps, + getattr(video_processor, "temporal_patch_size", 2), + ) + video_structure = "" + video_seqlen = int((video_grid[1:].prod() // video_merge_length).item()) + for frame_index in range(num_frames): + timestamp_sec = timestamps[frame_index] + frame_structure = ( + f"<{timestamp_sec:.1f} seconds>" + f"{self.vision_bos_token}{self.video_token * video_seqlen}{self.vision_eos_token}" + ) + video_structure += frame_structure + else: + video_structure = f"{self.vision_bos_token}{self.video_token}{self.vision_eos_token}" + + content = content.replace(VIDEO_PLACEHOLDER, video_structure, 1) + num_video_tokens += 1 + + message["content"] = content + + return messages + + +@dataclass +class GLM4VPlugin(Qwen2VLPlugin): + @override + def _get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "MMProcessor", + ) -> dict[str, "torch.Tensor"]: + image_processor: BaseImageProcessor = getattr(processor, "image_processor", None) + video_processor: BaseImageProcessor = getattr(processor, "video_processor", None) + mm_inputs = {} + if len(images) != 0: + images = self._regularize_images( + images, + image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768), + image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32), + )["images"] + mm_inputs.update(image_processor(images, return_tensors="pt")) + + if len(videos) != 0: + video_data = self._regularize_videos( + videos, + image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256), + image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16), + video_fps=getattr(processor, "video_fps", 2.0), + video_maxlen=getattr(processor, "video_maxlen", 128), + ) + # prepare video metadata + video_metadata = [ + {"fps": 2, "duration": duration, "total_frames": len(video)} + for video, duration in zip(video_data["videos"], video_data["durations"]) + ] + mm_inputs.update(video_processor(images=None, videos=video_data["videos"], video_metadata=video_metadata)) + + return mm_inputs + + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens, num_video_tokens = 0, 0 + messages = deepcopy(messages) + image_processor: BaseImageProcessor = getattr(processor, "image_processor") + + merge_length: int = getattr(image_processor, "merge_size") ** 2 + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + image_grid_thw = mm_inputs.get("image_grid_thw", []) + video_grid_thw = mm_inputs.get("video_grid_thw", []) + num_frames = video_grid_thw[0][0] if len(video_grid_thw) > 0 else 0 # hard code for now + timestamps = mm_inputs.get("timestamps", []) + + if hasattr(timestamps, "tolist"): + timestamps = timestamps.tolist() + + if not timestamps: + timestamps_list = [] + elif isinstance(timestamps[0], list): + timestamps_list = timestamps[0] + else: + timestamps_list = timestamps + + unique_timestamps = timestamps_list.copy() + selected_timestamps = unique_timestamps[:num_frames] + while len(selected_timestamps) < num_frames: + selected_timestamps.append(selected_timestamps[-1] if selected_timestamps else 0) + + else: + image_grid_thw = [None] * len(images) + video_grid_thw = [None] * len(videos) + num_frames = 0 + selected_timestamps = [0] + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + image_seqlen = image_grid_thw[num_image_tokens].prod() // merge_length if self.expand_mm_tokens else 1 + content = content.replace( + IMAGE_PLACEHOLDER, f"<|begin_of_image|>{self.image_token * image_seqlen}<|end_of_image|>", 1 + ) + num_image_tokens += 1 + + while VIDEO_PLACEHOLDER in content: + video_structure = "" + for frame_index in range(num_frames): + video_seqlen = ( + video_grid_thw[num_video_tokens][1:].prod() // merge_length if self.expand_mm_tokens else 1 + ) + timestamp_sec = selected_timestamps[frame_index] + frame_structure = ( + f"<|begin_of_image|>{self.image_token * video_seqlen}<|end_of_image|>{timestamp_sec}" + ) + video_structure += frame_structure + + if not self.expand_mm_tokens: + video_structure = self.video_token + + content = content.replace(VIDEO_PLACEHOLDER, f"<|begin_of_video|>{video_structure}<|end_of_video|>", 1) + num_video_tokens += 1 + + message["content"] = content + + return messages + + @override + def get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + imglens: list[int], + vidlens: list[int], + audlens: list[int], + batch_ids: list[list[int]], + processor: Optional["ProcessorMixin"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(processor, images, videos, audios) + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + mm_inputs.pop("timestamps", None) + return mm_inputs + + +@dataclass +class Qwen2OmniPlugin(Qwen2VLPlugin): + audio_bos_token: str = "<|audio_start|>" + audio_eos_token: str = "<|audio_end|>" + + @override + def _get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "MMProcessor", + ) -> dict[str, "torch.Tensor"]: + image_processor: BaseImageProcessor = getattr(processor, "image_processor", None) + video_processor: BaseVideoProcessor = getattr(processor, "video_processor", None) + feature_extractor: SequenceFeatureExtractor = getattr(processor, "feature_extractor", None) or getattr( + processor, "audio_processor", None + ) + mm_inputs = {} + if len(images) != 0: + images = self._regularize_images( + images, + image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768), + image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32), + )["images"] + mm_inputs.update(image_processor(images, return_tensors="pt")) + + if len(videos) != 0: + video_dict = self._regularize_videos( + videos, + image_max_pixels=getattr(processor, "video_max_pixels", 256 * 256), + image_min_pixels=getattr(processor, "video_min_pixels", 16 * 16), + video_fps=getattr(processor, "video_fps", 2.0), + video_maxlen=getattr(processor, "video_maxlen", 128), + ) + mm_inputs.update(video_processor(videos=video_dict["videos"], return_tensors="pt")) + temporal_patch_size: int = getattr(image_processor, "temporal_patch_size", 2) + mm_inputs["video_second_per_grid"] = torch.tensor( + [temporal_patch_size / fps for fps in video_dict["fps_per_video"]] + ) + + if len(audios) != 0: + audios = self._regularize_audios( + audios, + sampling_rate=getattr(processor, "audio_sampling_rate", 16000), + )["audios"] + mm_inputs.update( + feature_extractor( + audios, + sampling_rate=getattr(processor, "audio_sampling_rate", 16000), + return_attention_mask=True, + padding="max_length", + return_tensors="pt", + ) + ) + mm_inputs["feature_attention_mask"] = mm_inputs.pop("attention_mask") # prevent conflicts + + return mm_inputs + + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens, num_video_tokens, num_audio_tokens = 0, 0, 0 + messages = deepcopy(messages) + image_processor: BaseImageProcessor = getattr(processor, "image_processor", None) + + merge_length = processor.image_processor.merge_size**2 + use_audio_in_video = getattr(processor, "use_audio_in_video", False) + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + image_grid_thw = mm_inputs.get("image_grid_thw", []) + video_grid_thw = mm_inputs.get("video_grid_thw", []) + if "feature_attention_mask" in mm_inputs: + if processor.__class__.__name__ == "Qwen3OmniMoeProcessor": # for qwen3omni + input_lengths = mm_inputs["feature_attention_mask"].sum(-1) + input_lengths_leave = input_lengths % 100 + feature_lengths = (input_lengths_leave - 1) // 2 + 1 + audio_lengths = ((feature_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 + else: + input_lengths = (mm_inputs["feature_attention_mask"].sum(-1).numpy() - 1) // 2 + 1 + audio_lengths = (input_lengths - 2) // 2 + 1 + else: + mm_inputs = {} + image_grid_thw = [None] * len(images) + video_grid_thw = [None] * len(videos) + audio_lengths = [None] * len(audios) + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + image_seqlen = image_grid_thw[num_image_tokens].prod() // merge_length if self.expand_mm_tokens else 1 + content = content.replace( + IMAGE_PLACEHOLDER, + f"{self.vision_bos_token}{self.image_token * image_seqlen}{self.vision_eos_token}", + 1, + ) + num_image_tokens += 1 + + if ( + use_audio_in_video and len(audios) and len(videos) + ): # if use the audio of video # deal video token and audio token together + if len(videos) != len(audios): + raise ValueError( + f"Number of videos ({len(videos)}) must match number of audios ({len(audios)}) when using audio in video." + ) + + while VIDEO_PLACEHOLDER in content: + video_pos = content.find(VIDEO_PLACEHOLDER) + audio_pos = content.find(AUDIO_PLACEHOLDER, video_pos) + if audio_pos == -1 or audio_pos < video_pos: + raise ValueError( + f"Each {VIDEO_PLACEHOLDER} must be followed by an {AUDIO_PLACEHOLDER} when using audio in video." + ) + + position_id_per_seconds: int = getattr(processor, "position_id_per_seconds", 25) + audio_t_index = torch.arange(audio_lengths[num_audio_tokens]) + video_t_index = ( + torch.arange(video_grid_thw[num_video_tokens][0]) + .view(-1, 1, 1) + .expand( + -1, + video_grid_thw[num_video_tokens][1] // image_processor.merge_size, + video_grid_thw[num_video_tokens][2] // image_processor.merge_size, + ) + .flatten() + * mm_inputs["video_second_per_grid"][num_video_tokens] + * position_id_per_seconds + ).long() + t_ntoken_per_chunk = position_id_per_seconds * 2 + video_chunk_indices = processor.get_chunked_index(video_t_index, t_ntoken_per_chunk) + audio_chunk_indices = processor.get_chunked_index(audio_t_index, t_ntoken_per_chunk) + placeholder_string = "" + placeholder_string += self.vision_bos_token + self.audio_bos_token + for j in range(max(len(video_chunk_indices), len(audio_chunk_indices))): + video_chunk_index = video_chunk_indices[j] if j < len(video_chunk_indices) else None + audio_chunk_index = audio_chunk_indices[j] if j < len(audio_chunk_indices) else None + if video_chunk_index is not None: + placeholder_string += self.video_token * (video_chunk_index[1] - video_chunk_index[0]) + + if audio_chunk_index is not None: + placeholder_string += self.audio_token * (audio_chunk_index[1] - audio_chunk_index[0]) + + placeholder_string += self.audio_eos_token + self.vision_eos_token + content = content.replace(VIDEO_PLACEHOLDER, placeholder_string, 1) + content = content.replace(AUDIO_PLACEHOLDER, "", 1) + num_audio_tokens += 1 + num_video_tokens += 1 + else: + while AUDIO_PLACEHOLDER in content: + audio_seqlen = audio_lengths[num_audio_tokens] if self.expand_mm_tokens else 1 + content = content.replace( + AUDIO_PLACEHOLDER, + f"{self.audio_bos_token}{self.audio_token * audio_seqlen}{self.audio_eos_token}", + 1, + ) + num_audio_tokens += 1 + + while VIDEO_PLACEHOLDER in content: + video_seqlen = ( + video_grid_thw[num_video_tokens].prod() // merge_length if self.expand_mm_tokens else 1 + ) + content = content.replace( + VIDEO_PLACEHOLDER, + f"{self.vision_bos_token}{self.video_token * video_seqlen}{self.vision_eos_token}", + 1, + ) + num_video_tokens += 1 + + message["content"] = content + + return messages + + +@dataclass +class VideoLlavaPlugin(BasePlugin): + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens, num_video_tokens = 0, 0 + messages = deepcopy(messages) + num_frames = 0 + if self.expand_mm_tokens: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + if "pixel_values_images" in mm_inputs: + height, width = get_image_size(to_numpy_array(mm_inputs["pixel_values_images"][0])) + num_frames = 1 + + if "pixel_values_videos" in mm_inputs: + one_video = to_numpy_array(mm_inputs["pixel_values_videos"][0]) + height, width = get_image_size(one_video[0]) + num_frames = one_video.shape[0] # frame dim is always after batch dim + + if "pixel_values_images" in mm_inputs or "pixel_values_videos" in mm_inputs: + image_seqlen = (height // processor.patch_size) * ( + width // processor.patch_size + ) + processor.num_additional_image_tokens + video_seqlen = image_seqlen * num_frames + if processor.vision_feature_select_strategy == "default": + image_seqlen -= 1 + else: + image_seqlen, video_seqlen = 1, 1 + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + content = content.replace(IMAGE_PLACEHOLDER, "{{image}}" * image_seqlen, 1) + num_image_tokens += 1 + + while VIDEO_PLACEHOLDER in content: + content = content.replace(VIDEO_PLACEHOLDER, "{{video}}" * video_seqlen, 1) + num_video_tokens += 1 + + content = content.replace("{{image}}", self.image_token) + message["content"] = content.replace("{{video}}", self.video_token) + + return messages + + +@dataclass +class LFMVLPlugin(BasePlugin): + r"""Plugin for LFM2.5-VL vision-language models. + + LFM2.5-VL uses dynamic image token counts based on image resolution. + The image processor returns spatial_shapes tensor with [height, width] grid dimensions. + Token count per image = (spatial_h * spatial_w) / (downsample_factor^2) + """ + + @override + def _get_mm_inputs( + self, + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: "MMProcessor", + ) -> dict[str, "torch.Tensor"]: + image_processor: BaseImageProcessor = getattr(processor, "image_processor", None) + mm_inputs = {} + if len(images) != 0: + images = self._regularize_images( + images, + image_max_pixels=getattr(processor, "image_max_pixels", 768 * 768), + image_min_pixels=getattr(processor, "image_min_pixels", 32 * 32), + )["images"] + mm_inputs.update(image_processor(images, return_tensors="pt")) + return mm_inputs + + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + num_image_tokens = 0 + messages = deepcopy(messages) + image_processor: BaseImageProcessor = getattr(processor, "image_processor") + downsample_factor: int = getattr(image_processor, "downsample_factor", 2) + + if self.expand_mm_tokens and len(images) > 0: + mm_inputs = self._get_mm_inputs(images, videos, audios, processor) + spatial_shapes = mm_inputs.get("spatial_shapes", []) + else: + spatial_shapes = [] + + for message in messages: + content = message["content"] + while IMAGE_PLACEHOLDER in content: + if self.expand_mm_tokens and len(spatial_shapes) > num_image_tokens: + h, w = spatial_shapes[num_image_tokens].tolist() + image_seqlen = (h * w) // (downsample_factor * downsample_factor) + else: + image_seqlen = 1 + + content = content.replace(IMAGE_PLACEHOLDER, "{{image}}" * image_seqlen, 1) + num_image_tokens += 1 + + message["content"] = content.replace("{{image}}", self.image_token) + + return messages + + +@dataclass +class YoutuVLPlugin(BasePlugin): + r"""Plugin for Youtu-VL vision-language models.""" + + vision_bos_token: str = "<|vision_start|>" + vision_eos_token: str = "<|vision_end|>" + + @override + def process_messages( + self, + messages: list[dict[str, str]], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + processor: Optional["MMProcessor"], + ) -> list[dict[str, str]]: + self._validate_input(processor, images, videos, audios) + self._validate_messages(messages, images, videos, audios) + messages = deepcopy(messages) + + for message in messages: + content = message["content"] + content = content.replace( + IMAGE_PLACEHOLDER, f"{self.vision_bos_token}{self.image_token}{self.vision_eos_token}" + ) + content = content.replace( + VIDEO_PLACEHOLDER, f"{self.vision_bos_token}{self.video_token}{self.vision_eos_token}" + ) + + message["content"] = content + + return messages + + +PLUGINS = { + "base": BasePlugin, + "ernie_vl": ErnieVLPlugin, + "gemma3": Gemma3Plugin, + "gemma3n": Gemma3nPlugin, + "gemma4": Gemma4Plugin, + "glm4v": GLM4VPlugin, + "intern_vl": InternVLPlugin, + "kimi_vl": KimiVLPlugin, + "llama4": Llama4Plugin, + "llava": LlavaPlugin, + "llava_next": LlavaNextPlugin, + "llava_next_video": LlavaNextVideoPlugin, + "lfm2_vl": LFMVLPlugin, + "minicpm_v": MiniCPMVPlugin, + "minicpm_v_4_6": MiniCPMV4_6Plugin, + "mllama": MllamaPlugin, + "paligemma": PaliGemmaPlugin, + "pixtral": PixtralPlugin, + "qwen2_audio": Qwen2AudioPlugin, + "qwen2_omni": Qwen2OmniPlugin, + "qwen2_vl": Qwen2VLPlugin, + "qwen3_vl": Qwen3VLPlugin, + "video_llava": VideoLlavaPlugin, + "youtu_vl": YoutuVLPlugin, +} + + +def register_mm_plugin(name: str, plugin_class: type["BasePlugin"]) -> None: + r"""Register a multimodal plugin.""" + if name in PLUGINS: + raise ValueError(f"Multimodal plugin {name} already exists.") + + PLUGINS[name] = plugin_class + + +def get_mm_plugin( + name: str, + image_token: str | None = None, + video_token: str | None = None, + audio_token: str | None = None, + **kwargs, +) -> "BasePlugin": + r"""Get plugin for multimodal inputs.""" + if name not in PLUGINS: + raise ValueError(f"Multimodal plugin `{name}` not found.") + + return PLUGINS[name](image_token, video_token, audio_token, **kwargs) diff --git a/src/llamafactory/data/parser.py b/src/llamafactory/data/parser.py new file mode 100644 index 0000000000000000000000000000000000000000..5209da64954a363c59d232010111139b7b37fdfe --- /dev/null +++ b/src/llamafactory/data/parser.py @@ -0,0 +1,149 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 +from typing import Any, Literal + +from huggingface_hub import hf_hub_download + +from ..extras.constants import DATA_CONFIG +from ..extras.misc import use_modelscope, use_openmind + + +@dataclass +class DatasetAttr: + r"""Dataset attributes.""" + + # basic configs + load_from: Literal["hf_hub", "ms_hub", "om_hub", "script", "file"] + dataset_name: str + formatting: Literal["alpaca", "sharegpt", "openai"] = "alpaca" + ranking: bool = False + # extra configs + subset: str | None = None + split: str = "train" + folder: str | None = None + num_samples: int | None = None + # common columns + system: str | None = None + tools: str | None = None + images: str | None = None + videos: str | None = None + audios: str | None = None + # dpo columns + chosen: str | None = None + rejected: str | None = None + kto_tag: str | None = None + # alpaca columns + prompt: str | None = "instruction" + query: str | None = "input" + response: str | None = "output" + history: str | None = None + # sharegpt columns + messages: str | None = "conversations" + # sharegpt tags + role_tag: str | None = "from" + content_tag: str | None = "value" + user_tag: str | None = "human" + assistant_tag: str | None = "gpt" + observation_tag: str | None = "observation" + function_tag: str | None = "function_call" + system_tag: str | None = "system" + + def __repr__(self) -> str: + return self.dataset_name + + def set_attr(self, key: str, obj: dict[str, Any], default: Any | None = None) -> None: + setattr(self, key, obj.get(key, default)) + + def join(self, attr: dict[str, Any]) -> None: + self.set_attr("formatting", attr, default="alpaca") + self.set_attr("ranking", attr, default=False) + self.set_attr("subset", attr) + self.set_attr("split", attr, default="train") + self.set_attr("folder", attr) + self.set_attr("num_samples", attr) + + if "columns" in attr: + column_names = ["prompt", "query", "response", "history", "messages", "system", "tools"] + column_names += ["images", "videos", "audios", "chosen", "rejected", "kto_tag"] + for column_name in column_names: + self.set_attr(column_name, attr["columns"]) + + if "tags" in attr: + tag_names = ["role_tag", "content_tag"] + tag_names += ["user_tag", "assistant_tag", "observation_tag", "function_tag", "system_tag"] + for tag in tag_names: + self.set_attr(tag, attr["tags"]) + + +def get_dataset_list(dataset_names: list[str] | None, dataset_dir: str | dict) -> list["DatasetAttr"]: + r"""Get the attributes of the datasets.""" + if dataset_names is None: + dataset_names = [] + + if isinstance(dataset_dir, dict): + dataset_info = dataset_dir + elif dataset_dir == "ONLINE": + dataset_info = None + else: + if dataset_dir.startswith("REMOTE:"): + config_path = hf_hub_download(repo_id=dataset_dir[7:], filename=DATA_CONFIG, repo_type="dataset") + else: + config_path = os.path.join(dataset_dir, DATA_CONFIG) + + try: + with open(config_path) as f: + dataset_info = json.load(f) + except Exception as err: + if len(dataset_names) != 0: + raise ValueError(f"Cannot open {config_path} due to {str(err)}.") + + dataset_info = None + + dataset_list: list[DatasetAttr] = [] + for name in dataset_names: + if dataset_info is None: # dataset_dir is ONLINE + load_from = "ms_hub" if use_modelscope() else "om_hub" if use_openmind() else "hf_hub" + dataset_attr = DatasetAttr(load_from, dataset_name=name) + dataset_list.append(dataset_attr) + continue + + if name not in dataset_info: + raise ValueError(f"Undefined dataset {name} in {DATA_CONFIG}.") + + has_hf_url = "hf_hub_url" in dataset_info[name] + has_ms_url = "ms_hub_url" in dataset_info[name] + has_om_url = "om_hub_url" in dataset_info[name] + + if has_hf_url or has_ms_url or has_om_url: + if has_ms_url and (use_modelscope() or not has_hf_url): + dataset_attr = DatasetAttr("ms_hub", dataset_name=dataset_info[name]["ms_hub_url"]) + elif has_om_url and (use_openmind() or not has_hf_url): + dataset_attr = DatasetAttr("om_hub", dataset_name=dataset_info[name]["om_hub_url"]) + else: + dataset_attr = DatasetAttr("hf_hub", dataset_name=dataset_info[name]["hf_hub_url"]) + elif "script_url" in dataset_info[name]: + dataset_attr = DatasetAttr("script", dataset_name=dataset_info[name]["script_url"]) + elif "cloud_file_name" in dataset_info[name]: + dataset_attr = DatasetAttr("cloud_file", dataset_name=dataset_info[name]["cloud_file_name"]) + else: + dataset_attr = DatasetAttr("file", dataset_name=dataset_info[name]["file_name"]) + + dataset_attr.join(dataset_info[name]) + dataset_list.append(dataset_attr) + + return dataset_list diff --git a/src/llamafactory/data/processor/__init__.py b/src/llamafactory/data/processor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..357ab7899f9eecbd29344482d109b89af274ea2e --- /dev/null +++ b/src/llamafactory/data/processor/__init__.py @@ -0,0 +1,31 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 .feedback import FeedbackDatasetProcessor +from .pairwise import PairwiseDatasetProcessor +from .pretrain import PretrainDatasetProcessor +from .processor_utils import DatasetProcessor +from .supervised import PackedSupervisedDatasetProcessor, SupervisedDatasetProcessor +from .unsupervised import UnsupervisedDatasetProcessor + + +__all__ = [ + "DatasetProcessor", + "FeedbackDatasetProcessor", + "PackedSupervisedDatasetProcessor", + "PairwiseDatasetProcessor", + "PretrainDatasetProcessor", + "SupervisedDatasetProcessor", + "UnsupervisedDatasetProcessor", +] diff --git a/src/llamafactory/data/processor/__pycache__/__init__.cpython-312.pyc b/src/llamafactory/data/processor/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f1fd8ae7ceb80d66c70e64031dd63f62b25bc62f Binary files /dev/null and b/src/llamafactory/data/processor/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/llamafactory/data/processor/__pycache__/feedback.cpython-312.pyc b/src/llamafactory/data/processor/__pycache__/feedback.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..04b70010d8f9e5726613766979aa0924367fe8c4 Binary files /dev/null and b/src/llamafactory/data/processor/__pycache__/feedback.cpython-312.pyc differ diff --git a/src/llamafactory/data/processor/__pycache__/pairwise.cpython-312.pyc b/src/llamafactory/data/processor/__pycache__/pairwise.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9b1b7e5b835bf419c41f4dd869f32bb9485beac Binary files /dev/null and b/src/llamafactory/data/processor/__pycache__/pairwise.cpython-312.pyc differ diff --git a/src/llamafactory/data/processor/__pycache__/pretrain.cpython-312.pyc b/src/llamafactory/data/processor/__pycache__/pretrain.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05150771422ebce4e40ad598e80101837c4e0408 Binary files /dev/null and b/src/llamafactory/data/processor/__pycache__/pretrain.cpython-312.pyc differ diff --git a/src/llamafactory/data/processor/__pycache__/processor_utils.cpython-312.pyc b/src/llamafactory/data/processor/__pycache__/processor_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b787c818020177d82617c45d59602fca4aa6d038 Binary files /dev/null and b/src/llamafactory/data/processor/__pycache__/processor_utils.cpython-312.pyc differ diff --git a/src/llamafactory/data/processor/__pycache__/supervised.cpython-312.pyc b/src/llamafactory/data/processor/__pycache__/supervised.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..90d9bfa2e83ecf6df463262f390e89810ec49299 Binary files /dev/null and b/src/llamafactory/data/processor/__pycache__/supervised.cpython-312.pyc differ diff --git a/src/llamafactory/data/processor/__pycache__/unsupervised.cpython-312.pyc b/src/llamafactory/data/processor/__pycache__/unsupervised.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c95db7d6413355e8952b3525c7c1b8345e59662d Binary files /dev/null and b/src/llamafactory/data/processor/__pycache__/unsupervised.cpython-312.pyc differ diff --git a/src/llamafactory/data/processor/feedback.py b/src/llamafactory/data/processor/feedback.py new file mode 100644 index 0000000000000000000000000000000000000000..871615b9266e501f25f68e84e4536c0d24617803 --- /dev/null +++ b/src/llamafactory/data/processor/feedback.py @@ -0,0 +1,129 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 collections import defaultdict +from typing import TYPE_CHECKING, Any, Optional + +from ...extras import logging +from ...extras.constants import IGNORE_INDEX +from .processor_utils import DatasetProcessor, infer_seqlen + + +if TYPE_CHECKING: + from ..mm_plugin import AudioInput, ImageInput, VideoInput + + +logger = logging.get_logger(__name__) + + +class FeedbackDatasetProcessor(DatasetProcessor): + def _encode_data_example( + self, + prompt: list[dict[str, str]], + response: list[dict[str, str]], + kl_response: list[dict[str, str]], + system: Optional[str], + tools: Optional[str], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + ) -> tuple[list[int], list[int], list[int], list[int], bool]: + if response[0]["content"]: # desired example + kto_tag = True + messages = prompt + [response[0]] + else: # undesired example + kto_tag = False + messages = prompt + [response[1]] + + if kl_response[0]["content"]: + kl_messages = prompt + [kl_response[0]] + else: + kl_messages = prompt + [kl_response[1]] + + messages = self.template.mm_plugin.process_messages(messages, images, videos, audios, self.processor) + kl_messages = self.template.mm_plugin.process_messages(kl_messages, images, videos, audios, self.processor) + prompt_ids, response_ids = self.template.encode_oneturn(self.tokenizer, messages, system, tools) + kl_prompt_ids, kl_response_ids = self.template.encode_oneturn(self.tokenizer, kl_messages, system, tools) + + if self.template.efficient_eos: + response_ids += [self.tokenizer.eos_token_id] + kl_response_ids += [self.tokenizer.eos_token_id] + + prompt_ids, _ = self.template.mm_plugin.process_token_ids( + prompt_ids, None, images, videos, audios, self.tokenizer, self.processor + ) + kl_prompt_ids, _ = self.template.mm_plugin.process_token_ids( + kl_prompt_ids, None, images, videos, audios, self.tokenizer, self.processor + ) + + source_len, target_len = infer_seqlen(len(prompt_ids), len(response_ids), self.data_args.cutoff_len) + prompt_ids = prompt_ids[:source_len] + response_ids = response_ids[:target_len] + kl_source_len, kl_target_len = infer_seqlen( + len(kl_prompt_ids), len(kl_response_ids), self.data_args.cutoff_len + ) + kl_prompt_ids = kl_prompt_ids[:kl_source_len] + kl_response_ids = kl_response_ids[:kl_target_len] + + input_ids = prompt_ids + response_ids + labels = [IGNORE_INDEX] * source_len + response_ids + kl_input_ids = kl_prompt_ids + kl_response_ids + kl_labels = [IGNORE_INDEX] * kl_source_len + kl_response_ids + return input_ids, labels, kl_input_ids, kl_labels, kto_tag + + def preprocess_dataset(self, examples: dict[str, list[Any]]) -> dict[str, list[Any]]: + # Creates mismatched pairs of prompts and completions for the KL dataset by adding a +1 offset to the order of completions. + kl_response = [examples["_response"][-1]] + examples["_response"][:-1] + model_inputs = defaultdict(list) + for i in range(len(examples["_prompt"])): + if len(examples["_prompt"][i]) % 2 != 1 or len(examples["_response"][i]) < 2: + logger.warning_rank0( + "Dropped invalid example: {}".format(examples["_prompt"][i] + examples["_response"][i]) + ) + continue + + input_ids, labels, kl_input_ids, kl_labels, kto_tag = self._encode_data_example( + prompt=examples["_prompt"][i], + response=examples["_response"][i], + kl_response=kl_response[i], + system=examples["_system"][i], + tools=examples["_tools"][i], + images=examples["_images"][i] or [], + videos=examples["_videos"][i] or [], + audios=examples["_audios"][i] or [], + ) + model_inputs["input_ids"].append(input_ids) + model_inputs["attention_mask"].append([1] * len(input_ids)) + model_inputs["labels"].append(labels) + model_inputs["kl_input_ids"].append(kl_input_ids) + model_inputs["kl_attention_mask"].append([1] * len(kl_input_ids)) + model_inputs["kl_labels"].append(kl_labels) + model_inputs["kto_tags"].append(kto_tag) + model_inputs["images"].append(examples["_images"][i]) + model_inputs["videos"].append(examples["_videos"][i]) + model_inputs["audios"].append(examples["_audios"][i]) + + desirable_num = sum([1 for tag in model_inputs["kto_tags"] if tag]) + undesirable_num = len(model_inputs["kto_tags"]) - desirable_num + if desirable_num == 0 or undesirable_num == 0: + logger.warning_rank0("Your dataset only has one preference type.") + + return model_inputs + + def print_data_example(self, example: dict[str, list[int]]) -> None: + valid_labels = list(filter(lambda x: x != IGNORE_INDEX, example["labels"])) + print("input_ids:\n{}".format(example["input_ids"])) + print("inputs:\n{}".format(self.tokenizer.decode(example["input_ids"], skip_special_tokens=False))) + print("label_ids:\n{}".format(example["labels"])) + print(f"labels:\n{self.tokenizer.decode(valid_labels, skip_special_tokens=False)}") diff --git a/src/llamafactory/data/processor/pairwise.py b/src/llamafactory/data/processor/pairwise.py new file mode 100644 index 0000000000000000000000000000000000000000..94101deb8e75af73c1851720604994a11f2eb87d --- /dev/null +++ b/src/llamafactory/data/processor/pairwise.py @@ -0,0 +1,118 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 collections import defaultdict +from typing import TYPE_CHECKING, Any, Optional + +from ...extras import logging +from ...extras.constants import IGNORE_INDEX +from .processor_utils import DatasetProcessor, infer_seqlen + + +if TYPE_CHECKING: + from ..mm_plugin import AudioInput, ImageInput, VideoInput + + +logger = logging.get_logger(__name__) + + +class PairwiseDatasetProcessor(DatasetProcessor): + def _encode_data_example( + self, + prompt: list[dict[str, str]], + response: list[dict[str, str]], + system: Optional[str], + tools: Optional[str], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + ) -> tuple[list[int], list[int], list[int], list[int]]: + chosen_messages = self.template.mm_plugin.process_messages( + prompt + [response[0]], images, videos, audios, self.processor + ) + rejected_messages = self.template.mm_plugin.process_messages( + prompt + [response[1]], images, videos, audios, self.processor + ) + prompt_ids, chosen_ids = self.template.encode_oneturn(self.tokenizer, chosen_messages, system, tools) + _, rejected_ids = self.template.encode_oneturn(self.tokenizer, rejected_messages, system, tools) + + if self.template.efficient_eos: + chosen_ids += [self.tokenizer.eos_token_id] + rejected_ids += [self.tokenizer.eos_token_id] + + prompt_ids, _ = self.template.mm_plugin.process_token_ids( + prompt_ids, None, images, videos, audios, self.tokenizer, self.processor + ) + # consider the response is more important + source_len, target_len = infer_seqlen( + len(prompt_ids), max(len(chosen_ids), len(rejected_ids)), self.data_args.cutoff_len + ) + prompt_ids = prompt_ids[:source_len] + chosen_ids = chosen_ids[:target_len] + rejected_ids = rejected_ids[:target_len] + + chosen_input_ids = prompt_ids + chosen_ids + chosen_labels = [IGNORE_INDEX] * source_len + chosen_ids + rejected_input_ids = prompt_ids + rejected_ids + rejected_labels = [IGNORE_INDEX] * source_len + rejected_ids + return chosen_input_ids, chosen_labels, rejected_input_ids, rejected_labels + + def preprocess_dataset(self, examples: dict[str, list[Any]]) -> dict[str, list[Any]]: + # build input pairs with format ` X`, `Y1 ` and `Y2 ` + model_inputs = defaultdict(list) + for i in range(len(examples["_prompt"])): + if len(examples["_prompt"][i]) % 2 != 1 or len(examples["_response"][i]) < 2: + logger.warning_rank0( + "Dropped invalid example: {}".format(examples["_prompt"][i] + examples["_response"][i]) + ) + continue + + chosen_input_ids, chosen_labels, rejected_input_ids, rejected_labels = self._encode_data_example( + prompt=examples["_prompt"][i], + response=examples["_response"][i], + system=examples["_system"][i], + tools=examples["_tools"][i], + images=examples["_images"][i] or [], + videos=examples["_videos"][i] or [], + audios=examples["_audios"][i] or [], + ) + model_inputs["chosen_input_ids"].append(chosen_input_ids) + model_inputs["chosen_attention_mask"].append([1] * len(chosen_input_ids)) + model_inputs["chosen_labels"].append(chosen_labels) + model_inputs["rejected_input_ids"].append(rejected_input_ids) + model_inputs["rejected_attention_mask"].append([1] * len(rejected_input_ids)) + model_inputs["rejected_labels"].append(rejected_labels) + model_inputs["images"].append(examples["_images"][i]) + model_inputs["videos"].append(examples["_videos"][i]) + model_inputs["audios"].append(examples["_audios"][i]) + + return model_inputs + + def print_data_example(self, example: dict[str, list[int]]) -> None: + valid_chosen_labels = list(filter(lambda x: x != IGNORE_INDEX, example["chosen_labels"])) + valid_rejected_labels = list(filter(lambda x: x != IGNORE_INDEX, example["rejected_labels"])) + print("chosen_input_ids:\n{}".format(example["chosen_input_ids"])) + print( + "chosen_inputs:\n{}".format(self.tokenizer.decode(example["chosen_input_ids"], skip_special_tokens=False)) + ) + print("chosen_label_ids:\n{}".format(example["chosen_labels"])) + print(f"chosen_labels:\n{self.tokenizer.decode(valid_chosen_labels, skip_special_tokens=False)}") + print("rejected_input_ids:\n{}".format(example["rejected_input_ids"])) + print( + "rejected_inputs:\n{}".format( + self.tokenizer.decode(example["rejected_input_ids"], skip_special_tokens=False) + ) + ) + print("rejected_label_ids:\n{}".format(example["rejected_labels"])) + print(f"rejected_labels:\n{self.tokenizer.decode(valid_rejected_labels, skip_special_tokens=False)}") diff --git a/src/llamafactory/data/processor/pretrain.py b/src/llamafactory/data/processor/pretrain.py new file mode 100644 index 0000000000000000000000000000000000000000..3fa6b1ca58a8d59493cd4b43c51cb268080cc506 --- /dev/null +++ b/src/llamafactory/data/processor/pretrain.py @@ -0,0 +1,57 @@ +# Copyright 2025 HuggingFace Inc. and the LlamaFactory team. +# +# This code is inspired by the HuggingFace's transformers library. +# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/language-modeling/run_clm.py +# +# 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 +from itertools import chain +from typing import Any + +from .processor_utils import DatasetProcessor + + +@dataclass +class PretrainDatasetProcessor(DatasetProcessor): + def preprocess_dataset(self, examples: dict[str, list[Any]]) -> dict[str, list[Any]]: + # build grouped texts with format `X1 X2 X3 ...` if packing is enabled + eos_token = "<|end_of_text|>" if self.data_args.template == "llama3" else self.tokenizer.eos_token + text_examples = [messages[0]["content"] + eos_token for messages in examples["_prompt"]] + + if not self.data_args.packing: + if getattr(self.tokenizer, "add_bos_token", False): + text_examples = [self.tokenizer.bos_token + example for example in text_examples] + + result = self.tokenizer( + text_examples, add_special_tokens=False, truncation=True, max_length=self.data_args.cutoff_len + ) + else: + tokenized_examples = self.tokenizer(text_examples, add_special_tokens=False) + concatenated_examples = {k: list(chain(*tokenized_examples[k])) for k in tokenized_examples.keys()} + total_length = len(concatenated_examples[list(concatenated_examples.keys())[0]]) + block_size = self.data_args.cutoff_len + total_length = (total_length // block_size) * block_size + result = { + k: [t[i : i + block_size] for i in range(0, total_length, block_size)] + for k, t in concatenated_examples.items() + } + if getattr(self.tokenizer, "add_bos_token", False): + for i in range(len(result["input_ids"])): + result["input_ids"][i][0] = self.tokenizer.bos_token_id + + return result + + def print_data_example(self, example: dict[str, list[int]]) -> None: + print("input_ids:\n{}".format(example["input_ids"])) + print("inputs:\n{}".format(self.tokenizer.decode(example["input_ids"], skip_special_tokens=False))) diff --git a/src/llamafactory/data/processor/processor_utils.py b/src/llamafactory/data/processor/processor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..db44b19cf6fc84d6551fb7cce82283774ae72030 --- /dev/null +++ b/src/llamafactory/data/processor/processor_utils.py @@ -0,0 +1,88 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 bisect +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Optional + + +if TYPE_CHECKING: + from transformers import PreTrainedTokenizer, ProcessorMixin + + from ...hparams import DataArguments + from ..template import Template + + +@dataclass +class DatasetProcessor(ABC): + r"""A class for data processors.""" + + template: "Template" + tokenizer: "PreTrainedTokenizer" + processor: Optional["ProcessorMixin"] + data_args: "DataArguments" + + @abstractmethod + def preprocess_dataset(self, examples: dict[str, list[Any]]) -> dict[str, list[Any]]: + r"""Build model inputs from the examples.""" + ... + + @abstractmethod + def print_data_example(self, example: dict[str, list[int]]) -> None: + r"""Print a data example to stdout.""" + ... + + +def search_for_fit(numbers: list[int], capacity: int) -> int: + r"""Find the index of largest number that fits into the knapsack with the given capacity.""" + index = bisect.bisect(numbers, capacity) + return -1 if index == 0 else (index - 1) + + +def greedy_knapsack(numbers: list[int], capacity: int) -> list[list[int]]: + r"""Implement efficient greedy algorithm with binary search for the knapsack problem.""" + numbers.sort() # sort numbers in ascending order for binary search + knapsacks = [] + + while numbers: + current_knapsack = [] + remaining_capacity = capacity + + while True: + index = search_for_fit(numbers, remaining_capacity) + if index == -1: + break # no more numbers fit in this knapsack + + remaining_capacity -= numbers[index] # update the remaining capacity + current_knapsack.append(numbers.pop(index)) # add the number to knapsack + + knapsacks.append(current_knapsack) + + return knapsacks + + +def infer_seqlen(source_len: int, target_len: int, cutoff_len: int) -> tuple[int, int]: + r"""Compute the real sequence length after truncation by the cutoff_len.""" + if target_len * 2 < cutoff_len: # truncate source + max_target_len = cutoff_len + elif source_len * 2 < cutoff_len: # truncate target + max_target_len = cutoff_len - source_len + else: # truncate both + max_target_len = int(cutoff_len * (target_len / (source_len + target_len))) + + new_target_len = min(max_target_len, target_len) + max_source_len = max(cutoff_len - new_target_len, 0) + new_source_len = min(max_source_len, source_len) + return new_source_len, new_target_len diff --git a/src/llamafactory/data/processor/supervised.py b/src/llamafactory/data/processor/supervised.py new file mode 100644 index 0000000000000000000000000000000000000000..26f14c69a7a708fe8a0be712c8b4e27d939bde1f --- /dev/null +++ b/src/llamafactory/data/processor/supervised.py @@ -0,0 +1,252 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 collections import defaultdict +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Any, Optional + +from ...extras import logging +from ...extras.constants import IGNORE_INDEX +from .processor_utils import DatasetProcessor, greedy_knapsack, infer_seqlen + + +if TYPE_CHECKING: + from ..mm_plugin import AudioInput, ImageInput, VideoInput + + +logger = logging.get_logger(__name__) + +MAX_SU_SEQ_IDX = 2**32 # maximum sub-sequence index + + +@dataclass +class PackingParams: + r"""Metadata for a packed sequence: sub-sequence boundaries and multimodal data indices. + + - sequence_boundaries: cumulative token positions, e.g. [0, 100, 250, 512] means 3 sub-seqs + with token ranges [0,100), [100,250), [250,512). Length = num_sub_seqs + 1. + - image_subseq_ids / video_subseq_ids / audio_subseq_ids: for each mm item, the 0-based + sub-sequence index it belongs to. Length = total number of that mm type in the packed sample. + """ + + sequence_boundaries: list[int] + image_subseq_ids: list[int] + video_subseq_ids: list[int] + audio_subseq_ids: list[int] + right_padding_length: int + + +@dataclass +class SupervisedDatasetProcessor(DatasetProcessor): + def _encode_data_example( + self, + prompt: list[dict[str, str]], + response: list[dict[str, str]], + system: Optional[str], + tools: Optional[str], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + ) -> tuple[list[int], list[int]]: + messages = self.template.mm_plugin.process_messages(prompt + response, images, videos, audios, self.processor) + input_ids, labels = self.template.mm_plugin.process_token_ids( + [], [], images, videos, audios, self.tokenizer, self.processor + ) + discarding_history_cot = self.data_args.mask_history and not self.template.preserve_thinking + encoded_pairs = self.template.encode_multiturn(self.tokenizer, messages, system, tools, discarding_history_cot) + total_length = len(input_ids) + (1 if self.template.efficient_eos else 0) + if self.data_args.mask_history: + encoded_pairs = encoded_pairs[::-1] # high priority for last turns + + for turn_idx, (source_ids, target_ids) in enumerate(encoded_pairs): + if total_length >= self.data_args.cutoff_len: + break + + source_len, target_len = infer_seqlen( + len(source_ids), len(target_ids), self.data_args.cutoff_len - total_length + ) + source_ids = source_ids[:source_len] + target_ids = target_ids[:target_len] + total_length += source_len + target_len + + if self.data_args.train_on_prompt: + source_label = source_ids + elif self.template.efficient_eos and turn_idx != 0: + source_label = [self.tokenizer.eos_token_id] + [IGNORE_INDEX] * (source_len - 1) + else: + source_label = [IGNORE_INDEX] * source_len + + if self.data_args.mask_history and turn_idx != 0: # train on the last turn only + target_label = [IGNORE_INDEX] * target_len + else: + target_label = target_ids + + if self.data_args.mask_history: # reversed sequences + input_ids = source_ids + target_ids + input_ids + labels = source_label + target_label + labels + else: + input_ids += source_ids + target_ids + labels += source_label + target_label + + if self.template.efficient_eos: + input_ids += [self.tokenizer.eos_token_id] + labels += [self.tokenizer.eos_token_id] + + return input_ids, labels + + def preprocess_dataset(self, examples: dict[str, list[Any]]) -> dict[str, list[Any]]: + # build inputs with format ` X Y ` and labels with format ` ... Y ` + # for multiturn examples, we only mask the prompt part in each prompt-response pair. + model_inputs = defaultdict(list) + for i in range(len(examples["_prompt"])): + if len(examples["_prompt"][i]) % 2 != 1 or len(examples["_response"][i]) != 1: + logger.warning_rank0( + "Dropped invalid example: {}".format(examples["_prompt"][i] + examples["_response"][i]) + ) + continue + + input_ids, labels = self._encode_data_example( + prompt=examples["_prompt"][i], + response=examples["_response"][i], + system=examples["_system"][i], + tools=examples["_tools"][i], + images=examples["_images"][i] or [], + videos=examples["_videos"][i] or [], + audios=examples["_audios"][i] or [], + ) + model_inputs["input_ids"].append(input_ids) + model_inputs["attention_mask"].append([1] * len(input_ids)) + model_inputs["labels"].append(labels) + model_inputs["images"].append(examples["_images"][i]) + model_inputs["videos"].append(examples["_videos"][i]) + model_inputs["audios"].append(examples["_audios"][i]) + + return model_inputs + + def print_data_example(self, example: dict[str, list[int]]) -> None: + valid_labels = list(filter(lambda x: x != IGNORE_INDEX, example["labels"])) + print("input_ids:\n{}".format(example["input_ids"])) + print("inputs:\n{}".format(self.tokenizer.decode(example["input_ids"], skip_special_tokens=False))) + print("label_ids:\n{}".format(example["labels"])) + print(f"labels:\n{self.tokenizer.decode(valid_labels, skip_special_tokens=False)}") + + +@dataclass +class PackedSupervisedDatasetProcessor(SupervisedDatasetProcessor): + def preprocess_dataset(self, examples: dict[str, list[Any]]) -> dict[str, list[Any]]: + # TODO: use `position_ids` to achieve packing + # build inputs with format ` X1 Y1 X2 Y2 ` + # and labels with format ` ... Y1 ... Y2 ` + valid_num = 0 + batch_input_ids, batch_labels, batch_images, batch_videos, batch_audios = [], [], [], [], [] + lengths = [] + length2indexes = defaultdict(list) + for i in range(len(examples["_prompt"])): + if len(examples["_prompt"][i]) % 2 != 1 or len(examples["_response"][i]) != 1: + logger.warning_rank0( + "Dropped invalid example: {}".format(examples["_prompt"][i] + examples["_response"][i]) + ) + continue + + input_ids, labels = self._encode_data_example( + prompt=examples["_prompt"][i], + response=examples["_response"][i], + system=examples["_system"][i], + tools=examples["_tools"][i], + images=examples["_images"][i] or [], + videos=examples["_videos"][i] or [], + audios=examples["_audios"][i] or [], + ) + length = len(input_ids) + if length > self.data_args.cutoff_len: + logger.warning_rank0(f"Dropped lengthy example with length {length} > {self.data_args.cutoff_len}.") + else: + lengths.append(length) + length2indexes[length].append(valid_num) + batch_input_ids.append(input_ids) + batch_labels.append(labels) + batch_images.append(examples["_images"][i] or []) + batch_videos.append(examples["_videos"][i] or []) + batch_audios.append(examples["_audios"][i] or []) + valid_num += 1 + + model_inputs = defaultdict(list) + requires_packing_params = self.data_args.neat_packing + knapsacks = greedy_knapsack(lengths, self.data_args.cutoff_len) + for knapsack in knapsacks: + packed_input_ids, packed_attention_masks, packed_position_ids, packed_labels = [], [], [], [] + packed_images, packed_videos, packed_audios = [], [], [] + if requires_packing_params: + sequence_boundaries = [0] + image_subseq_ids: list[int] = [] + video_subseq_ids: list[int] = [] + audio_subseq_ids: list[int] = [] + + for i, length in enumerate(knapsack): + index = length2indexes[length].pop() + packed_input_ids += batch_input_ids[index] + packed_position_ids += list(range(len(batch_input_ids[index]))) # NOTE: pad_to_multiple_of ignore this + packed_labels += batch_labels[index] + packed_images += batch_images[index] + packed_videos += batch_videos[index] + packed_audios += batch_audios[index] + if requires_packing_params: + n_img = len(batch_images[index]) + n_vid = len(batch_videos[index]) + n_aud = len(batch_audios[index]) + sequence_boundaries.append(sequence_boundaries[-1] + len(batch_input_ids[index])) + image_subseq_ids.extend([i] * n_img) + video_subseq_ids.extend([i] * n_vid) + audio_subseq_ids.extend([i] * n_aud) + + if self.data_args.neat_packing: + packed_attention_masks += [i + 1] * len(batch_input_ids[index]) # start from 1 + else: + packed_attention_masks += [1] * len(batch_input_ids[index]) + + if len(packed_input_ids) < self.data_args.cutoff_len + 1: # avoid flash_attn drops attn mask + pad_length = self.data_args.cutoff_len - len(packed_input_ids) + 1 + packed_input_ids += [self.tokenizer.pad_token_id] * pad_length + packed_position_ids += [0] * pad_length + packed_labels += [IGNORE_INDEX] * pad_length + if self.data_args.neat_packing: + packed_attention_masks += [0] * pad_length + else: + packed_attention_masks += [1] * pad_length # more efficient flash_attn + + if requires_packing_params: + sequence_boundaries.append(sequence_boundaries[-1] + pad_length) + + if len(packed_input_ids) != self.data_args.cutoff_len + 1: + raise ValueError("The length of packed example should be identical to the cutoff length.") + + model_inputs["input_ids"].append(packed_input_ids) + if requires_packing_params: + packing_params = PackingParams( + sequence_boundaries=sequence_boundaries, + image_subseq_ids=image_subseq_ids or [MAX_SU_SEQ_IDX], # avoid dataset concat error + video_subseq_ids=video_subseq_ids or [MAX_SU_SEQ_IDX], + audio_subseq_ids=audio_subseq_ids or [MAX_SU_SEQ_IDX], + right_padding_length=pad_length, + ) + model_inputs["packing_params"].append(asdict(packing_params)) + + model_inputs["attention_mask"].append(packed_attention_masks) + model_inputs["position_ids"].append(packed_position_ids) + model_inputs["labels"].append(packed_labels) + model_inputs["images"].append(packed_images or None) + model_inputs["videos"].append(packed_videos or None) + model_inputs["audios"].append(packed_audios or None) + + return model_inputs diff --git a/src/llamafactory/data/processor/unsupervised.py b/src/llamafactory/data/processor/unsupervised.py new file mode 100644 index 0000000000000000000000000000000000000000..256174b6dd38696b5b180501102af40ff395d0a9 --- /dev/null +++ b/src/llamafactory/data/processor/unsupervised.py @@ -0,0 +1,91 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 collections import defaultdict +from typing import TYPE_CHECKING, Any, Optional + +from ...extras import logging +from ..data_utils import Role +from .processor_utils import DatasetProcessor, infer_seqlen + + +if TYPE_CHECKING: + from ..mm_plugin import AudioInput, ImageInput, VideoInput + + +logger = logging.get_logger(__name__) + + +class UnsupervisedDatasetProcessor(DatasetProcessor): + def _encode_data_example( + self, + prompt: list[dict[str, str]], + response: list[dict[str, str]], + system: Optional[str], + tools: Optional[str], + images: list["ImageInput"], + videos: list["VideoInput"], + audios: list["AudioInput"], + ) -> tuple[list[int], list[int]]: + if len(response) == 1: + messages = prompt + response + else: + messages = prompt + [{"role": Role.ASSISTANT.value, "content": ""}] + + messages = self.template.mm_plugin.process_messages(messages, images, videos, audios, self.processor) + input_ids, labels = self.template.encode_oneturn(self.tokenizer, messages, system, tools) + if self.template.efficient_eos: + labels += [self.tokenizer.eos_token_id] + + input_ids, _ = self.template.mm_plugin.process_token_ids( + input_ids, None, images, videos, audios, self.tokenizer, self.processor + ) + source_len, target_len = infer_seqlen(len(input_ids), len(labels), self.data_args.cutoff_len) + input_ids = input_ids[:source_len] + labels = labels[:target_len] + return input_ids, labels + + def preprocess_dataset(self, examples: dict[str, list[Any]]) -> dict[str, list[Any]]: + # build inputs with format ` X` and labels with format `Y ` + model_inputs = defaultdict(list) + for i in range(len(examples["_prompt"])): + if len(examples["_prompt"][i]) % 2 != 1: + logger.warning_rank0( + "Dropped invalid example: {}".format(examples["_prompt"][i] + examples["_response"][i]) + ) + continue + + input_ids, labels = self._encode_data_example( + prompt=examples["_prompt"][i], + response=examples["_response"][i], + system=examples["_system"][i], + tools=examples["_tools"][i], + images=examples["_images"][i] or [], + videos=examples["_videos"][i] or [], + audios=examples["_audios"][i] or [], + ) + model_inputs["input_ids"].append(input_ids) + model_inputs["attention_mask"].append([1] * len(input_ids)) + model_inputs["labels"].append(labels) + model_inputs["images"].append(examples["_images"][i]) + model_inputs["videos"].append(examples["_videos"][i]) + model_inputs["audios"].append(examples["_audios"][i]) + + return model_inputs + + def print_data_example(self, example: dict[str, list[int]]) -> None: + print("input_ids:\n{}".format(example["input_ids"])) + print("inputs:\n{}".format(self.tokenizer.decode(example["input_ids"], skip_special_tokens=False))) + print("label_ids:\n{}".format(example["labels"])) + print("labels:\n{}".format(self.tokenizer.decode(example["labels"], skip_special_tokens=False))) diff --git a/src/llamafactory/data/template.py b/src/llamafactory/data/template.py new file mode 100644 index 0000000000000000000000000000000000000000..492909eea8f947d994cb4a5832bf6593a89177aa --- /dev/null +++ b/src/llamafactory/data/template.py @@ -0,0 +1,2410 @@ +# Copyright 2025 the LlamaFactory team. +# +# 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 re +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Union + +from typing_extensions import override + +from ..extras import logging +from .data_utils import Role +from .formatter import EmptyFormatter, FunctionFormatter, StringFormatter, ToolFormatter +from .mm_plugin import get_mm_plugin + + +if TYPE_CHECKING: + from transformers import PreTrainedTokenizer + + from ..hparams import DataArguments + from .formatter import SLOTS, Formatter + from .mm_plugin import BasePlugin + from .tool_utils import FunctionCall + + +logger = logging.get_logger(__name__) + + +@dataclass +class Template: + format_user: "Formatter" + format_assistant: "Formatter" + format_system: "Formatter" + format_function: "Formatter" + format_observation: "Formatter" + format_tools: "Formatter" + format_prefix: "Formatter" + default_system: str + stop_words: list[str] + thought_words: tuple[str, str] + tool_call_words: tuple[str, str] + efficient_eos: bool + replace_eos: bool + replace_jinja_template: bool + enable_thinking: Optional[bool] + preserve_thinking: bool + mm_plugin: "BasePlugin" + + def encode_oneturn( + self, + tokenizer: "PreTrainedTokenizer", + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + ) -> tuple[list[int], list[int]]: + r"""Return a single pair of token ids representing prompt and response respectively.""" + encoded_messages = self._encode(tokenizer, messages, system, tools) + prompt_ids = [] + for encoded_ids in encoded_messages[:-1]: + prompt_ids += encoded_ids + + response_ids = encoded_messages[-1] + return prompt_ids, response_ids + + def encode_multiturn( + self, + tokenizer: "PreTrainedTokenizer", + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + discarding_history_cot: bool = False, # only effect reasoning template + ) -> list[tuple[list[int], list[int]]]: + r"""Return multiple pairs of token ids representing prompts and responses respectively.""" + encoded_messages = self._encode(tokenizer, messages, system, tools) + return [(encoded_messages[i], encoded_messages[i + 1]) for i in range(0, len(encoded_messages), 2)] + + def extract_tool(self, content: str) -> Union[str, list["FunctionCall"]]: + r"""Extract tool message.""" + return self.format_tools.extract(content) + + def get_stop_token_ids(self, tokenizer: "PreTrainedTokenizer") -> list[int]: + r"""Return stop token ids.""" + stop_token_ids = {tokenizer.eos_token_id} + for token in self.stop_words: + stop_token_ids.add(tokenizer.convert_tokens_to_ids(token)) + + return list(stop_token_ids) + + def add_thought(self, content: str = "") -> str: + r"""Add empty thought to assistant message.""" + return f"{self.thought_words[0]}{self.thought_words[1]}" + content + + def remove_thought(self, content: str) -> str: + r"""Remove thought from assistant message.""" + pattern = re.compile(f"{re.escape(self.thought_words[0])}(.*?){re.escape(self.thought_words[1])}", re.DOTALL) + return re.sub(pattern, "", content).lstrip("\n") + + def get_thought_word_ids(self, tokenizer: "PreTrainedTokenizer") -> list[int]: + r"""Get the token ids of thought words.""" + return tokenizer.encode(self.add_thought(), add_special_tokens=False) + + def _convert_elements_to_ids(self, tokenizer: "PreTrainedTokenizer", elements: "SLOTS") -> list[int]: + r"""Convert elements to token ids.""" + token_ids = [] + for elem in elements: + if isinstance(elem, str): + if len(elem) != 0: + token_ids += tokenizer.encode(elem, add_special_tokens=False) + elif isinstance(elem, dict): + token_ids += [tokenizer.convert_tokens_to_ids(elem.get("token"))] + elif isinstance(elem, set): + if "bos_token" in elem and tokenizer.bos_token_id is not None: + token_ids += [tokenizer.bos_token_id] + elif "eos_token" in elem and tokenizer.eos_token_id is not None: + token_ids += [tokenizer.eos_token_id] + else: + raise ValueError(f"Input must be string, set[str] or dict[str, str], got {type(elem)}") + + return token_ids + + def _encode( + self, + tokenizer: "PreTrainedTokenizer", + messages: list[dict[str, str]], + system: Optional[str], + tools: Optional[str], + ) -> list[list[int]]: + r"""Encode formatted inputs to pairs of token ids. + + Turn 0: prefix + system + query resp + Turn t: query resp. + """ + system = system or self.default_system + encoded_messages = [] + for i, message in enumerate(messages): + elements = [] + + if i == 0: + elements += self.format_prefix.apply() + if system or tools: + tool_text = self.format_tools.apply(content=tools)[0] if tools else "" + elements += self.format_system.apply(content=(system + tool_text)) + + if message["role"] == Role.USER: + elements += self.format_user.apply(content=message["content"], idx=str(i // 2)) + elif message["role"] == Role.ASSISTANT: + elements += self.format_assistant.apply(content=message["content"]) + elif message["role"] == Role.OBSERVATION: + elements += self.format_observation.apply(content=message["content"]) + elif message["role"] == Role.FUNCTION: + elements += self.format_function.apply( + content=message["content"], thought_words=self.thought_words, tool_call_words=self.tool_call_words + ) + else: + raise NotImplementedError("Unexpected role: {}".format(message["role"])) + + encoded_messages.append(self._convert_elements_to_ids(tokenizer, elements)) + + return encoded_messages + + @staticmethod + def _add_or_replace_eos_token(tokenizer: "PreTrainedTokenizer", eos_token: str) -> None: + r"""Add or replace eos token to the tokenizer.""" + if tokenizer.eos_token == eos_token: + return + + is_added = tokenizer.eos_token_id is None + num_added_tokens = tokenizer.add_special_tokens({"eos_token": eos_token}) + + if is_added: + logger.info_rank0(f"Add eos token: {tokenizer.eos_token}.") + else: + logger.info_rank0(f"Replace eos token: {tokenizer.eos_token}.") + + if num_added_tokens > 0: + logger.warning_rank0("New tokens have been added, make sure `resize_vocab` is True.") + + def fix_special_tokens(self, tokenizer: "PreTrainedTokenizer") -> None: + r"""Add eos token and pad token to the tokenizer.""" + stop_words = self.stop_words + if self.replace_eos: + if not stop_words: + raise ValueError("Stop words are required to replace the EOS token.") + + self._add_or_replace_eos_token(tokenizer, eos_token=stop_words[0]) + stop_words = stop_words[1:] + + if tokenizer.eos_token_id is None: + self._add_or_replace_eos_token(tokenizer, eos_token="<|endoftext|>") + + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + logger.info_rank0(f"Add pad token: {tokenizer.pad_token}") + + if stop_words: + try: + num_added_tokens = tokenizer.add_special_tokens( + dict(additional_special_tokens=stop_words), replace_additional_special_tokens=False + ) + except TypeError: + num_added_tokens = tokenizer.add_special_tokens(dict(additional_special_tokens=stop_words)) + logger.info_rank0("Add {} to stop words.".format(",".join(stop_words))) + if num_added_tokens > 0: + logger.warning_rank0("New tokens have been added, make sure `resize_vocab` is True.") + + @staticmethod + def _jinja_escape(content: str) -> str: + r"""Escape single quotes in content.""" + return content.replace("'", r"\'") + + @staticmethod + def _convert_slots_to_jinja(slots: "SLOTS", tokenizer: "PreTrainedTokenizer", placeholder: str = "content") -> str: + r"""Convert slots to jinja template.""" + slot_items = [] + for slot in slots: + if isinstance(slot, str): + slot_pieces = slot.split("{{content}}") + if slot_pieces[0]: + slot_items.append("'" + Template._jinja_escape(slot_pieces[0]) + "'") + if len(slot_pieces) > 1: + slot_items.append(placeholder) + if slot_pieces[1]: + slot_items.append("'" + Template._jinja_escape(slot_pieces[1]) + "'") + elif isinstance(slot, set): # do not use {{ eos_token }} since it may be replaced + if "bos_token" in slot and tokenizer.bos_token_id is not None: + slot_items.append("'" + tokenizer.bos_token + "'") + elif "eos_token" in slot and tokenizer.eos_token_id is not None: + slot_items.append("'" + tokenizer.eos_token + "'") + elif isinstance(slot, dict): + raise ValueError("Dict is not supported.") + + return " + ".join(slot_items) + + def _get_jinja_template(self, tokenizer: "PreTrainedTokenizer") -> str: + r"""Return the jinja template.""" + prefix = self._convert_slots_to_jinja(self.format_prefix.apply(), tokenizer) + system = self._convert_slots_to_jinja(self.format_system.apply(), tokenizer, placeholder="system_message") + user = self._convert_slots_to_jinja(self.format_user.apply(), tokenizer) + assistant = self._convert_slots_to_jinja(self.format_assistant.apply(), tokenizer) + jinja_template = "" + if prefix: + jinja_template += "{{ " + prefix + " }}" + + if self.default_system: + jinja_template += "{% set system_message = '" + self._jinja_escape(self.default_system) + "' %}" + + jinja_template += ( + "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}" + "{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% endif %}" + "{% if system_message is defined %}{{ " + system + " }}{% endif %}" + "{% for message in loop_messages %}" + "{% set content = message['content'] %}" + "{% if message['role'] == 'user' %}" + "{{ " + user + " }}" + "{% elif message['role'] == 'assistant' %}" + "{{ " + assistant + " }}" + "{% endif %}" + "{% endfor %}" + ) + return jinja_template + + def fix_jinja_template(self, tokenizer: "PreTrainedTokenizer") -> None: + r"""Replace the jinja template in the tokenizer.""" + if tokenizer.chat_template is None or self.replace_jinja_template: + try: + tokenizer.chat_template = self._get_jinja_template(tokenizer) + except ValueError as e: + logger.info_rank0(f"Cannot add this chat template to tokenizer: {e}.") + + @staticmethod + def _convert_slots_to_ollama( + slots: "SLOTS", tokenizer: "PreTrainedTokenizer", placeholder: str = "content" + ) -> str: + r"""Convert slots to ollama template.""" + slot_items = [] + for slot in slots: + if isinstance(slot, str): + slot_pieces = slot.split("{{content}}") + if slot_pieces[0]: + slot_items.append(slot_pieces[0]) + if len(slot_pieces) > 1: + slot_items.append("{{ " + placeholder + " }}") + if slot_pieces[1]: + slot_items.append(slot_pieces[1]) + elif isinstance(slot, set): # do not use {{ eos_token }} since it may be replaced + if "bos_token" in slot and tokenizer.bos_token_id is not None: + slot_items.append(tokenizer.bos_token) + elif "eos_token" in slot and tokenizer.eos_token_id is not None: + slot_items.append(tokenizer.eos_token) + elif isinstance(slot, dict): + raise ValueError("Dict is not supported.") + + return "".join(slot_items) + + def _get_ollama_template(self, tokenizer: "PreTrainedTokenizer") -> str: + r"""Return the ollama template.""" + prefix = self._convert_slots_to_ollama(self.format_prefix.apply(), tokenizer) + system = self._convert_slots_to_ollama(self.format_system.apply(), tokenizer, placeholder=".System") + user = self._convert_slots_to_ollama(self.format_user.apply(), tokenizer, placeholder=".Content") + assistant = self._convert_slots_to_ollama(self.format_assistant.apply(), tokenizer, placeholder=".Content") + return ( + f"{prefix}{{{{ if .System }}}}{system}{{{{ end }}}}" + f"""{{{{ range .Messages }}}}{{{{ if eq .Role "user" }}}}{user}""" + f"""{{{{ else if eq .Role "assistant" }}}}{assistant}{{{{ end }}}}{{{{ end }}}}""" + ) + + def get_ollama_modelfile(self, tokenizer: "PreTrainedTokenizer") -> str: + r"""Return the ollama modelfile. + + TODO: support function calling. + """ + modelfile = "# ollama modelfile auto-generated by llamafactory\n\n" + modelfile += f'FROM .\n\nTEMPLATE """{self._get_ollama_template(tokenizer)}"""\n\n' + + if self.default_system: + modelfile += f'SYSTEM """{self.default_system}"""\n\n' + + for stop_token_id in self.get_stop_token_ids(tokenizer): + modelfile += f'PARAMETER stop "{tokenizer.convert_ids_to_tokens(stop_token_id)}"\n' + + modelfile += "PARAMETER num_ctx 4096\n" + return modelfile + + +@dataclass +class Llama2Template(Template): + r"""A template that fuse the system message to first user message.""" + + @override + def _encode( + self, + tokenizer: "PreTrainedTokenizer", + messages: list[dict[str, str]], + system: str, + tools: str, + ) -> list[list[int]]: + system = system or self.default_system + encoded_messages = [] + for i, message in enumerate(messages): + elements = [] + + system_text = "" + if i == 0: + elements += self.format_prefix.apply() + if system or tools: + tool_text = self.format_tools.apply(content=tools)[0] if tools else "" + system_text = self.format_system.apply(content=(system + tool_text))[0] + + if message["role"] == Role.USER: + elements += self.format_user.apply(content=system_text + message["content"]) + elif message["role"] == Role.ASSISTANT: + elements += self.format_assistant.apply(content=message["content"]) + elif message["role"] == Role.OBSERVATION: + elements += self.format_observation.apply(content=message["content"]) + elif message["role"] == Role.FUNCTION: + elements += self.format_function.apply(content=message["content"]) + else: + raise NotImplementedError("Unexpected role: {}".format(message["role"])) + + encoded_messages.append(self._convert_elements_to_ids(tokenizer, elements)) + + return encoded_messages + + def _get_jinja_template(self, tokenizer: "PreTrainedTokenizer") -> str: + prefix = self._convert_slots_to_jinja(self.format_prefix.apply(), tokenizer) + system_message = self._convert_slots_to_jinja( + self.format_system.apply(), tokenizer, placeholder="system_message" + ) + user_message = self._convert_slots_to_jinja(self.format_user.apply(), tokenizer) + assistant_message = self._convert_slots_to_jinja(self.format_assistant.apply(), tokenizer) + jinja_template = "" + if prefix: + jinja_template += "{{ " + prefix + " }}" + + if self.default_system: + jinja_template += "{% set system_message = '" + self._jinja_escape(self.default_system) + "' %}" + + jinja_template += ( + "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}" + "{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% endif %}" + "{% for message in loop_messages %}" + "{% if loop.index0 == 0 and system_message is defined %}" + "{% set content = " + system_message + " + message['content'] %}" + "{% else %}{% set content = message['content'] %}{% endif %}" + "{% if message['role'] == 'user' %}" + "{{ " + user_message + " }}" + "{% elif message['role'] == 'assistant' %}" + "{{ " + assistant_message + " }}" + "{% endif %}" + "{% endfor %}" + ) + return jinja_template + + +@dataclass +class ReasoningTemplate(Template): + r"""A template that add thought to assistant message.""" + + @override + def encode_oneturn( + self, + tokenizer: "PreTrainedTokenizer", + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + ) -> tuple[list[int], list[int]]: + messages = deepcopy(messages) + if not self.preserve_thinking: + for i in range(1, len(messages) - 2, 2): + messages[i]["content"] = self.remove_thought(messages[i]["content"]) + + if self.enable_thinking is False: # remove all cot + messages[-1]["content"] = self.remove_thought(messages[-1]["content"]) + + prompt_ids, response_ids = super().encode_oneturn(tokenizer, messages, system, tools) + if ( + self.thought_words[0].strip() not in messages[-1]["content"] + and self.thought_words[1].strip() not in messages[-1]["content"] + ): # add empty cot + if not self.enable_thinking: # do not compute loss + prompt_ids += self.get_thought_word_ids(tokenizer) + else: # do compute loss + response_ids = self.get_thought_word_ids(tokenizer) + response_ids + + return prompt_ids, response_ids + + @override + def encode_multiturn( + self, + tokenizer: "PreTrainedTokenizer", + messages: list[dict[str, str]], + system: Optional[str] = None, + tools: Optional[str] = None, + discarding_history_cot: bool = False, + ) -> list[tuple[list[int], list[int]]]: + messages = deepcopy(messages) + if self.enable_thinking is False: # remove all cot + for i in range(1, len(messages), 2): + messages[i]["content"] = self.remove_thought(messages[i]["content"]) + + if discarding_history_cot: + for i in range(1, len(messages) - 2, 2): # preserve the last cot + messages[i]["content"] = self.remove_thought(messages[i]["content"]) + + encoded_messages = self._encode(tokenizer, messages, system, tools) + if discarding_history_cot: + turn_indices = [len(messages) - 2] + else: + turn_indices = range(0, len(messages), 2) + + for i in turn_indices: + if ( + self.thought_words[0].strip() not in messages[i + 1]["content"] + and self.thought_words[1].strip() not in messages[i + 1]["content"] + ): # add empty cot + if not self.enable_thinking: # do not compute loss + encoded_messages[i] += self.get_thought_word_ids(tokenizer) + else: # do compute loss + encoded_messages[i + 1] = self.get_thought_word_ids(tokenizer) + encoded_messages[i + 1] + + return [(encoded_messages[i], encoded_messages[i + 1]) for i in range(0, len(encoded_messages), 2)] + + +@dataclass +class Glm47ReasoningTemplate(ReasoningTemplate): + r"""GLM-4.7 uses only the closing tag for empty thinking blocks.""" + + @override + def add_thought(self, content: str = "") -> str: + if not content: + return self.thought_words[1] + + return self.thought_words[0] + content + self.thought_words[1] + + +TEMPLATES: dict[str, "Template"] = {} + + +def register_template( + name: str, + format_user: Optional["Formatter"] = None, + format_assistant: Optional["Formatter"] = None, + format_system: Optional["Formatter"] = None, + format_function: Optional["Formatter"] = None, + format_observation: Optional["Formatter"] = None, + format_tools: Optional["Formatter"] = None, + format_prefix: Optional["Formatter"] = None, + default_system: str = "", + stop_words: Optional[list[str]] = None, + thought_words: Optional[tuple[str, str]] = None, + tool_call_words: Optional[tuple[str, str]] = None, + efficient_eos: bool = False, + replace_eos: bool = False, + replace_jinja_template: bool = False, + enable_thinking: Optional[bool] = True, + preserve_thinking: bool = False, + mm_plugin: "BasePlugin" = get_mm_plugin(name="base"), + template_class: type["Template"] = Template, +) -> None: + r"""Register a chat template. + + To add the following chat template: + ``` + user prompt here + model response here + user prompt here + model response here + ``` + + The corresponding code should be: + ``` + register_template( + name="custom", + format_user=StringFormatter(slots=["{{content}}\n"]), + format_assistant=StringFormatter(slots=["{{content}}\n"]), + format_prefix=EmptyFormatter(""), + ) + ``` + """ + if name in TEMPLATES: + raise ValueError(f"Template {name} already exists.") + + default_slots = ["{{content}}"] if efficient_eos else ["{{content}}", {"eos_token"}] + default_user_formatter = StringFormatter(slots=["{{content}}"]) + default_assistant_formatter = StringFormatter(slots=default_slots) + if format_assistant is not None: + default_function_formatter = FunctionFormatter(slots=format_assistant.slots, tool_format="default") + else: + default_function_formatter = FunctionFormatter(slots=default_slots, tool_format="default") + + default_tool_formatter = ToolFormatter(tool_format="default") + default_prefix_formatter = EmptyFormatter() + TEMPLATES[name] = template_class( + format_user=format_user or default_user_formatter, + format_assistant=format_assistant or default_assistant_formatter, + format_system=format_system or default_user_formatter, + format_function=format_function or default_function_formatter, + format_observation=format_observation or format_user or default_user_formatter, + format_tools=format_tools or default_tool_formatter, + format_prefix=format_prefix or default_prefix_formatter, + default_system=default_system, + stop_words=stop_words or [], + thought_words=thought_words or ("\n", "\n\n\n"), + tool_call_words=tool_call_words or ("", ""), + efficient_eos=efficient_eos, + replace_eos=replace_eos, + replace_jinja_template=replace_jinja_template, + enable_thinking=enable_thinking, + preserve_thinking=preserve_thinking, + mm_plugin=mm_plugin, + ) + + +def parse_template(tokenizer: "PreTrainedTokenizer") -> "Template": + r"""Extract a chat template from the tokenizer.""" + + def find_diff(short_str: str, long_str: str) -> str: + i, j = 0, 0 + diff = "" + while i < len(short_str) and j < len(long_str): + if short_str[i] == long_str[j]: + i += 1 + j += 1 + else: + diff += long_str[j] + j += 1 + + return diff + + prefix = tokenizer.decode(tokenizer.encode("")) + + messages = [{"role": "system", "content": "{{content}}"}] + system_slot = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False)[len(prefix) :] + + messages = [{"role": "system", "content": ""}, {"role": "user", "content": "{{content}}"}] + user_slot_empty_system = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + user_slot_empty_system = user_slot_empty_system[len(prefix) :] + + messages = [{"role": "user", "content": "{{content}}"}] + user_slot = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + user_slot = user_slot[len(prefix) :] + + messages = [{"role": "user", "content": "{{content}}"}, {"role": "assistant", "content": "{{content}}"}] + assistant_slot = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False) + assistant_slot = assistant_slot[len(prefix) + len(user_slot) :] + template_class = ReasoningTemplate if "" in assistant_slot else Template + assistant_slot = assistant_slot.replace("", "").replace("", "").lstrip("\n") # remove thought tags + + if len(user_slot) > len(user_slot_empty_system): + default_system = find_diff(user_slot_empty_system, user_slot) + sole_system = system_slot.replace("{{content}}", default_system, 1) + user_slot = user_slot[len(sole_system) :] + else: # if defaut_system is empty, user_slot_empty_system will be longer than user_slot + default_system = "" + + return template_class( + format_user=StringFormatter(slots=[user_slot]), + format_assistant=StringFormatter(slots=[assistant_slot]), + format_system=StringFormatter(slots=[system_slot]), + format_function=FunctionFormatter(slots=[assistant_slot], tool_format="default"), + format_observation=StringFormatter(slots=[user_slot]), + format_tools=ToolFormatter(tool_format="default"), + format_prefix=EmptyFormatter(slots=[prefix]) if prefix else EmptyFormatter(), + default_system=default_system, + stop_words=[], + thought_words=("\n", "\n\n\n"), + tool_call_words=("", ""), + efficient_eos=False, + replace_eos=False, + replace_jinja_template=False, + enable_thinking=True, + preserve_thinking=False, + mm_plugin=get_mm_plugin(name="base"), + ) + + +def get_template_and_fix_tokenizer(tokenizer: "PreTrainedTokenizer", data_args: "DataArguments") -> "Template": + r"""Get chat template and fixes the tokenizer.""" + if data_args.template is None: + if isinstance(tokenizer.chat_template, str): + logger.warning_rank0("`template` was not specified, try parsing the chat template from the tokenizer.") + template = parse_template(tokenizer) + else: + logger.warning_rank0("`template` was not specified, use `empty` template.") + template = TEMPLATES["empty"] # placeholder + else: + if data_args.template not in TEMPLATES: + raise ValueError(f"Template {data_args.template} does not exist.") + + template = TEMPLATES[data_args.template] + + if data_args.train_on_prompt and template.efficient_eos: + raise ValueError("Current template does not support `train_on_prompt`.") + + if data_args.tool_format is not None: + logger.info_rank0(f"Using tool format: {data_args.tool_format}.") + default_slots = ["{{content}}"] if template.efficient_eos else ["{{content}}", {"eos_token"}] + template.format_function = FunctionFormatter(slots=default_slots, tool_format=data_args.tool_format) + template.format_tools = ToolFormatter(tool_format=data_args.tool_format) + + if data_args.default_system is not None: + logger.info_rank0(f"Using default system message: {data_args.default_system}.") + template.default_system = data_args.default_system + + if isinstance(template, ReasoningTemplate): + logger.warning_rank0( + "You are using reasoning template. " + "If the base model is NOT a reasoning model (i.e., it has a separate Instruct variant), " + "please add `_nothink` suffix to disable thinking. " + "For reasoning-only model families (e.g., Qwen3.6), the suffix is not needed. " + "e.g., qwen3_vl_nothink" + ) + template.enable_thinking = data_args.enable_thinking + template.preserve_thinking = data_args.preserve_thinking + + template.fix_special_tokens(tokenizer) + template.fix_jinja_template(tokenizer) + return template + + +register_template( + name="alpaca", + format_user=StringFormatter(slots=["### Instruction:\n{{content}}\n\n### Response:\n"]), + format_assistant=StringFormatter(slots=["{{content}}", {"eos_token"}, "\n\n"]), + default_system=( + "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n" + ), + replace_jinja_template=True, +) + + +register_template( + name="bailing", + format_user=StringFormatter(slots=["HUMAN{{content}}ASSISTANT"]), + format_system=StringFormatter(slots=["SYSTEM{{content}}"]), + format_observation=StringFormatter(slots=["OBSERVATION{{content}}ASSISTANT"]), + stop_words=["<|endoftext|>"], + efficient_eos=True, +) + + +register_template( + name="bailing_v2", + format_user=StringFormatter(slots=["HUMAN{{content}}<|role_end|>ASSISTANT"]), + format_system=StringFormatter(slots=["SYSTEM{{content}}<|role_end|>"]), + format_assistant=StringFormatter(slots=["{{content}}<|role_end|>"]), + format_observation=StringFormatter( + slots=[ + "OBSERVATION\n\n{{content}}\n<|role_end|>ASSISTANT" + ] + ), + format_function=FunctionFormatter(slots=["{{content}}<|role_end|>"], tool_format="ling"), + format_tools=ToolFormatter(tool_format="ling"), + stop_words=["<|endoftext|>"], + efficient_eos=True, +) + + +register_template( + name="breeze", + format_user=StringFormatter(slots=["[INST] {{content}} [/INST] "]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + efficient_eos=True, +) + + +register_template( + name="chatglm3", + format_user=StringFormatter(slots=[{"token": "<|user|>"}, "\n", "{{content}}", {"token": "<|assistant|>"}]), + format_assistant=StringFormatter(slots=["\n", "{{content}}"]), + format_system=StringFormatter(slots=[{"token": "<|system|>"}, "\n", "{{content}}"]), + format_function=FunctionFormatter(slots=["{{content}}"], tool_format="glm4"), + format_observation=StringFormatter( + slots=[{"token": "<|observation|>"}, "\n", "{{content}}", {"token": "<|assistant|>"}] + ), + format_tools=ToolFormatter(tool_format="glm4"), + format_prefix=EmptyFormatter(slots=[{"token": "[gMASK]"}, {"token": "sop"}]), + stop_words=["<|user|>", "<|observation|>"], + efficient_eos=True, +) + + +register_template( + name="chatml", + format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + format_assistant=StringFormatter(slots=["{{content}}<|im_end|>\n"]), + format_system=StringFormatter(slots=["<|im_start|>system\n{{content}}<|im_end|>\n"]), + format_observation=StringFormatter(slots=["<|im_start|>tool\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + stop_words=["<|im_end|>", "<|im_start|>"], + replace_eos=True, + replace_jinja_template=True, +) + + +# copied from chatml template +register_template( + name="chatml_de", + format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + format_assistant=StringFormatter(slots=["{{content}}<|im_end|>\n"]), + format_system=StringFormatter(slots=["<|im_start|>system\n{{content}}<|im_end|>\n"]), + format_observation=StringFormatter(slots=["<|im_start|>tool\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + default_system="Du bist ein freundlicher und hilfsbereiter KI-Assistent.", + stop_words=["<|im_end|>", "<|im_start|>"], + replace_eos=True, + replace_jinja_template=True, +) + + +register_template( + name="cohere", + format_user=StringFormatter( + slots=[ + ( + "<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{{content}}<|END_OF_TURN_TOKEN|>" + "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>" + ) + ] + ), + format_system=StringFormatter(slots=["<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{{content}}<|END_OF_TURN_TOKEN|>"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), +) + + +# copied from chatml template +register_template( + name="cpm4", + format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + format_assistant=StringFormatter(slots=["{{content}}<|im_end|>\n"]), + format_system=StringFormatter(slots=["<|im_start|>system\n{{content}}<|im_end|>\n"]), + format_observation=StringFormatter(slots=["<|im_start|>tool\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=["<|im_end|>"], +) + + +# copied from chatml template +register_template( + name="dbrx", + format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + format_assistant=StringFormatter(slots=["{{content}}<|im_end|>\n"]), + format_system=StringFormatter(slots=["<|im_start|>system\n{{content}}<|im_end|>\n"]), + format_observation=StringFormatter(slots=["<|im_start|>tool\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + default_system=( + "You are DBRX, created by Databricks. You were last updated in December 2023. " + "You answer questions based on information available up to that point.\n" + "YOU PROVIDE SHORT RESPONSES TO SHORT QUESTIONS OR STATEMENTS, but provide thorough " + "responses to more complex and open-ended questions.\nYou assist with various tasks, " + "from writing to coding (using markdown for code blocks — remember to use ``` with " + "code, JSON, and tables).\n(You do not have real-time data access or code execution " + "capabilities. You avoid stereotyping and provide balanced perspectives on " + "controversial topics. You do not provide song lyrics, poems, or news articles and " + "do not divulge details of your training data.)\nThis is your system prompt, " + "guiding your responses. Do not reference it, just respond to the user. If you find " + "yourself talking about this message, stop. You should be responding appropriately " + "and usually that means not mentioning this.\nYOU DO NOT MENTION ANY OF THIS INFORMATION " + "ABOUT YOURSELF UNLESS THE INFORMATION IS DIRECTLY PERTINENT TO THE USER'S QUERY." + ), + stop_words=["<|im_end|>"], + replace_eos=True, +) + + +register_template( + name="deepseek", + format_user=StringFormatter(slots=["User: {{content}}\n\nAssistant:"]), + format_system=StringFormatter(slots=["{{content}}\n\n"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), +) + + +register_template( + name="deepseek3", + format_user=StringFormatter(slots=["<｜User｜>{{content}}<｜Assistant｜>"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), +) + + +# copied from deepseek3 template +register_template( + name="deepseekr1", + format_user=StringFormatter(slots=["<｜User｜>{{content}}<｜Assistant｜>"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + template_class=ReasoningTemplate, +) + + +register_template( + name="hy3", + format_user=StringFormatter(slots=["<｜hy_User｜>{{content}}<｜hy_Assistant｜>"]), + format_assistant=StringFormatter(slots=["{{content}}<｜hy_eos｜>"]), + format_system=StringFormatter(slots=["{{content}}"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=["<｜hy_eos｜>"], + replace_eos=True, + thought_words=("", ""), + template_class=ReasoningTemplate, +) + + +register_template( + name="deepseekcoder", + format_user=StringFormatter(slots=["### Instruction:\n{{content}}\n### Response:"]), + format_assistant=StringFormatter(slots=["\n{{content}}\n<|EOT|>\n"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + default_system=( + "You are an AI programming assistant, utilizing the DeepSeek Coder model, " + "developed by DeepSeek Company, and you only answer questions related to computer science. " + "For politically sensitive questions, security and privacy issues, " + "and other non-computer science questions, you will refuse to answer.\n" + ), +) + + +register_template( + name="default", + format_user=StringFormatter(slots=["Human: {{content}}", {"eos_token"}, "\nAssistant:"]), + format_assistant=StringFormatter(slots=["{{content}}", {"eos_token"}, "\n"]), + format_system=StringFormatter(slots=["System: {{content}}", {"eos_token"}, "\n"]), + replace_jinja_template=True, +) + + +register_template( + name="dots_ocr", + format_user=StringFormatter(slots=["<|user|>{{content}}<|endofuser|><|assistant|>"]), + format_assistant=StringFormatter(slots=["{{content}}<|endofassistant|>"]), + format_system=StringFormatter(slots=["<|system|>{{content}}<|endofsystem|>\n"]), + stop_words=["<|endofassistant|>"], + efficient_eos=True, + mm_plugin=get_mm_plugin( + name="qwen2_vl", + image_token="<|imgpad|>", + video_token="<|vidpad|>", + vision_bos_token="<|img|>", + vision_eos_token="<|endofimg|>", + ), +) + + +register_template( + name="empty", + format_assistant=StringFormatter(slots=["{{content}}"]), +) + + +# copied from chatml template +register_template( + name="ernie", + format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n\n<|im_start|>assistant\n"]), + format_assistant=StringFormatter(slots=["{{content}}<|im_end|>\n\n"]), + format_system=StringFormatter(slots=["<|im_start|>system\n{{content}}<|im_end|>\n\n"]), + format_observation=StringFormatter(slots=["<|im_start|>tool\n{{content}}<|im_end|>\n\n<|im_start|>assistant\n"]), + default_system="\nthink_mode=True\n", + stop_words=["<|im_end|>"], +) + + +register_template( + name="ernie_nothink", + format_user=StringFormatter(slots=["User: {{content}}\nAssistant: "]), + format_assistant=StringFormatter(slots=["{{content}}<|end_of_sentence|>"]), + format_system=StringFormatter(slots=["{{content}}\n"]), + format_prefix=EmptyFormatter(slots=["<|begin_of_sentence|>"]), + stop_words=["<|end_of_sentence|>"], +) + + +register_template( + name="ernie_vl", + format_user=StringFormatter(slots=["User: {{content}}"]), + format_assistant=StringFormatter(slots=["\nAssistant: {{content}}<|end_of_sentence|>"]), + format_system=StringFormatter(slots=["{{content}}\n"]), + stop_words=["<|end_of_sentence|>"], + replace_eos=True, + replace_jinja_template=True, + template_class=ReasoningTemplate, + mm_plugin=get_mm_plugin(name="ernie_vl", image_token="<|IMAGE_PLACEHOLDER|>", video_token="<|VIDEO_PLACEHOLDER|>"), +) + + +register_template( + name="exaone", + format_user=StringFormatter(slots=["[|user|]{{content}}\n[|assistant|]"]), + format_assistant=StringFormatter(slots=["{{content}}", {"eos_token"}, "\n"]), + format_system=StringFormatter(slots=["[|system|]{{content}}[|endofturn|]\n"]), +) + + +register_template( + name="falcon", + format_user=StringFormatter(slots=["User: {{content}}\nFalcon:"]), + format_assistant=StringFormatter(slots=["{{content}}\n"]), + efficient_eos=True, +) + + +# copied from chatml template +register_template( + name="falcon_h1", + format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + format_assistant=StringFormatter(slots=["{{content}}<|im_end|>\n"]), + format_system=StringFormatter(slots=["<|im_start|>system\n{{content}}<|im_end|>\n"]), + format_observation=StringFormatter(slots=["<|im_start|>tool\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=["<|im_end|>", "<|end_of_text|>"], +) + + +register_template( + name="fewshot", + format_assistant=StringFormatter(slots=["{{content}}\n\n"]), + efficient_eos=True, + replace_jinja_template=True, +) + + +register_template( + name="gemma", + format_user=StringFormatter(slots=["user\n{{content}}\nmodel\n"]), + format_assistant=StringFormatter(slots=["{{content}}\n"]), + format_system=StringFormatter(slots=["{{content}}\n\n"]), + format_observation=StringFormatter( + slots=["tool\n{{content}}\nmodel\n"] + ), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=[""], + replace_eos=True, + template_class=Llama2Template, +) + + +# copied from gemma template +register_template( + name="gemma2", + format_user=StringFormatter(slots=["user\n{{content}}\nmodel\n"]), + format_assistant=StringFormatter(slots=["{{content}}\n"]), + format_system=StringFormatter(slots=["{{content}}\n\n"]), + format_observation=StringFormatter( + slots=["tool\n{{content}}\nmodel\n"] + ), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=["", ""], + efficient_eos=True, + template_class=Llama2Template, +) + + +# copied from gemma template +register_template( + name="gemma3", + format_user=StringFormatter(slots=["user\n{{content}}\nmodel\n"]), + format_assistant=StringFormatter(slots=["{{content}}\n"]), + format_system=StringFormatter(slots=["{{content}}\n\n"]), + format_observation=StringFormatter( + slots=["tool\n{{content}}\nmodel\n"] + ), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=[""], + replace_eos=True, + mm_plugin=get_mm_plugin("gemma3", image_token=""), + template_class=Llama2Template, +) + + +register_template( + name="gemma3n", + format_user=StringFormatter(slots=["user\n{{content}}\nmodel\n"]), + format_assistant=StringFormatter(slots=["{{content}}\n"]), + format_system=StringFormatter(slots=["{{content}}\n\n"]), + format_observation=StringFormatter( + slots=["tool\n{{content}}\nmodel\n"] + ), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=[""], + replace_eos=True, + mm_plugin=get_mm_plugin("gemma3n", image_token="", audio_token=""), + template_class=Llama2Template, +) + + +register_template( + name="gemma4", + format_user=StringFormatter(slots=["<|turn>user\n{{content}}\n<|turn>model\n"]), + format_assistant=StringFormatter(slots=["{{content}}\n"]), + format_system=StringFormatter( + slots=["<|turn>system\n<|think|>{{content}}\n"] + ), # default thought signal contained + format_observation=StringFormatter( + slots=["<|turn>tool\n{{content}}\n<|turn>model\n"] + ), # seem not consistent with the chattemplate + format_tools=ToolFormatter(tool_format="gemma4"), + format_function=FunctionFormatter(slots=["<|tool>{{content}}"], tool_format="gemma4"), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=[""], + default_system="You are a helpful assistant.", # important for thinking + thought_words=("<|channel>thought\n", ""), + replace_eos=True, + mm_plugin=get_mm_plugin( + "gemma4", + image_token="<|image|>", + video_token="<|video|>", + ), + template_class=ReasoningTemplate, +) + + +register_template( + name="gemma4n", + format_user=StringFormatter(slots=["<|turn>user\n{{content}}\n<|turn>model\n"]), + format_assistant=StringFormatter(slots=["{{content}}\n"]), + format_system=StringFormatter( + slots=["<|turn>system\n<|think|>{{content}}\n"] + ), # default thought signal contained + format_observation=StringFormatter(slots=["<|turn>tool\n{{content}}\n<|turn>model\n"]), + format_tools=ToolFormatter(tool_format="gemma4"), + format_function=FunctionFormatter(slots=["<|tool>{{content}}"], tool_format="gemma4"), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=[""], + default_system="You are a helpful assistant.", # important for thinking + thought_words=("<|channel>thought\n", ""), + replace_eos=True, + mm_plugin=get_mm_plugin( + "gemma4", + image_token="<|image|>", + video_token="<|video|>", + audio_token="<|audio|>", + ), + template_class=ReasoningTemplate, +) + + +register_template( + name="glm4", + format_user=StringFormatter(slots=["<|user|>\n{{content}}<|assistant|>"]), + format_assistant=StringFormatter(slots=["\n{{content}}"]), + format_system=StringFormatter(slots=["<|system|>\n{{content}}"]), + format_function=FunctionFormatter(slots=["{{content}}"], tool_format="glm4"), + format_observation=StringFormatter(slots=["<|observation|>\n{{content}}<|assistant|>"]), + format_tools=ToolFormatter(tool_format="glm4"), + format_prefix=EmptyFormatter(slots=["[gMASK]"]), + stop_words=["<|user|>", "<|observation|>"], + efficient_eos=True, +) + + +# copied from glm4 template +register_template( + name="glm4_moe", + format_user=StringFormatter(slots=["<|user|>\n{{content}}<|assistant|>"]), + format_assistant=StringFormatter(slots=["\n{{content}}"]), + format_system=StringFormatter(slots=["<|system|>\n{{content}}"]), + format_function=FunctionFormatter(slots=["{{content}}"], tool_format="glm4_moe"), + format_observation=StringFormatter(slots=["<|observation|>\n{{content}}<|assistant|>"]), + format_tools=ToolFormatter(tool_format="glm4_moe"), + format_prefix=EmptyFormatter(slots=["[gMASK]"]), + stop_words=["<|user|>", "<|observation|>"], + efficient_eos=True, + template_class=ReasoningTemplate, +) + + +# copied from glm4 template +register_template( + name="glm4v", + format_user=StringFormatter(slots=["<|user|>\n{{content}}<|assistant|>"]), + format_assistant=StringFormatter(slots=["\n{{content}}"]), + format_system=StringFormatter(slots=["<|system|>\n{{content}}"]), + format_function=FunctionFormatter(slots=["{{content}}"], tool_format="glm4"), + format_observation=StringFormatter(slots=["<|observation|>\n{{content}}<|assistant|>"]), + format_tools=ToolFormatter(tool_format="glm4"), + format_prefix=EmptyFormatter(slots=["[gMASK]"]), + stop_words=["<|user|>", "<|observation|>", ""], + efficient_eos=True, + mm_plugin=get_mm_plugin(name="glm4v", image_token="<|image|>", video_token="<|video|>"), + template_class=ReasoningTemplate, +) + + +# copied from glm4 template +register_template( + name="glm4_5v", + format_user=StringFormatter(slots=["<|user|>\n{{content}}<|assistant|>"]), + format_assistant=StringFormatter(slots=["\n{{content}}"]), + format_system=StringFormatter(slots=["<|system|>\n{{content}}"]), + format_function=FunctionFormatter(slots=["{{content}}"], tool_format="glm4_moe"), + format_observation=StringFormatter(slots=["<|observation|>\n{{content}}<|assistant|>"]), + format_tools=ToolFormatter(tool_format="glm4_moe"), + format_prefix=EmptyFormatter(slots=["[gMASK]"]), + stop_words=["<|user|>", "<|observation|>", ""], + efficient_eos=True, + mm_plugin=get_mm_plugin(name="glm4v", image_token="<|image|>", video_token="<|video|>"), + template_class=ReasoningTemplate, +) + + +# copied from glm4 template +register_template( + name="glm_ocr", + format_user=StringFormatter(slots=["<|user|>\n{{content}}<|assistant|>"]), + format_assistant=StringFormatter(slots=["\n{{content}}"]), + format_system=StringFormatter(slots=["<|system|>\n{{content}}"]), + format_function=FunctionFormatter(slots=["{{content}}"], tool_format="glm4"), + format_observation=StringFormatter(slots=["<|observation|>\n{{content}}<|assistant|>"]), + format_tools=ToolFormatter(tool_format="glm4"), + format_prefix=EmptyFormatter(slots=["[gMASK]"]), + stop_words=["<|user|>", "<|observation|>"], + efficient_eos=True, + mm_plugin=get_mm_plugin(name="glm4v", image_token="<|image|>", video_token="<|video|>"), +) + + +# copied from glm4_moe template +register_template( + name="glm4_7", + format_user=StringFormatter(slots=["<|user|>\n{{content}}<|assistant|>"]), + format_assistant=StringFormatter(slots=["\n{{content}}"]), + format_system=StringFormatter(slots=["<|system|>\n{{content}}"]), + format_function=FunctionFormatter(slots=["{{content}}"], tool_format="glm4_moe"), + format_observation=StringFormatter(slots=["<|observation|>\n{{content}}<|assistant|>"]), + format_tools=ToolFormatter(tool_format="glm4_moe"), + format_prefix=EmptyFormatter(slots=["[gMASK]"]), + stop_words=["<|user|>", "<|observation|>"], + thought_words=("", ""), + efficient_eos=True, + template_class=Glm47ReasoningTemplate, +) + + +# copied from glm4 template +register_template( + name="glmz1", + format_user=StringFormatter(slots=["<|user|>\n{{content}}<|assistant|>"]), + format_assistant=StringFormatter(slots=["\n{{content}}"]), + format_system=StringFormatter(slots=["<|system|>\n{{content}}"]), + format_function=FunctionFormatter(slots=["{{content}}"], tool_format="glm4"), + format_observation=StringFormatter(slots=["<|observation|>\n{{content}}<|assistant|>"]), + format_tools=ToolFormatter(tool_format="glm4"), + format_prefix=EmptyFormatter(slots=["[gMASK]"]), + stop_words=["<|user|>", "<|observation|>"], + efficient_eos=True, + template_class=ReasoningTemplate, +) + + +register_template( + name="gpt_oss", + format_user=StringFormatter(slots=["<|start|>user<|message|>{{content}}<|end|><|start|>assistant"]), + format_assistant=StringFormatter(slots=["{{content}}"]), + format_system=StringFormatter(slots=["<|start|>system<|message|>{{content}}<|end|>"]), + default_system="You are ChatGPT, a large language model trained by OpenAI.", + thought_words=("<|channel|>analysis<|message|>", "<|end|><|start|>assistant<|channel|>final<|message|>"), + efficient_eos=True, + template_class=ReasoningTemplate, +) + + +register_template( + name="granite3", + format_user=StringFormatter( + slots=[ + "<|start_of_role|>user<|end_of_role|>{{content}}<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>" + ] + ), + format_assistant=StringFormatter(slots=["{{content}}<|end_of_text|>\n"]), + format_system=StringFormatter(slots=["<|start_of_role|>system<|end_of_role|>{{content}}<|end_of_text|>\n"]), +) + + +register_template( + name="granite3_vision", + format_user=StringFormatter(slots=["<|user|>\n{{content}}\n<|assistant|>\n"]), + format_system=StringFormatter(slots=["<|system|>\n{{content}}\n"]), + default_system=( + "A chat between a curious user and an artificial intelligence assistant. " + "The assistant gives helpful, detailed, and polite answers to the user's questions." + ), + mm_plugin=get_mm_plugin(name="llava_next", image_token=""), +) + + +register_template( + name="granite4", + format_user=StringFormatter( + slots=[ + "<|start_of_role|>user<|end_of_role|>{{content}}<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>" + ] + ), + format_assistant=StringFormatter(slots=["{{content}}<|end_of_text|>\n"]), + format_system=StringFormatter(slots=["<|start_of_role|>system<|end_of_role|>{{content}}<|end_of_text|>\n"]), + format_function=FunctionFormatter(slots=["{{content}}<|end_of_text|>\n"], tool_format="default"), + format_observation=StringFormatter( + slots=["<|start_of_role|>tool<|end_of_role|>{{content}}<|end_of_text|>\n<|start_of_role|>assistant\n"] + ), + format_tools=ToolFormatter(tool_format="default"), + stop_words=["<|end_of_text|>"], + default_system="You are Granite, developed by IBM. You are a helpful AI assistant.", +) + + +register_template( + name="index", + format_user=StringFormatter(slots=["reserved_0{{content}}reserved_1"]), + format_system=StringFormatter(slots=["{{content}}"]), + efficient_eos=True, +) + + +register_template( + name="hunyuan", + format_user=StringFormatter(slots=["{{content}}<|extra_0|>"]), + format_assistant=StringFormatter(slots=["{{content}}<|eos|>"]), + format_system=StringFormatter(slots=["{{content}}<|extra_4|>"]), + format_prefix=EmptyFormatter(slots=["<|startoftext|>"]), + stop_words=["<|eos|>"], +) + + +register_template( + name="hunyuan_small", + format_user=StringFormatter(slots=["<｜hy_User｜>{{content}}<｜hy_place▁holder▁no▁8｜>"]), + format_assistant=StringFormatter(slots=["{{content}}<｜hy_place▁holder▁no▁2｜>"]), + format_system=StringFormatter(slots=["{{content}}<｜hy_place▁holder▁no▁3｜>"]), + format_prefix=EmptyFormatter(slots=["<｜hy_begin▁of▁sentence｜>"]), + stop_words=["<｜hy_place▁holder▁no▁2｜>"], +) + + +# The following two templates are copied from the official Hy-MT2 chat templates: +# https://github.com/Tencent-Hunyuan/Hy-MT2/blob/main/train/llama_factory_support/hy_dense_template.py +register_template( + name="hy_dense_1_8b", + format_user=StringFormatter(slots=["<｜hy_User｜>{{content}}"]), + format_assistant=StringFormatter(slots=["<｜hy_Assistant｜>{{content}}"]), + format_system=StringFormatter(slots=["{{content}}<｜hy_place▁holder▁no▁3｜>"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=["<｜hy_place▁holder▁no▁2｜>"], + efficient_eos=True, +) + + +register_template( + name="hy_dense_7b", + format_user=StringFormatter(slots=["{{content}}<|extra_0|>"]), + format_assistant=StringFormatter(slots=["{{content}}"]), + format_system=StringFormatter(slots=["{{content}}<|extra_4|>"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + stop_words=["<|eos|>"], + efficient_eos=True, +) + + +register_template( + name="intern2", + format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + format_assistant=StringFormatter(slots=["{{content}}<|im_end|>\n"]), + format_system=StringFormatter(slots=["<|im_start|>system\n{{content}}<|im_end|>\n"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + default_system=( + "You are an AI assistant whose name is InternLM (书生·浦语).\n" + "- InternLM (书生·浦语) is a conversational language model that is developed by Shanghai AI Laboratory " + "(上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n" + "- InternLM (书生·浦语) can understand and communicate fluently in the language " + "chosen by the user such as English and 中文." + ), + stop_words=["<|im_end|>"], +) + + +register_template( + name="intern_vl", + format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]), + format_assistant=StringFormatter(slots=["{{content}}<|im_end|>\n"]), + format_system=StringFormatter(slots=["<|im_start|>system\n{{content}}<|im_end|>\n"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + default_system=( + "你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。" + ), + stop_words=["<|im_end|>"], + mm_plugin=get_mm_plugin(name="intern_vl", image_token="", video_token="