Update README.md

README.md

@@ -43,31 +43,31 @@ pip -r requirements.txt
 ```python
 from src.model import MMEBModel
 from src.arguments import ModelArguments
-from src.utils import load_processor
-
-import torch
-from transformers import HfArgumentParser, AutoProcessor
+from src.model_utils import load_processor, QWEN2_VL, vlm_image_tokens
 from PIL import Image
-import numpy as np
-
+import torch
 
-model_args = (
-    model_name='TIGER-Lab/VLM2Vec-Qwen2VL',
+model_args = ModelArguments(
+    model_name='Qwen/Qwen2-VL-7B-Instruct',
+    checkpoint_path='TIGER-Lab/VLM2Vec-Qwen2VL-7B',
     pooling='last',
     normalize=True,
-    model_backbone='qwen2_vl')
+    model_backbone='qwen2_vl',
+    lora=True
+)
 
 processor = load_processor(model_args)
-
 model = MMEBModel.load(model_args)
-model.eval()
 model = model.to('cuda', dtype=torch.bfloat16)
+model.eval()
 
 # Image + Text -> Text
-inputs = processor(text='<image> Represent the given image with the following question: What is in the image',
+inputs = processor(text=f'{vlm_image_tokens[QWEN2_VL]} Represent the given image with the following question: What is in the image',
                    images=Image.open('figures/example.jpg'),
                    return_tensors="pt")
 inputs = {key: value.to('cuda') for key, value in inputs.items()}
+inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)
+inputs['image_grid_thw'] = inputs['image_grid_thw'].unsqueeze(0)
 qry_output = model(qry=inputs)["qry_reps"]
 
 string = 'A cat and a dog'
@@ -77,7 +77,7 @@ inputs = processor(text=string,
 inputs = {key: value.to('cuda') for key, value in inputs.items()}
 tgt_output = model(tgt=inputs)["tgt_reps"]
 print(string, '=', model.compute_similarity(qry_output, tgt_output))
-## A cat and a dog = tensor([[0.4414]], device='cuda:0', dtype=torch.bfloat16)
+## A cat and a dog = tensor([[0.3301]], device='cuda:0', dtype=torch.bfloat16)
 
 string = 'A cat and a tiger'
 inputs = processor(text=string,
@@ -86,8 +86,7 @@ inputs = processor(text=string,
 inputs = {key: value.to('cuda') for key, value in inputs.items()}
 tgt_output = model(tgt=inputs)["tgt_reps"]
 print(string, '=', model.compute_similarity(qry_output, tgt_output))
-## A cat and a tiger = tensor([[0.3555]], device='cuda:0', dtype=torch.bfloat16)
-
+## A cat and a tiger = tensor([[0.2891]], device='cuda:0', dtype=torch.bfloat16)
 ```