Model update

README.md

@@ -275,92 +275,510 @@ Done ✅
 ### Using Transformers 🤗
 
 ```python
-from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
-from qwen_vl_utils import process_vision_info
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
 
-# default: Load the model on the available device(s)
-model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    "Qwen/Qwen2.5-VL-3B-Instruct", torch_dtype="auto", device_map="auto"
+# Load the processor
+# We dont currently support a fast image processor
+processor = AutoProcessor.from_pretrained(
+    'jinaai/jina-vlm-v1', use_fast=False, trust_remote_code=True
 )
 
-# We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
-# model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-#     "Qwen/Qwen2.5-VL-3B-Instruct",
+# Load the model on the available device(s)
+model = AutoModelForCausalLM.from_pretrained(
+    'jinaai/jina-vlm-v1',
+    device_map='auto', 
+    trust_remote_code=True
+)
+
+# You can specify a different model dtype and/or attention implementation
+# Available attention implementations: 
+#   'flash_attention_2', 'sdpa', 'eager'
+# Flash Attention 2 is recommended for improved inference speed and memory efficiency
+# For more details, see https://github.com/Dao-AILab/flash-attention
+# Flash Attention requires a CUDA device with compute capability >= 12.0
+# and dtype=torch.bfloat16 or torch.float16
+# SDPA and Eager are available on CPU and GPU, on all dtypes
+#
+# model = AutoModelForCausalLM.from_pretrained(
+#     'jinaai/jina-vlm-v1',
 #     torch_dtype=torch.bfloat16,
-#     attn_implementation="flash_attention_2",
-#     device_map="auto",
+#     attn_implementation='flash_attention_2',
+#     device_map='auto',
+#     trust_remote_code=True
 # )
 
-# default processer
-processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-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 range of 256-1280, to balance performance and cost.
-# min_pixels = 256*28*28
-# max_pixels = 1280*28*28
-# processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
-
-messages = [
+image = './assets/the_persistence_of_memory.jpg'
+conversation = [
     {
-        "role": "user",
-        "content": [
+        'role': 'user',
+        'content': [
             {
-                "type": "image",
-                "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
+                'type': 'image',
+                'image': image,
             },
-            {"type": "text", "text": "Describe this image."},
+            {'type': 'text', 'text': 'Describe this image.'},
         ],
     }
 ]
 
-# Preparation for inference
-text = processor.apply_chat_template(
-    messages, tokenize=False, add_generation_prompt=True
-)
-image_inputs, video_inputs = process_vision_info(messages)
+text = processor.apply_chat_template(conversation, add_generation_prompt=True)
 inputs = processor(
     text=[text],
-    images=image_inputs,
-    videos=video_inputs,
-    padding=True,
-    return_tensors="pt",
+    images=[image],
+    padding='longest',
+    return_tensors='pt',
 )
-inputs = inputs.to("cuda")
+# Configure max_pixels and max_crops when calling the processor
+# max_pixels if passed resizes all images that exceed the max number of pixels while 
+# preserving the aspect ratio. Less pixels == less visual tokens
+# max_crops specifies the max number of crops to generate for each image, also
+# reducing the number of visual tokens.
+# inputs = processor(
+#     text=[text],
+#     images=[image],
+#     padding='longest',
+#     max_length=1024,
+#     max_crops=8,
+#     max_pixels=100_000,
+#     do_resize=True,
+#     return_tensors='pt',
+# )
 
-# Inference: Generation of the output
-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)
-]
-output_text = processor.batch_decode(
-    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+# Move the inputs to the appropriate device and/or dtype
+device = torch.device('cuda')
+dtype = torch.float16
+model_inputs = {}
+for k, v in inputs.items():
+    if isinstance(v, torch.Tensor):
+        if v.is_floating_point():
+            model_inputs[k] = v.to(device, dtype=dtype, non_blocking=True)
+        else:
+            model_inputs[k] = v.to(device, non_blocking=True)
+    else:
+        model_inputs[k] = v
+
+# Inference
+output = model.generate(
+    **model_inputs,
+    generation_config=GenerationConfig(
+        max_new_tokens=20, do_sample=False,
+    ),
+    return_dict_in_generate=True,
+    use_model_defaults=True,
 )
-print(output_text)
+
+# Decode the output sequences and print the generated text
+# Input prompts will be skipped
+input_sequence_length = inputs.input_ids.shape[-1]
+for idx in range(len(output.sequences)):
+    gen_ids = output.sequences[idx][input_sequence_length:]
+    response = processor.tokenizer.decode(gen_ids, skip_special_tokens=True)
+    print(response)
 ```
 
 <details>
 <summary>Batch inference</summary>
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
+
+processor = AutoProcessor.from_pretrained(
+    'jinaai/jina-vlm-v1', use_fast=False, trust_remote_code=True
+)
+model = AutoModelForCausalLM.from_pretrained(
+    'jinaai/jina-vlm-v1',
+    device_map='auto', 
+    torch_dtype=torch.bfloat16,
+    attn_implementation='flash_attention_2',
+    trust_remote_code=True
+)
+images = [
+    'https://picsum.photos/id/22/4434/3729',
+    'https://picsum.photos/id/49/1280/792'
+]
+conversations = [
+    [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'image', 'image': images[0]},
+                {'type': 'text', 'text': 'What is the man doing in this image?'},
+            ],
+        }
+    ],
+    [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'image', 'image': images[1]},
+                {'type': 'text', 'text': 'What country\'s flag is in this image?'},
+            ],
+        }
+    ],
+    
+]
+texts = processor.apply_chat_template(conversations, add_generation_prompt=True)
+inputs = processor(
+    text=texts,
+    images=images,
+    padding='longest',
+    return_tensors='pt',
+)
+device = torch.device('cuda')
+dtype = torch.bfloat16
+model_inputs = {}
+for k, v in inputs.items():
+    if isinstance(v, torch.Tensor):
+        if v.is_floating_point():
+            model_inputs[k] = v.to(device, dtype=dtype, non_blocking=True)
+        else:
+            model_inputs[k] = v.to(device, non_blocking=True)
+    else:
+        model_inputs[k] = v
+
+output = model.generate(
+    **model_inputs,
+    generation_config=GenerationConfig(
+        max_new_tokens=20, do_sample=False,
+    ),
+    return_dict_in_generate=True,
+    use_model_defaults=True,
+)
+input_sequence_length = inputs.input_ids.shape[-1]
+for idx in range(len(output.sequences)):
+    gen_ids = output.sequences[idx][input_sequence_length:]
+    response = processor.tokenizer.decode(gen_ids, skip_special_tokens=True)
+    print(response)
+```
+
 </details>
 
 <details>
 <summary>Multi-image inference</summary>
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
+
+processor = AutoProcessor.from_pretrained(
+    'jinaai/jina-vlm-v1', use_fast=False, trust_remote_code=True
+)
+model = AutoModelForCausalLM.from_pretrained(
+    'jinaai/jina-vlm-v1',
+    device_map='auto', 
+    torch_dtype=torch.bfloat16,
+    attn_implementation='flash_attention_2',
+    trust_remote_code=True
+)
+images = [
+    'https://picsum.photos/id/0/5000/3333',
+    'https://picsum.photos/id/2/5000/3333'
+]
+conversation = [
+    {
+        'role': 'user',
+        'content': [
+            {'type': 'image', 'image': images[0]},
+            {'type': 'image', 'image': images[1]},
+            {'type': 'text', 'text': 'What is the difference between these two images?'},
+        ],
+    }
+]
+text = processor.apply_chat_template(conversation, add_generation_prompt=True)
+inputs = processor(
+    text=[text],
+    images=images,
+    padding='longest',
+    return_tensors='pt',
+)
+device = torch.device('cuda')
+dtype = torch.bfloat16
+model_inputs = {}
+for k, v in inputs.items():
+    if isinstance(v, torch.Tensor):
+        if v.is_floating_point():
+            model_inputs[k] = v.to(device, dtype=dtype, non_blocking=True)
+        else:
+            model_inputs[k] = v.to(device, non_blocking=True)
+    else:
+        model_inputs[k] = v
+
+output = model.generate(
+    **model_inputs,
+    generation_config=GenerationConfig(
+        max_new_tokens=20, do_sample=False,
+    ),
+    return_dict_in_generate=True,
+    use_model_defaults=True,
+)
+input_sequence_length = inputs.input_ids.shape[-1]
+for idx in range(len(output.sequences)):
+    gen_ids = output.sequences[idx][input_sequence_length:]
+    response = processor.tokenizer.decode(gen_ids, skip_special_tokens=True)
+    print(response)
+```
+
 </details>
 
 <details>
 <summary>Text-only inference</summary>
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
+
+processor = AutoProcessor.from_pretrained(
+    'jinaai/jina-vlm-v1', use_fast=False, trust_remote_code=True
+)
+model = AutoModelForCausalLM.from_pretrained(
+    'jinaai/jina-vlm-v1',
+    device_map='auto', 
+    torch_dtype=torch.bfloat16,
+    attn_implementation='flash_attention_2',
+    trust_remote_code=True
+)
+conversation = [
+    {
+        'role': 'user',
+        'content': [
+            {
+                'type': 'text',
+                'text': 'Describe the concept of polymorphism in Computer Science'
+            },
+        ],
+    }
+]
+text = processor.apply_chat_template(conversation, add_generation_prompt=True)
+inputs = processor(
+    text=[text],
+    images=None,
+    padding='longest',
+    return_tensors='pt',
+)
+device = torch.device('cuda')
+dtype = torch.bfloat16
+model_inputs = {}
+for k, v in inputs.items():
+    if isinstance(v, torch.Tensor):
+        if v.is_floating_point():
+            model_inputs[k] = v.to(device, dtype=dtype, non_blocking=True)
+        else:
+            model_inputs[k] = v.to(device, non_blocking=True)
+    else:
+        model_inputs[k] = v
+
+output = model.generate(
+    **model_inputs,
+    generation_config=GenerationConfig(
+        max_new_tokens=20, do_sample=False,
+    ),
+    return_dict_in_generate=True,
+    use_model_defaults=True,
+)
+input_sequence_length = inputs.input_ids.shape[-1]
+for idx in range(len(output.sequences)):
+    gen_ids = output.sequences[idx][input_sequence_length:]
+    response = processor.tokenizer.decode(gen_ids, skip_special_tokens=True)
+    print(response)
+```
+
 </details>
 
 <details>
-<summary>Mixed-batch inference</summary>
+<summary>Batch inference with mixed examples</summary>
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
+
+processor = AutoProcessor.from_pretrained(
+    'jinaai/jina-vlm-v1', use_fast=False, trust_remote_code=True
+)
+model = AutoModelForCausalLM.from_pretrained(
+    'jinaai/jina-vlm-v1',
+    device_map='auto', 
+    torch_dtype=torch.bfloat16,
+    attn_implementation='flash_attention_2',
+    trust_remote_code=True
+)
+images = [
+    ['https://picsum.photos/id/22/4434/3729'],
+    ['https://picsum.photos/id/49/1280/792'],
+    [
+        'https://picsum.photos/id/0/5000/3333',
+        'https://picsum.photos/id/2/5000/3333',
+    ]
+]
+conversations = [
+    [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'image', 'image': images[0][0]},
+                {'type': 'text', 'text': 'What is the man doing in this image?'},
+            ],
+        }
+    ],
+    [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'image', 'image': images[1][0]},
+                {'type': 'text', 'text': 'What country\'s flag is in this image?'},
+            ],
+        }
+    ],
+    [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'image', 'image': images[2][0]},
+                {'type': 'image', 'image': images[2][1]},
+                {'type': 'text', 'text': 'What is the difference between these two images?'},
+            ],
+        }
+    ],
+    [
+        {
+            'role': 'user',
+            'content': [
+                {
+                    'type': 'text',
+                    'text': 'Describe the concept of polymorphism in Computer Science'
+                },
+            ],
+        }
+    ],
+]
+texts = processor.apply_chat_template(conversations, add_generation_prompt=True)
+inputs = processor(
+    text=texts,
+    images=images,
+    padding='longest',
+    return_tensors='pt',
+)
+device = torch.device('cuda')
+dtype = torch.bfloat16
+model_inputs = {}
+for k, v in inputs.items():
+    if isinstance(v, torch.Tensor):
+        if v.is_floating_point():
+            model_inputs[k] = v.to(device, dtype=dtype, non_blocking=True)
+        else:
+            model_inputs[k] = v.to(device, non_blocking=True)
+    else:
+        model_inputs[k] = v
+
+output = model.generate(
+    **model_inputs,
+    generation_config=GenerationConfig(
+        max_new_tokens=20, do_sample=False,
+    ),
+    return_dict_in_generate=True,
+    use_model_defaults=True,
+)
+input_sequence_length = inputs.input_ids.shape[-1]
+for idx in range(len(output.sequences)):
+    gen_ids = output.sequences[idx][input_sequence_length:]
+    response = processor.tokenizer.decode(gen_ids, skip_special_tokens=True)
+    print(response)
+```
+
 </details>
 
 <details>
 <summary>Feature extraction</summary>
