TQThinh thêm importance tags

ram/inference.py

@@ -32,10 +32,14 @@ def inference_tag2text(image, model, input_tag="None"):
 
 def inference_ram(image, model):
 
-    with torch.no_grad():
-        tags, tags_chinese = model.generate_tag(image)
+    # with torch.no_grad():
+    #     tags, tags_chinese = model.generate_tag(image)
 
-    return tags[0],tags_chinese[0]
+    # return tags[0],tags_chinese[0]
+    with torch.no_grad():
+        tags, chinese_tags, tag_importance = model.generate_tag(image)
+
+    return tags, chinese_tags, tag_importance
 
 
 def inference_ram_openset(image, model):

ram/models/bert.py

@@ -36,11 +36,11 @@
     SequenceClassifierOutput,
     TokenClassifierOutput,
 )
+from transformers.pytorch_utils import apply_chunking_to_forward
+from .helper import find_pruneable_heads_and_indices
+from .helper2 import _prune_linear_layer
 from transformers.modeling_utils import (
     PreTrainedModel,
-    apply_chunking_to_forward,
-    find_pruneable_heads_and_indices,
-    prune_linear_layer,
 )
 from transformers.utils import logging
 from transformers.models.bert.configuration_bert import BertConfig
@@ -310,10 +310,10 @@ def prune_heads(self, heads):
         )
 
         # Prune linear layers
-        self.self.query = prune_linear_layer(self.self.query, index)
-        self.self.key = prune_linear_layer(self.self.key, index)
-        self.self.value = prune_linear_layer(self.self.value, index)
-        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+        self.self.query = _prune_linear_layer(self.self.query, index)
+        self.self.key = _prune_linear_layer(self.self.key, index)
+        self.self.value = _prune_linear_layer(self.self.value, index)
+        self.output.dense = _prune_linear_layer(self.output.dense, index, dim=1)
 
         # Update hyper params and store pruned heads
         self.self.num_attention_heads = self.self.num_attention_heads - len(heads)

ram/models/helper.py

@@ -0,0 +1,40 @@
+import torch
+
+def find_pruneable_heads_and_indices(
+    heads: list[int],
+    n_heads: int,
+    head_size: int,
+    already_pruned_heads: set[int],
+) -> tuple[set[int], torch.LongTensor]:
+    """
+    Finds the heads and the flattened indices to keep, taking already-pruned heads
+    into account.
+
+    Parameters
+    ----------
+    heads : list[int]
+        Head indices requested for pruning.
+    n_heads : int
+        Total number of attention heads before this pruning step.
+    head_size : int
+        Size of each attention head.
+    already_pruned_heads : set[int]
+        Heads that were pruned in earlier steps.
+
+    Returns
+    -------
+    tuple[set[int], torch.LongTensor]
+        (new_heads_to_prune, flattened_indices_to_keep)
+    """
+    mask = torch.ones(n_heads, head_size, dtype=torch.bool)
+
+    heads = set(heads) - already_pruned_heads
+
+    for head in heads:
+        # Shift the head index left by however many smaller heads
+        # were already removed earlier.
+        shifted_head = head - sum(1 for h in already_pruned_heads if h < head)
+        mask[shifted_head] = False
+
+    index = torch.arange(n_heads * head_size)[mask.view(-1)].long()
+    return heads, index

ram/models/helper2.py

