updata to v0.9.1_stable

tests/data/processors/test_unsupervised.py

@@ -22,11 +22,11 @@ from transformers import AutoTokenizer
 from llamafactory.train.test_utils import load_train_dataset
 
 
-DEMO_DATA = os.environ.get("DEMO_DATA", "llamafactory/demo_data")
+DEMO_DATA = os.getenv("DEMO_DATA", "llamafactory/demo_data")
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
-TINY_DATA = os.environ.get("TINY_DATA", "llamafactory/tiny-supervised-dataset")
+TINY_DATA = os.getenv("TINY_DATA", "llamafactory/tiny-supervised-dataset")
 
 TRAIN_ARGS = {
     "model_name_or_path": TINY_LLAMA,

tests/data/test_mm_plugin.py

@@ -31,9 +31,9 @@ if TYPE_CHECKING:
     from llamafactory.data.mm_plugin import BasePlugin
 
 
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
 MM_MESSAGES = [
     {"role": "user", "content": "<image>What is in this image?"},
@@ -61,7 +61,7 @@ INPUT_IDS = [0, 1, 2, 3, 4]
 
 LABELS = [0, 1, 2, 3, 4]
 
-SEQLENS = [1024]
+BATCH_IDS = [[1] * 1024]
 
 
 def _get_mm_inputs(processor: "ProcessorMixin") -> Dict[str, "torch.Tensor"]:
@@ -74,6 +74,10 @@ def _is_close(batch_a: Dict[str, Any], batch_b: Dict[str, Any]) -> None:
     for key in batch_a.keys():
         if isinstance(batch_a[key], torch.Tensor):
             assert torch.allclose(batch_a[key], batch_b[key], rtol=1e-4, atol=1e-5)
+        elif isinstance(batch_a[key], list) and all(isinstance(item, torch.Tensor) for item in batch_a[key]):
+            assert len(batch_a[key]) == len(batch_b[key])
+            for tensor_a, tensor_b in zip(batch_a[key], batch_b[key]):
+                assert torch.allclose(tensor_a, tensor_b, rtol=1e-4, atol=1e-5)
         else:
             assert batch_a[key] == batch_b[key]
 
@@ -101,7 +105,7 @@ def _check_plugin(
         expected_labels,
     )
     _is_close(
-        plugin.get_mm_inputs(IMAGES, NO_VIDEOS, IMGLENS, NO_VIDLENS, SEQLENS, processor),
+        plugin.get_mm_inputs(IMAGES, NO_VIDEOS, IMGLENS, NO_VIDLENS, BATCH_IDS, processor),
         expected_mm_inputs,
     )
     # test text_messages
@@ -111,7 +115,7 @@ def _check_plugin(
         LABELS,
     )
     _is_close(
-        plugin.get_mm_inputs(NO_IMAGES, NO_VIDEOS, NO_IMGLENS, NO_VIDLENS, SEQLENS, processor),
+        plugin.get_mm_inputs(NO_IMAGES, NO_VIDEOS, NO_IMGLENS, NO_VIDLENS, BATCH_IDS, processor),
         expected_no_mm_inputs,
     )
 
@@ -179,6 +183,28 @@ def test_paligemma_plugin():
     _check_plugin(**check_inputs)
 
 
+def test_pixtral_plugin():
+    tokenizer, processor = _load_tokenizer_module(model_name_or_path="mistral-community/pixtral-12b")
+    pixtral_plugin = get_mm_plugin(name="pixtral", image_token="[IMG]")
+    image_slice_height, image_slice_width = 2, 2
+    check_inputs = {"plugin": pixtral_plugin, "tokenizer": tokenizer, "processor": processor}
+    check_inputs["expected_mm_messages"] = [
+        {
+            key: value.replace(
+                "<image>",
+                ("{}[IMG_BREAK]".format("[IMG]" * image_slice_width) * image_slice_height).rsplit("[IMG_BREAK]", 1)[0]
+                + "[IMG_END]",
+            )
+            for key, value in message.items()
+        }
+        for message in MM_MESSAGES
+    ]
+    check_inputs["expected_mm_inputs"] = _get_mm_inputs(processor)
+    check_inputs["expected_mm_inputs"].pop("image_sizes")
+    check_inputs["expected_mm_inputs"]["pixel_values"] = check_inputs["expected_mm_inputs"]["pixel_values"][0]
+    _check_plugin(**check_inputs)
+
+
 def test_qwen2_vl_plugin():
     tokenizer, processor = _load_tokenizer_module(model_name_or_path="Qwen/Qwen2-VL-7B-Instruct")
     qwen2_vl_plugin = get_mm_plugin(name="qwen2_vl", image_token="<|image_pad|>")

tests/data/test_template.py

@@ -27,9 +27,9 @@ if TYPE_CHECKING:
     from transformers import PreTrainedTokenizer
 
 
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
 MESSAGES = [
     {"role": "user", "content": "How are you"},

tests/e2e/test_chat.py

@@ -17,7 +17,7 @@ import os
 from llamafactory.chat import ChatModel
 
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
 INFER_ARGS = {
     "model_name_or_path": TINY_LLAMA,

tests/e2e/test_train.py

@@ -19,11 +19,11 @@ import pytest
 from llamafactory.train.tuner import export_model, run_exp
 
 
-DEMO_DATA = os.environ.get("DEMO_DATA", "llamafactory/demo_data")
+DEMO_DATA = os.getenv("DEMO_DATA", "llamafactory/demo_data")
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
-TINY_LLAMA_ADAPTER = os.environ.get("TINY_LLAMA_ADAPTER", "llamafactory/tiny-random-Llama-3-lora")
+TINY_LLAMA_ADAPTER = os.getenv("TINY_LLAMA_ADAPTER", "llamafactory/tiny-random-Llama-3-lora")
 
 TRAIN_ARGS = {
     "model_name_or_path": TINY_LLAMA,
@@ -44,10 +44,9 @@ INFER_ARGS = {
     "finetuning_type": "lora",
     "template": "llama3",
     "infer_dtype": "float16",
-    "export_dir": "llama3_export",
 }
 
-OS_NAME = os.environ.get("OS_NAME", "")
+OS_NAME = os.getenv("OS_NAME", "")
 
 
 @pytest.mark.parametrize(
@@ -61,11 +60,12 @@ OS_NAME = os.environ.get("OS_NAME", "")
     ],
 )
 def test_run_exp(stage: str, dataset: str):
-    output_dir = "train_{}".format(stage)
+    output_dir = os.path.join("output", f"train_{stage}")
     run_exp({"stage": stage, "dataset": dataset, "output_dir": output_dir, **TRAIN_ARGS})
     assert os.path.exists(output_dir)
 
 
 def test_export():
-    export_model(INFER_ARGS)
-    assert os.path.exists("llama3_export")
+    export_dir = os.path.join("output", "llama3_export")
+    export_model({"export_dir": export_dir, **INFER_ARGS})
+    assert os.path.exists(export_dir)

tests/model/model_utils/test_attention.py

@@ -19,7 +19,7 @@ from transformers.utils import is_flash_attn_2_available, is_torch_sdpa_availabl
 from llamafactory.train.test_utils import load_infer_model
 
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
 INFER_ARGS = {
     "model_name_or_path": TINY_LLAMA,

tests/model/model_utils/test_checkpointing.py

@@ -20,7 +20,7 @@ from llamafactory.extras.misc import get_current_device
 from llamafactory.train.test_utils import load_train_model
 
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
 TRAIN_ARGS = {
     "model_name_or_path": TINY_LLAMA,
@@ -54,7 +54,7 @@ def test_checkpointing_disable():
 def test_unsloth_gradient_checkpointing():
     model = load_train_model(use_unsloth_gc=True, **TRAIN_ARGS)
     for module in filter(lambda m: hasattr(m, "gradient_checkpointing"), model.modules()):
-        assert module._gradient_checkpointing_func.__self__.__name__ == "UnslothGradientCheckpointing"  # classmethod
+        assert module._gradient_checkpointing_func.__self__.__name__ == "UnslothGradientCheckpointing"
 
 
 def test_upcast_layernorm():

tests/model/test_base.py

@@ -16,17 +16,12 @@ import os
 
 import pytest
 
-from llamafactory.train.test_utils import (
-    compare_model,
-    load_infer_model,
-    load_reference_model,
-    patch_valuehead_model,
-)
+from llamafactory.train.test_utils import compare_model, load_infer_model, load_reference_model, patch_valuehead_model
 
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
-TINY_LLAMA_VALUEHEAD = os.environ.get("TINY_LLAMA_VALUEHEAD", "llamafactory/tiny-random-Llama-3-valuehead")
+TINY_LLAMA_VALUEHEAD = os.getenv("TINY_LLAMA_VALUEHEAD", "llamafactory/tiny-random-Llama-3-valuehead")
 
 INFER_ARGS = {
     "model_name_or_path": TINY_LLAMA,

tests/model/test_freeze.py

@@ -19,7 +19,7 @@ import torch
 from llamafactory.train.test_utils import load_infer_model, load_train_model
 
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
 TRAIN_ARGS = {
     "model_name_or_path": TINY_LLAMA,

tests/model/test_full.py

@@ -19,7 +19,7 @@ import torch
 from llamafactory.train.test_utils import load_infer_model, load_train_model
 
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
 TRAIN_ARGS = {
     "model_name_or_path": TINY_LLAMA,

tests/model/test_lora.py

@@ -27,11 +27,11 @@ from llamafactory.train.test_utils import (
 )
 
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
-TINY_LLAMA_ADAPTER = os.environ.get("TINY_LLAMA_ADAPTER", "llamafactory/tiny-random-Llama-3-lora")
+TINY_LLAMA_ADAPTER = os.getenv("TINY_LLAMA_ADAPTER", "llamafactory/tiny-random-Llama-3-lora")
 
-TINY_LLAMA_VALUEHEAD = os.environ.get("TINY_LLAMA_VALUEHEAD", "llamafactory/tiny-random-Llama-3-valuehead")
+TINY_LLAMA_VALUEHEAD = os.getenv("TINY_LLAMA_VALUEHEAD", "llamafactory/tiny-random-Llama-3-valuehead")
 
 TRAIN_ARGS = {
     "model_name_or_path": TINY_LLAMA,

tests/model/test_pissa.py

@@ -19,9 +19,9 @@ import pytest
 from llamafactory.train.test_utils import compare_model, load_infer_model, load_reference_model, load_train_model
 
 
-TINY_LLAMA = os.environ.get("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
+TINY_LLAMA = os.getenv("TINY_LLAMA", "llamafactory/tiny-random-Llama-3")
 
-TINY_LLAMA_PISSA = os.environ.get("TINY_LLAMA_ADAPTER", "llamafactory/tiny-random-Llama-3-pissa")
+TINY_LLAMA_PISSA = os.getenv("TINY_LLAMA_ADAPTER", "llamafactory/tiny-random-Llama-3-pissa")
 
 TRAIN_ARGS = {
     "model_name_or_path": TINY_LLAMA,
@@ -49,10 +49,10 @@ INFER_ARGS = {
     "infer_dtype": "float16",
 }
 
-OS_NAME = os.environ.get("OS_NAME", "")
+OS_NAME = os.getenv("OS_NAME", "")
 
 
-@pytest.mark.xfail(OS_NAME.startswith("windows"), reason="Known connection error on Windows.")
+@pytest.mark.xfail(reason="PiSSA initialization is not stable in different platform.")
 def test_pissa_train():
     model = load_train_model(**TRAIN_ARGS)
     ref_model = load_reference_model(TINY_LLAMA_PISSA, TINY_LLAMA_PISSA, use_pissa=True, is_trainable=True)