-</details>
 
-### Using vLLM
+```python
+import torch
+from transformers import AutoModel, AutoProcessor
 
-Coming soon!
+processor = AutoProcessor.from_pretrained(
+    'jinaai/jina-vlm-v1', use_fast=False, trust_remote_code=True
+)
+model = AutoModel.from_pretrained(
+    'jinaai/jina-vlm-v1',
+    device_map='auto', 
+    torch_dtype=torch.bfloat16,
+    attn_implementation='flash_attention_2',
+    trust_remote_code=True
+)
+images = [
+    ['https://picsum.photos/id/22/4434/3729'],
+    ['https://picsum.photos/id/49/1280/792'],
+    [
+        'https://picsum.photos/id/0/5000/3333',
+        'https://picsum.photos/id/2/5000/3333',
+    ]
+]
+conversations = [
+    [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'image', 'image': images[0][0]},
+                {'type': 'text', 'text': 'What is the man doing in this image?'},
+            ],
+        }
+    ],
+    [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'image', 'image': images[1][0]},
+                {'type': 'text', 'text': 'What country\'s flag is in this image?'},
+            ],
+        }
+    ],
+    [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'image', 'image': images[2][0]},
+                {'type': 'image', 'image': images[2][1]},
+                {'type': 'text', 'text': 'What is the difference between these two images?'},
+            ],
+        }
+    ],
+    [
+        {
+            'role': 'user',
+            'content': [
+                {
+                    'type': 'text',
+                    'text': 'Describe the concept of polymorphism in Computer Science'
+                },
+            ],
+        }
+    ],
+]
+texts = processor.apply_chat_template(conversations, add_generation_prompt=True)
+inputs = processor(
+    text=texts,
+    images=images,
+    padding='longest',
+    return_tensors='pt',
+)
+device = torch.device('cuda')
+dtype = torch.bfloat16
+model_inputs = {}
+for k, v in inputs.items():
+    if isinstance(v, torch.Tensor):
+        if v.is_floating_point():
+            model_inputs[k] = v.to(device, dtype=dtype, non_blocking=True)
+        else:
+            model_inputs[k] = v.to(device, non_blocking=True)
+    else:
+        model_inputs[k] = v
+
+output = model(**model_inputs)
+hidden_states = output.hidden_states
+last_hidden_states = output.last_hidden_state
+```
+
+</details>
 
 
 ## License

