From a8a0ab10a2ed6efa70258c7c7a5279e93fabd7f4 Mon Sep 17 00:00:00 2001 From: Manuel de Prada Corral Date: Tue, 24 Jun 2025 20:24:24 +0200 Subject: [PATCH 1/2] create test for #38916 --- tests/generation/test_utils.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index e92d1e1ec77a..5ec354dfb90e 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -16,12 +16,15 @@ import collections import copy import datetime +import io import gc import inspect import random import tempfile import unittest import warnings +from contextlib import redirect_stdout +from pathlib import Path import numpy as np import pytest @@ -4985,6 +4988,34 @@ def test_custom_generate_requires_trust_remote_code(self): with self.assertRaises(ValueError): model.generate(**model_inputs, custom_generate="transformers-community/custom_generate_example") + def test_custom_generate_local_directory(self): + """Tests that custom_generate works with local directories containing importable relative modules""" + with tempfile.TemporaryDirectory() as tmp_dir: + custom_generate_dir = Path(tmp_dir) / "custom_generate" + custom_generate_dir.mkdir() + with open(custom_generate_dir / "generate.py", "w") as f: + f.write( + "from .helper import ret_success\n" + "def generate(*args, **kwargs):\n" + " return ret_success()\n" + ) + with open(custom_generate_dir / "helper.py", "w") as f: + f.write( + "def ret_success():\n" + " return \"success\"\n" + ) + model = AutoModelForCausalLM.from_pretrained( + "hf-internal-testing/tiny-random-MistralForCausalLM", device_map="auto" + ) + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-MistralForCausalLM") + model_inputs = tokenizer("Hello, world!", return_tensors="pt").to(model.device) + value = model.generate( + **model_inputs, + custom_generate=str(tmp_dir), + trust_remote_code=True, + ) + assert value == "success" + @require_torch class TokenHealingTestCase(unittest.TestCase): From 63365ca9e436d73387c9710c38e376ee6594f49c Mon Sep 17 00:00:00 2001 From: Manuel de Prada Corral Date: Tue, 24 Jun 2025 20:27:18 +0200 Subject: [PATCH 2/2] ruff --- tests/generation/test_utils.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index 5ec354dfb90e..973ce7ab422c 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -16,14 +16,12 @@ import collections import copy import datetime -import io import gc import inspect import random import tempfile import unittest import warnings -from contextlib import redirect_stdout from pathlib import Path import numpy as np @@ -4994,16 +4992,9 @@ def test_custom_generate_local_directory(self): custom_generate_dir = Path(tmp_dir) / "custom_generate" custom_generate_dir.mkdir() with open(custom_generate_dir / "generate.py", "w") as f: - f.write( - "from .helper import ret_success\n" - "def generate(*args, **kwargs):\n" - " return ret_success()\n" - ) + f.write("from .helper import ret_success\ndef generate(*args, **kwargs):\n return ret_success()\n") with open(custom_generate_dir / "helper.py", "w") as f: - f.write( - "def ret_success():\n" - " return \"success\"\n" - ) + f.write('def ret_success():\n return "success"\n') model = AutoModelForCausalLM.from_pretrained( "hf-internal-testing/tiny-random-MistralForCausalLM", device_map="auto" )