@@ -0,0 +1,92 @@
+import logging
+import torch
+import torch.nn as nn
+from transformers import modeling_utils
+
+# ------------------------------------------------------------------------------
+# 1. Define Fallback Implementations
+#    (Mirrors logic removed in transformers v4.45+)
+# ------------------------------------------------------------------------------
+
+def _prune_linear_layer(layer, index, dim=0):
+    """Fallback implementation for prune_linear_layer"""
+    index = index.to(layer.weight.device)
+    W = layer.weight.index_select(dim, index).clone().detach()
+    if layer.bias is not None:
+        if dim == 0:
+            b = layer.bias.index_select(dim, index).clone().detach()
+        else:
+            b = layer.bias.clone().detach()
+    new_size = list(layer.weight.size())
+    new_size[dim] = len(index)
+    new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device)
+    new_layer.weight.requires_grad = False
+    new_layer.weight.copy_(W.contiguous())
+    new_layer.weight.requires_grad = True
+    if layer.bias is not None:
+        new_layer.bias.requires_grad = False
+        new_layer.bias.copy_(b.contiguous())
+        new_layer.bias.requires_grad = True
+    return new_layer
+
+def _prune_layer(layer, index, dim=None):
+    """Fallback for prune_layer (generic)"""
+    if isinstance(layer, nn.Linear):
+        return _prune_linear_layer(layer, index, dim=0 if dim is None else dim)
+
+    # Handle Conv1D if available
+    try:
+        from transformers.pytorch_utils import prune_conv1d_layer
+        if hasattr(modeling_utils, "Conv1D") and isinstance(layer, modeling_utils.Conv1D):
+            return prune_conv1d_layer(layer, index, dim=1 if dim is None else dim)
+    except ImportError:
+        pass
+
+    return _prune_linear_layer(layer, index, dim=0 if dim is None else dim)
+
+def _find_pruneable_heads_and_indices(heads, n_heads, head_size, already_pruned_heads):
+    """Fallback for finding heads to prune"""
+    mask = torch.ones(n_heads, head_size)
+    heads = set(heads) - already_pruned_heads
+    for head in heads:
+        head = head - sum(1 if h < head else 0 for h in already_pruned_heads)
+        mask[head] = 0
+    mask = mask.view(-1).contiguous().eq(1)
+    index = torch.arange(len(mask))[mask].long()
+    return heads, index
+
+# ------------------------------------------------------------------------------
+# 2. Inject Missing Functions into transformers.modeling_utils
+# ------------------------------------------------------------------------------
+
+# Patch 'Conv1D' if missing
+if not hasattr(modeling_utils, "Conv1D"):
+    try:
+        from transformers.pytorch_utils import Conv1D
+        modeling_utils.Conv1D = Conv1D
+    except ImportError:
+        pass
+
+# Patch 'prune_linear_layer'
+if not hasattr(modeling_utils, "prune_linear_layer"):
+    try:
+        from transformers.pytorch_utils import prune_linear_layer
+        modeling_utils.prune_linear_layer = prune_linear_layer
+    except ImportError:
+        modeling_utils.prune_linear_layer = _prune_linear_layer
+
+# Patch 'prune_layer' (Critical for MeshGraphormer)
+if not hasattr(modeling_utils, "prune_layer"):
+    modeling_utils.prune_layer = _prune_layer
+
+# Patch 'find_pruneable_heads_and_indices'
+if not hasattr(modeling_utils, "find_pruneable_heads_and_indices"):
+    try:
+        from transformers.pytorch_utils import find_pruneable_heads_and_indices
+        modeling_utils.find_pruneable_heads_and_indices = find_pruneable_heads_and_indices
+    except ImportError:
+        modeling_utils.find_pruneable_heads_and_indices = _find_pruneable_heads_and_indices
+
+logging.info("\033[32m[Transformers Fix] Successfully injected missing pruning functions for MeshGraphormer compatibility.\033[0m")
+
+NODE_CLASS_MAPPINGS = {}

ram/models/ram.py

@@ -303,19 +303,59 @@ def forward(self, image, caption, image_tag, parse_tag, clip_feature):
 
         return loss_t2t, loss_tag, loss_dis
 
-    def generate_tag(self,
-                 image,
-                 threshold=0.68,
-                 tag_input=None,
-                 ):
+    # def generate_tag(self,
+    #              image,
+    #              threshold=0.68,
+    #              tag_input=None,
+    #              ):
 