processing_jvlm.py

@@ -39,7 +39,7 @@ class JinaVLMTextKwargs(TypedDict, total=False):
 
 
 class JinaVLProcessingKwargs(JinaVLMTextKwargs, JinaVLMImagesKwargs, CommonKwargs):
-    pass
+    return_labels: Optional[bool]
 
 
 class JinaVLMProcessor(ProcessorMixin):
@@ -259,6 +259,7 @@ class JinaVLMProcessor(ProcessorMixin):
         image_tokens: List[np.ndarray],
         image_input_idx: List[np.ndarray],
         image_padding_mask: List[np.ndarray],
+        return_labels: bool = False,
         add_empty_image_features: bool = False,
     ):
         """Interleave images and text tokens into multi-modal features for the model."""
@@ -282,8 +283,9 @@ class JinaVLMProcessor(ProcessorMixin):
             data = {
                 'input_ids': input_ids,
                 'position_ids': position_ids,
-                'labels': target_tokens,
             }
+            if return_labels:
+                data['labels'] = target_tokens
             if add_empty_image_features:
                 # Add size-zero image features, this can be useful to make sure all
                 # devices get an image input when the image ViT is FSDP wrapped
@@ -367,14 +369,16 @@ class JinaVLMProcessor(ProcessorMixin):
             image_input_idx < 0, image_input_idx, image_input_idx + 1
         )
         position_ids = np.arange(len(input_ids), dtype=np.int64)
-        return {
+        data = {
             'input_ids': input_ids,
             'position_ids': position_ids,
             'images': images,
             'image_input_idx': image_input_idx,
             'image_masks': image_masks,
-            'labels': target_tokens,
         }
+        if return_labels:
+            data['labels'] = target_tokens
+        return data
 
     def __call__(
         self,
@@ -425,6 +429,7 @@ class JinaVLMProcessor(ProcessorMixin):
             raise ValueError('Processor requires text input.')
 
         return_tensors = kwargs.pop('return_tensors', None)
+        return_labels = kwargs.pop('return_labels', False)
         padding = kwargs.pop('padding', PaddingStrategy.LONGEST)
         padding_side = kwargs.pop('padding_side', 'left')
         max_length = kwargs.pop('max_length', None)
@@ -498,6 +503,7 @@ class JinaVLMProcessor(ProcessorMixin):
                 image_input_idx,
                 image_padding_mask if image_padding_mask is not None else [],
                 add_empty_image_features=(batch_size > 1),
+                return_labels=return_labels,
             )
             for k, v in output.items():
                 outputs[k].append(v)