Update README.md

README.md

@@ -58,16 +58,6 @@ import torch
 
 device = 'cuda:0'
 
-# This function is borrowed from https://huggingface.co/intfloat/e5-mistral-7b-instruct
-def last_token_pool(last_hidden_states, attention_mask):
-    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
-    if left_padding:
-        return last_hidden_states[:, -1]
-    else:
-        sequence_lengths = attention_mask.sum(dim=1) - 1
-        batch_size = last_hidden_states.shape[0]
-        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
-
 # Load model, be sure to substitute `model_path` by your model path 
 model_path = '/local/path/to/model'
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
@@ -86,19 +76,15 @@ query_full = query_instruction + query
 
 # Embed image documents
 with torch.no_grad():
-    p_outputs = model(text=['', '', ''], image=[image_1, image_2, image_3], tokenizer=tokenizer)
-    p_reps = last_token_pool(p_outputs.last_hidden_state, p_outputs.attention_mask)
+    p_reps = model(text=['', '', ''], image=[image_1, image_2, image_3], tokenizer=tokenizer)
 
 # Embed text queries
 with torch.no_grad():
-    q_outputs = model(text=[query_full], image=[None], tokenizer=tokenizer) # [B, s, d]
-    q_reps = last_token_pool(q_outputs.last_hidden_state, q_outputs.attention_mask) # [B, d]
+    q_reps = model(text=[query_full], image=[None], tokenizer=tokenizer) # [B, s, d]
 
 # Calculate similarities
 scores = torch.matmul(q_reps, p_reps.T)
 print(scores)
-
-# tensor([[0.6506, 4.9630, 3.8614]], device='cuda:0')
 ```
 
 # Limitations