3939class 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>\n What is shown in this image?<|im_end|><|im_start|>assistant\n "
0 commit comments