Update README.md

README.md

@@ -39,9 +39,10 @@ def get_embedding(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor)
 ### Encode Text Query
 
 ```python
-queries = ["query: Where can we see Llama?", "query: What is LLaMA model?"]
+queries = ["query: Where can we see Llama?</s>", "query: What is LLaMA model?</s>"]
 query_inputs = processor(queries, return_tensors="pt", padding="longest", max_length=128, truncation=True).to('cuda:0')
-output = model(**query_inputs, return_dict=True, output_hidden_states=True)
+with torch.no_grad():
+    output = model(**query_inputs, return_dict=True, output_hidden_states=True)
 query_embeddings = get_embedding(output.hidden_states[-1], query_inputs["attention_mask"])
 ```
 
@@ -53,8 +54,8 @@ import requests
 from io import BytesIO
 
 # URLs of the images
-url1 = "https://huggingface.co/Tevatron/dse-phi3-docmatix-v1.0/blob/main/animal-llama.png"
-url2 = "https://huggingface.co/Tevatron/dse-phi3-docmatix-v1.0/blob/main/meta-llama.png"
+url1 = "https://huggingface.co/Tevatron/dse-phi3-docmatix-v1.0/resolve/main/animal-llama.png"
+url2 = "https://huggingface.co/Tevatron/dse-phi3-docmatix-v1.0/resolve/main/meta-llama.png"
 
 # Download and open images
 response1 = requests.get(url1)
@@ -64,21 +65,30 @@ passage_image1 = Image.open(BytesIO(response1.content))
 passage_image2 = Image.open(BytesIO(response2.content))
 
 passage_images = [passage_image1, passage_image2]
-passage_prompts = ["\nWhat is shown in this image?</s>", "\nWhat is shown in this image?</s>"]
+passage_prompts = ["<|image_1|>\nWhat is shown in this image?</s>", "<|image_2|>\nWhat is shown in this image?</s>"]
 
 # Process inputs and get embeddings
 passage_inputs = processor(passage_prompts, images=passage_images, return_tensors="pt", padding="longest", max_length=4096, truncation=True).to('cuda:0')
-output = model(**passage_inputs, return_dict=True, output_hidden_states=True)
+passage_inputs['input_ids'] = passage_inputs['input_ids'].squeeze(0)
+passage_inputs['attention_mask'] = passage_inputs['attention_mask'].squeeze(0)
+passage_inputs['image_sizes'] = passage_inputs['image_sizes'].squeeze(0)
+with torch.no_grad():
+    output = model(**passage_inputs, return_dict=True, output_hidden_states=True)
 doc_embeddings = get_embedding(output.hidden_states[-1], passage_inputs["attention_mask"])
+
 ```
 
 ### Compute Similarity
 
 ```python
 from torch.nn.functional import cosine_similarity
+num_queries = query_embeddings.size(0)
+num_passages = doc_embeddings.size(0)
 
-similarities = cosine_similarity(query_embeddings, doc_embeddings)
-print(similarities)
+for i in range(num_queries):
+    query_embedding = query_embeddings[i].unsqueeze(0)
+    similarities = cosine_similarity(query_embedding, doc_embeddings)
+    print(f"Similarities for Query {i+1}: {similarities.cpu().float().numpy()}")
 ```
 
 ### Encode Document Text
@@ -90,9 +100,12 @@ passage_prompts = [
 ]
 
 passage_inputs = processor(passage_prompts, images=None, return_tensors="pt", padding="longest", max_length=4096, truncation=True).to('cuda:0')
-output = model(**passage_inputs, return_dict=True, output_hidden_states=True)
+with torch.no_grad():
+    output = model(**passage_inputs, return_dict=True, output_hidden_states=True)
 doc_embeddings = get_embedding(output.hidden_states[-1], passage_inputs["attention_mask"])
 
-similarities = cosine_similarity(query_embeddings, doc_embeddings)
-print(similarities)
+for i in range(num_queries):
+    query_embedding = query_embeddings[i].unsqueeze(0)
+    similarities = cosine_similarity(query_embedding, doc_embeddings)
+    print(f"Similarities for Query {i+1}: {similarities.cpu().float().numpy()}")
 ```