+    #     label_embed = torch.nn.functional.relu(self.wordvec_proj(self.label_embed))
+
+    #     image_embeds = self.image_proj(self.visual_encoder(image))
+    #     image_atts = torch.ones(image_embeds.size()[:-1],
+    #                             dtype=torch.long).to(image.device)
+
+    #     # recognized image tags using image-tag recogntiion decoder
+    #     image_cls_embeds = image_embeds[:, 0, :]
+    #     image_spatial_embeds = image_embeds[:, 1:, :]
+
+    #     bs = image_spatial_embeds.shape[0]
+    #     label_embed = label_embed.unsqueeze(0).repeat(bs, 1, 1)
+    #     tagging_embed = self.tagging_head(
+    #         encoder_embeds=label_embed,
+    #         encoder_hidden_states=image_embeds,
+    #         encoder_attention_mask=image_atts,
+    #         return_dict=False,
+    #         mode='tagging',
+    #     )
+
+    #     logits = self.fc(tagging_embed[0]).squeeze(-1)
+
+    #     targets = torch.where(
+    #         torch.sigmoid(logits) > self.class_threshold.to(image.device),
+    #         torch.tensor(1.0).to(image.device),
+    #         torch.zeros(self.num_class).to(image.device))
+
+    #     tag = targets.cpu().numpy()
+    #     tag[:,self.delete_tag_index] = 0
+    #     tag_output = []
+    #     tag_output_chinese = []
+    #     for b in range(bs):
+    #         index = np.argwhere(tag[b] == 1)
+    #         token = self.tag_list[index].squeeze(axis=1)
+    #         tag_output.append(' | '.join(token))
+    #         token_chinese = self.tag_list_chinese[index].squeeze(axis=1)
+    #         tag_output_chinese.append(' | '.join(token_chinese))
+
+
+    #     return tag_output, tag_output_chinese
+    def generate_tag(self, image, threshold=0.68, tag_input=None):
         label_embed = torch.nn.functional.relu(self.wordvec_proj(self.label_embed))
 
         image_embeds = self.image_proj(self.visual_encoder(image))
-        image_atts = torch.ones(image_embeds.size()[:-1],
-                                dtype=torch.long).to(image.device)
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(image.device)
 
-        # recognized image tags using image-tag recogntiion decoder
+        # Recognized image tags using image-tag recognition decoder
         image_cls_embeds = image_embeds[:, 0, :]
         image_spatial_embeds = image_embeds[:, 1:, :]
 
@@ -330,25 +370,31 @@ def generate_tag(self,
         )
 
         logits = self.fc(tagging_embed[0]).squeeze(-1)
+        probabilities = torch.sigmoid(logits)  # Tính xác suất bằng sigmoid
 
         targets = torch.where(
-            torch.sigmoid(logits) > self.class_threshold.to(image.device),
+            probabilities > self.class_threshold.to(image.device),
             torch.tensor(1.0).to(image.device),
-            torch.zeros(self.num_class).to(image.device))
+            torch.zeros(self.num_class).to(image.device)
+        )
 
         tag = targets.cpu().numpy()
-        tag[:,self.delete_tag_index] = 0
+        tag[:, self.delete_tag_index] = 0
         tag_output = []
         tag_output_chinese = []
+        tag_importance = []  # Lưu độ quan trọng của các tag
+
         for b in range(bs):
             index = np.argwhere(tag[b] == 1)
             token = self.tag_list[index].squeeze(axis=1)
             tag_output.append(' | '.join(token))
             token_chinese = self.tag_list_chinese[index].squeeze(axis=1)
             tag_output_chinese.append(' | '.join(token_chinese))
+            # Lưu xác suất của các tag được chọn
+            importance = probabilities[b][index].squeeze(axis=1).cpu().numpy()
+            tag_importance.append(dict(zip(token, importance)))
 
-
-        return tag_output, tag_output_chinese
+        return tag_output, tag_output_chinese, tag_importance    
 
     def generate_tag_openset(self,
                  image,

ram/models/utils.py

@@ -131,7 +131,7 @@ def init_tokenizer(text_encoder_type='bert-base-uncased'):
     tokenizer = BertTokenizer.from_pretrained(text_encoder_type)
     tokenizer.add_special_tokens({'bos_token': '[DEC]'})
     tokenizer.add_special_tokens({'additional_special_tokens': ['[ENC]']})
-    tokenizer.enc_token_id = tokenizer.additional_special_tokens_ids[0]
+    tokenizer.enc_token_id = tokenizer.convert_tokens_to_ids('[ENC]')
     return tokenizer
 
 

requirements.txt

@@ -1,5 +1,5 @@
 timm==0.4.12
-transformers>=4.25.1
+transformers==4.57.1
 fairscale==0.4.4
 pycocoevalcap
 torch