youssefedweqd/working

1# Copyright 2025 the LlamaFactory team.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import os16 17import pytest18 19from llamafactory.train.test_utils import compare_model, load_infer_model, load_reference_model, patch_valuehead_model20 21 22TINY_LLAMA3 = os.getenv("TINY_LLAMA3", "llamafactory/tiny-random-Llama-3")23 24TINY_LLAMA_VALUEHEAD = os.getenv("TINY_LLAMA_VALUEHEAD", "llamafactory/tiny-random-Llama-3-valuehead")25 26INFER_ARGS = {27    "model_name_or_path": TINY_LLAMA3,28    "template": "llama3",29    "infer_dtype": "float16",30}31 32 33@pytest.fixture34def fix_valuehead_cpu_loading():35    patch_valuehead_model()36 37 38def test_base():39    model = load_infer_model(**INFER_ARGS)40    ref_model = load_reference_model(TINY_LLAMA3)41    compare_model(model, ref_model)42 43 44@pytest.mark.usefixtures("fix_valuehead_cpu_loading")45def test_valuehead():46    model = load_infer_model(add_valuehead=True, **INFER_ARGS)47    ref_model = load_reference_model(TINY_LLAMA_VALUEHEAD, add_valuehead=True)48    compare_model(model, ref_model)49