Update README.md

README.md

@@ -30,7 +30,7 @@ The proposed solutions offer a more nuanced understanding of memes for accurate
 
 
 
-# Finetuned Gemma 3 Embedding Model
+# Finetuned Gemma 3 Embedding Model with mean pooling
 
 ``` python
 
@@ -39,13 +39,17 @@ import torch
 import torch._dynamo
 from tqdm import tqdm  # Progress bar library
 from unsloth import FastVisionModel
+
 from datasets import load_dataset
 
-def convert_to_conversation(sample):    
+instruction = "classify meme into Hateful or Not"
+
+def convert_to_conversation(sample):
+    
     lis=[]
     lis.append({"type": "text", "text": sample["text"]})
     lis.append({"type": "image", "image": sample["image"]})
-
+                    
     conversation = [
         {
             "role": "system",
@@ -61,11 +65,12 @@ def convert_to_conversation(sample):
 pass
 
 dataset = load_dataset("QCRI/ArGuard-Task1",  split="dev_test")
+
 converted_dataset = [convert_to_conversation(sample) for sample in dataset]
 
 # 1. Disable compiler optimization conflicts
-torch._dynamo.config.disable = True
-torch._dynamo.reset()
+#torch._dynamo.config.disable = True
+#torch._dynamo.reset()
 
 # 2. Load your fine-tuned model and processor
 model_path = "NYUAD-ComNets/Gemma3_meme_classification"
@@ -81,7 +86,7 @@ instruction = "classify meme into Hateful or Not"
 all_embeddings = []
 labels_list = []  
 
-num_iterations = min(500, len(converted_dataset))
+num_iterations = len(converted_dataset)
 
 print(f"Starting embedding extraction for {num_iterations} items...")
 
@@ -104,24 +109,26 @@ for idx in tqdm(range(num_iterations), desc="Extracting Embeddings"):
         templated_text = processor.apply_chat_template(conversation, tokenize=False)
         inputs = processor(text=templated_text, images=sample_image, return_tensors="pt").to("cuda")
 
-        # 5. Safe Extraction via .generate()
+
+        # 5. Forward Pass
         with torch.no_grad():
             with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
-                outputs = model.generate(
-                    **inputs, 
-                    max_new_tokens=1,              
-                    output_hidden_states=True,     
-                    return_dict_in_generate=True   
-                )
-
-        # 6. Extract final hidden layer states
-        last_hidden_states = outputs.hidden_states[0][-1]
-
-        # 7. Perform Mean Pooling and send directly to CPU memory
-        multimodal_embedding = last_hidden_states.mean(dim=1).squeeze(0).float()
+                outputs = model(**inputs, output_hidden_states=True, return_dict=True)
+
+        # 6. Extract final layer and attention mask
+        last_hidden_states = outputs.hidden_states[-1]  # [batch_size, seq_len, hidden_dim]
+        attention_mask = inputs["attention_mask"]       # [batch_size, seq_len]
+
+        # 7. Masked Mean Pooling (Ignores padding tokens entirely)
+        input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_states.size()).float()
+        sum_embeddings = torch.sum(last_hidden_states * input_mask_expanded, dim=1)
+        sum_mask = torch.clamp(input_mask_expanded.sum(dim=1), min=1e-9)
+        
+        all_embedding = (sum_embeddings / sum_mask).squeeze(0).float()
         
-        # Convert torch tensor to a clean numpy row vector and append
-        all_embeddings.append(multimodal_embedding.cpu().numpy())
+        all_embeddings.append(all_embedding.cpu().numpy())
+
+
         labels_list.append(sample_label)
 
     except Exception as e:
@@ -131,7 +138,8 @@ for idx in tqdm(range(num_iterations), desc="Extracting Embeddings"):
 embedding_matrix = np.vstack(all_embeddings)
 
 print("Final Concatenated Array Shape:", embedding_matrix.shape)
-np.save("test_embeddings.npy", embedding_matrix)
+
+np.save("test_gemma3_mean_embeddings.npy", embedding_matrix)
 ```