Update tools.py

tools.py

@@ -74,54 +74,32 @@ class translate_everything(Tool):
         return f"The translated sentence is : {translated_sentence}"
 
 class multimodal_interpreter(Tool):
-    name="multimodal_tool"
+    name = "multimodal_tool"
     description = "Allows you to answer any question which relies on image or video input."
     inputs = {
-        'image': {"type": "image", "description": "the image or video of interest"},
-        'prompt': {"type": "string", "description": "Any specific question you have on the image. For example, the prompt can be : Describe this image."}
+        'image': {"type": "image", "description": "The image or video of interest"},
+        'prompt': {"type": "string", "description": "Any specific question you have on the image. For example: Describe this image."}
     }
     output_type = "string"
-    
+
     def forward(self, prompt, image):
-        
-        # default: Load the model on the available device(s)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         model = Qwen2VLForConditionalGeneration.from_pretrained(
-            "Qwen/Qwen2-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
+        "Qwen/Qwen2-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
         )
-        
-        # We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
-        # model = Qwen2VLForConditionalGeneration.from_pretrained(
-        #     "Qwen/Qwen2-VL-7B-Instruct",
-        #     torch_dtype=torch.bfloat16,
-        #     attn_implementation="flash_attention_2",
-        #     device_map="auto",
-        # )
-        
-        # default processer
         processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
         
-        # The default range for the number of visual tokens per image in the model is 4-16384. You can set min_pixels and max_pixels according to your needs, such as a token count range of 256-1280, to balance speed and memory usage.
-        # min_pixels = 256*28*28
-        # max_pixels = 1280*28*28
-        # processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
-        
         messages = [
             {
                 "role": "user",
                 "content": [
-                    {
-                        "type": "image",
-                        "image": {image},
-                    },
-                    {"type": "text", "text": {prompt}},
+                    {"type": "image", "image": image},
+                    {"type": "text", "text": prompt},
                 ],
             }
         ]
-        
-        # Preparation for inference
-        text = processor.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
+
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         image_inputs, video_inputs = process_vision_info(messages)
         inputs = processor(
             text=[text],
@@ -129,18 +107,18 @@ class multimodal_interpreter(Tool):
             videos=video_inputs,
             padding=True,
             return_tensors="pt",
-        )
-        inputs = inputs.to("cuda")
-        
-        # Inference: Generation of the output
+        ).to(device)
+
         generated_ids = model.generate(**inputs, max_new_tokens=128)
         generated_ids_trimmed = [
-            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
         ]
         output_text = processor.batch_decode(
             generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
         )
-        return output_text
+
+        return output_text[0]
+
 
 class audio_or_mp3__interpreter(Tool):
     name="audio_tool"