Skip to content

Commit 8ebc435

Browse files
Fix llava_onevision tests (#37280)
* Fix llava_onevision tests * Trigger tests
1 parent ad3d157 commit 8ebc435

File tree

1 file changed

+8
-8
lines changed

1 file changed

+8
-8
lines changed

tests/models/llava_onevision/test_processor_llava_onevision.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -39,17 +39,19 @@
3939
class LlavaOnevisionProcessorTest(ProcessorTesterMixin, unittest.TestCase):
4040
processor_class = LlavaOnevisionProcessor
4141

42-
def setUp(self):
43-
self.tmpdirname = tempfile.mkdtemp()
42+
@classmethod
43+
def setUpClass(cls):
44+
cls.tmpdirname = tempfile.mkdtemp()
45+
cls.addClassCleanup(lambda tempdir=cls.tmpdirname: shutil.rmtree(tempdir))
4446
image_processor = LlavaOnevisionImageProcessor()
4547
video_processor = LlavaOnevisionVideoProcessor()
4648
tokenizer = Qwen2TokenizerFast.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
47-
processor_kwargs = self.prepare_processor_dict()
49+
processor_kwargs = cls.prepare_processor_dict()
4850

4951
processor = LlavaOnevisionProcessor(
5052
video_processor=video_processor, image_processor=image_processor, tokenizer=tokenizer, **processor_kwargs
5153
)
52-
processor.save_pretrained(self.tmpdirname)
54+
processor.save_pretrained(cls.tmpdirname)
5355

5456
def get_tokenizer(self, **kwargs):
5557
return AutoProcessor.from_pretrained(self.tmpdirname, **kwargs).tokenizer
@@ -60,7 +62,8 @@ def get_image_processor(self, **kwargs):
6062
def get_video_processor(self, **kwargs):
6163
return AutoProcessor.from_pretrained(self.tmpdirname, **kwargs).video_processor
6264

63-
def prepare_processor_dict(self):
65+
@staticmethod
66+
def prepare_processor_dict():
6467
return {
6568
"chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + ' '}}{# Render all images first #}{% for content in message['content'] | selectattr('type', 'equalto', 'image') %}{{ '<image>' }}{% endfor %}{# Render all video then #}{% for content in message['content'] | selectattr('type', 'equalto', 'video') %}{{ '<video>' }}{% endfor %}{# Render all text next #}{% if message['role'] != 'assistant' %}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{{ '\n' + content['text'] }}{% endfor %}{% else %}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{% generation %}{{ '\n' + content['text'] }}{% endgeneration %}{% endfor %}{% endif %}{{'<|im_end|>'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
6669
"num_image_tokens": 6,
@@ -88,9 +91,6 @@ def test_chat_template_is_saved(self):
8891
processor_dict = self.prepare_processor_dict()
8992
self.assertTrue(processor_loaded.chat_template == processor_dict.get("chat_template", None))
9093

91-
def tearDown(self):
92-
shutil.rmtree(self.tmpdirname)
93-
9494
def test_chat_template(self):
9595
processor = AutoProcessor.from_pretrained("llava-hf/llava-onevision-qwen2-7b-ov-hf")
9696
expected_prompt = "<|im_start|>user <image>\nWhat is shown in this image?<|im_end|><|im_start|>assistant\n"

0 commit comments

Comments
 (0)