Upload 2 files

detree/model/simclr.py

@@ -0,0 +1,284 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .text_embedding import TextEmbeddingModel
+
+class Tree():
+    def __init__(self,path):
+        self.name = {}
+        self.childs = {}
+        self.father = {}
+        self.dep = {}
+        self.root = None
+        self.max_dep = 0
+        self.subtree = {}
+        self.grad_fa = {} # the node closest to the root for each leaf
+        with open(path, 'r') as f:
+            lines = f.readlines()
+            for line in lines:
+                parts = line.strip().split()
+                assert len(parts) == 3, "Each line must have exactly three parts"
+                
+                now,fa,name = parts
+                now,fa = int(now),int(fa)
+                if name != 'none':
+                    self.name[now] = name.split(',')
+                if fa != -1:
+                    self.childs[fa] = self.childs.get(fa, []) + [now]
+                    self.father[now] = fa
+                else:
+                    self.root = now
+        self.fa_pos = torch.zeros((len(self.father),len(self.father)),dtype=torch.bool)
+
+        self.dfs(self.root)
+        #max_dep,N,N+K 0/1
+        self.pos_down2up = torch.zeros((self.max_dep,len(self.name),len(self.father)),dtype=torch.bool)
+        self.neg_down2up = torch.zeros((self.max_dep,len(self.name),len(self.father)),dtype=torch.bool)
+
+        self.pos_up2down = torch.zeros((self.max_dep,len(self.name),len(self.father)),dtype=torch.bool)
+        self.neg_up2down = torch.zeros((self.max_dep,len(self.name),len(self.father)),dtype=torch.bool)
+
+        self.pos_center = torch.zeros((self.max_dep,len(self.name)),dtype=torch.long)
+        self.mask_center = torch.zeros((self.max_dep,len(self.name),len(self.father)),dtype=torch.bool)
+
+        #max_dep,N 0/1
+        self.mask = torch.zeros((self.max_dep,len(self.name)),dtype=torch.bool)
+        self.depth = torch.zeros(len(self.name))
+        self.labels = torch.zeros(len(self.name),dtype=torch.long)
+        self.vis_leaf()
+        label_value = list(set(self.grad_fa.values()))
+        for key, value in self.grad_fa.items():
+            self.labels[key] = label_value.index(value)
+
+    def dfs(self, node, depth=0,grfa=-1):
+        self.dep[node] = depth
+        self.max_dep = max(self.max_dep, depth)
+        if node!=self.root:
+            self.subtree[node] = torch.zeros(len(self.father),dtype=torch.bool)
+            self.subtree[node][node] = 1
+
+        # if self.fa_pos.get(node) is None:
+            if self.father[node] != self.root:
+                self.fa_pos[node] = self.fa_pos[self.father[node]].clone()
+            self.fa_pos[node][node] = 1
+            if grfa == -1:
+                grfa = node
+        if self.childs.get(node) is None:
+            self.grad_fa[node] = grfa
+        for child in self.childs.get(node, []):
+            self.dfs(child, depth + 1,grfa)
+            if node!=self.root:
+                self.subtree[node] = torch.logical_or(self.subtree[node], self.subtree[child])
+
+    def gen_leaf_item(self,node):
+        last_node = -1
+        leaf_id = node
+        self.depth[node] = self.dep[node]
+        while node != self.root:
+            now_dep=self.dep[node]-1
+            self.mask[now_dep,leaf_id] = 1
+            self.pos_center[now_dep,leaf_id] = node
+            self.mask_center[now_dep,leaf_id] = torch.logical_not(torch.logical_or(self.fa_pos[node],self.subtree[node]))
+            self.mask_center[now_dep,leaf_id,node] = 1
+            if last_node == -1:
+                self.pos_down2up[now_dep,leaf_id] = self.subtree[node]
+            else:
+                self.pos_down2up[now_dep,leaf_id]=torch.logical_xor(self.subtree[node],self.subtree[last_node])
+            self.neg_down2up[now_dep,leaf_id]=torch.logical_not(self.subtree[node])
+
+            if self.father[node] == self.root:
+                self.neg_up2down[now_dep,leaf_id] = torch.logical_not(self.subtree[node])
+            else:
+                self.neg_up2down[now_dep,leaf_id] = torch.logical_xor(self.subtree[node],self.subtree[self.father[node]])
+            self.pos_up2down[now_dep,leaf_id] = self.subtree[node]
+
+            last_node = node
+            node = self.father[node]
+
+    def vis_leaf(self):
+        for node, name in self.name.items():
+            self.gen_leaf_item(node)
+
+
+    def display(self):
+        for node, name in self.name.items():
+            depth = self.dep[node]
+            print(f"{depth}- {name} {self.father[node]}")
+                
+class SimCLR_Tree(nn.Module):
+    def __init__(self, opt, fabric):
+        super(SimCLR_Tree, self).__init__()
+
+        self.temperature = opt.temperature
+        self.opt = opt
+        self.fabric = fabric
+
+        adapter_path = getattr(opt, "adapter_path", None)
+        self.model = TextEmbeddingModel(
+            opt.model_name,
+            lora=opt.lora,
+            use_pooling=opt.pooling,
+            lora_r=opt.lora_r,
+            lora_alpha=opt.lora_alpha,
+            lora_dropout=opt.lora_dropout,
+            adapter_path=adapter_path,
+        )
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.tree = Tree(opt.tree_txt)
+
+        self.pos_down2up = self.tree.pos_down2up.to(self.device)
+        self.neg_down2up = self.tree.neg_down2up.to(self.device)
+        self.pos_up2down = self.tree.pos_up2down.to(self.device)
+        self.neg_up2down = self.tree.neg_up2down.to(self.device)
+        self.pos_center = self.tree.pos_center.to(self.device)
+        self.mask_center = self.tree.mask_center.to(self.device)
+
+        self.K = self.pos_down2up.shape[0]
+
+        self.mask = self.tree.mask.to(self.device)
+        self.depth = self.tree.depth.to(self.device)
+        self.root_labels = self.tree.labels.to(self.device)
+        self.esp = torch.tensor(1e-6, device=self.device)
+        self.max_dep = self.tree.max_dep
+        self.leaf_cnt = len(self.tree.name)
+
+        self.names2id = {}
+        for key, value in self.tree.name.items():
+            for item in value:
+                self.names2id[item] = key
+
+        self.vitual_center = nn.Parameter(
+            torch.randn((len(self.tree.father), opt.projection_size), device=self.device),
+            requires_grad=True,
+        )
+        nn.init.xavier_uniform_(self.vitual_center)
+        self.center_labels = torch.arange(len(self.tree.father), dtype=torch.long, device=self.device)
+        if adapter_path is not None:
+            self.load_tree_state(adapter_path)
+
+
+    def get_encoder(self):
+        return self.model
+
+    def save_pretrained(self, save_directory: str, save_tokenizer: bool = True):
+        os.makedirs(save_directory, exist_ok=True)
+        self.model.save_pretrained(save_directory, save_tokenizer=save_tokenizer)
+        torch.save(
+            {"vitual_center": self.vitual_center.detach().cpu()},
+            os.path.join(save_directory, "tree_state.pt"),
+        )
+
+    def load_tree_state(self, directory: str):
+        state_path = os.path.join(directory, "tree_state.pt")
+        if not os.path.exists(state_path):
+            return
+        state = torch.load(state_path, map_location=self.vitual_center.device)
+        self.vitual_center.data.copy_(state["vitual_center"].to(self.vitual_center.device))
+
+    def load_from_directory(self, directory: str, is_trainable: bool = True):
+        if getattr(self.opt, "lora", False):
+            self.model.load_adapter(directory, is_trainable=is_trainable)
+        else:
+            self.model = TextEmbeddingModel(
+                directory,
+                lora=False,
+                use_pooling=self.opt.pooling,
+                output_hidden_states=False,
+            )
+        self.load_tree_state(directory)
+
+    def _compute_logits(self, q,q_labels,k,k_labels,pos_mask,neg_mask):
+        def cosine_similarity_matrix(q, k):
+            q_norm = F.normalize(q,dim=-1)
+            k_norm = F.normalize(k,dim=-1)
+            cosine_similarity = q_norm@k_norm.T
+            
+            return cosine_similarity
+        
+        def gen_label_mask(relation_matrix,q_labels, k_labels):
+
+            N1 = q_labels.shape[0]
+            N2 = k_labels.shape[0] 
+
+            q_labels_expanded = q_labels.unsqueeze(1).expand(-1, N2)  # N1 x N2
+            k_labels_expanded = k_labels.unsqueeze(0).expand(N1, -1)  # N1 x N2
+
+            result_matrix = relation_matrix[:,q_labels_expanded, k_labels_expanded]
+
+            return result_matrix
+
+        logits=cosine_similarity_matrix(q,k)
+        logits=logits/self.temperature
+        logits = logits.unsqueeze(0).expand(self.K,-1,-1) #K,N1,N2
+
+        pos_mask = gen_label_mask(pos_mask,q_labels, k_labels)
+        neg_mask = gen_label_mask(neg_mask,q_labels, k_labels) #K,N1,N2
+
+        pos_logits = torch.sum(logits*pos_mask,dim=-1)/torch.max(torch.sum(pos_mask,dim=-1),self.esp)#K,N1
+        pos_logits = pos_logits.unsqueeze(-1)#K,N1,1
+        neg_logits = logits*neg_mask#K,N1,N2
+
+        logits = torch.cat((pos_logits, neg_logits), dim=-1)#K,N1,N2+1
+
+        #model:model set
+        # pos_logits_model = torch.sum(logits*same_model,dim=1)/torch.max(torch.sum(same_model,dim=1),self.esp)# N
+        # neg_logits_model=logits*torch.logical_not(same_model)# N,N+K
+        # logits_model=torch.cat((pos_logits_model.unsqueeze(1), neg_logits_model), dim=1)
+
+        return logits
+
+    def forward(self, encoded_batch, labels):
+        q = self.model(encoded_batch)
+        N1 = q.shape[0]
+        k = q.clone().detach()
+        k = self.fabric.all_gather(k).view(-1, k.size(1))
+        k_labels = self.fabric.all_gather(labels).view(-1)
+        
+        now_depth = self.depth[labels].unsqueeze(0).expand(self.K,-1)
+        now_mask = self.mask[:,labels]
+        # leaf_labels = self.root_labels[labels]
+
+        k = torch.concat((k,self.vitual_center),dim=0)
+        k_labels = torch.concat((k_labels,self.center_labels),dim=0)
+
+        logits_sample = self._compute_logits(q,labels,k,k_labels,self.pos_down2up,self.neg_down2up)#K,N1,N2+1
+        gt_sample = torch.zeros(logits_sample.shape[:-1], dtype=torch.long,device=logits_sample.device)
+        logits_sample = logits_sample.permute(0,2,1)
+        loss_smaple1 = F.cross_entropy(logits_sample, gt_sample, reduction='none') #K,N1
+        loss_smaple1 = torch.sum((loss_smaple1/now_depth)*now_mask)/N1*self.max_dep
+
+        # out = self.root_classfier(q)
+        # loss_classfiy = F.cross_entropy(out, leaf_labels)
+
+        loss = loss_smaple1
+
+        return loss,loss_smaple1
+
+    # def forward(self, encoded_batch, labels):
+    #     q = self.model(encoded_batch)
+    #     # N1 = q.shape[0]
+    #     # k = q.clone().detach()
+    #     # k = self.fabric.all_gather(k).view(-1, k.size(1))
+    #     # k_labels = self.fabric.all_gather(labels).view(-1)
+        
+    #     # now_depth = self.depth[labels].unsqueeze(0).expand(self.K,-1)
+    #     # now_mask = self.mask[:,labels]
+    #     leaf_labels = self.root_labels[labels]
+
+    #     # k = torch.concat((k,self.vitual_center),dim=0)
+    #     # k_labels = torch.concat((k_labels,self.center_labels),dim=0)
+
+    #     # logits_sample = self._compute_logits(q,labels,k,k_labels,self.pos_down2up,self.neg_down2up)#K,N1,N2+1
+    #     # gt_sample = torch.zeros(logits_sample.shape[:-1], dtype=torch.long,device=logits_sample.device)
+    #     # logits_sample = logits_sample.permute(0,2,1)
+    #     # loss_smaple1 = F.cross_entropy(logits_sample, gt_sample, reduction='none') #K,N1
+    #     # loss_smaple1 = torch.sum((loss_smaple1/now_depth)*now_mask)/N1*self.max_dep
+
+    #     out = self.root_classfier(q)
+    #     loss_classfiy = F.cross_entropy(out, leaf_labels)
+
+    #     loss = loss_classfiy
+
+    #     return loss,loss_classfiy
+

detree/model/text_embedding.py

@@ -0,0 +1,152 @@
+import os
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, AutoModel
+from peft import LoraConfig, TaskType, PeftModel, get_peft_model
+
+
+class TextEmbeddingModel(nn.Module):
+    """Wrapper around a Hugging Face model with optional LoRA adapters."""
+
+    def __init__(
+        self,
+        model_name,
+        output_hidden_states=False,
+        lora=False,
+        infer=False,
+        use_pooling="average",
+        lora_r=128,
+        lora_alpha=256,
+        lora_dropout=0,
+        adapter_path=None,
+    ):
+        super(TextEmbeddingModel, self).__init__()
+        self.model_name = model_name
+        self.use_pooling = use_pooling
+        self.lora = lora
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+        model_kwargs = {"trust_remote_code": True}
+        if output_hidden_states:
+            model_kwargs["output_hidden_states"] = True
+        self.model = AutoModel.from_pretrained(model_name, **model_kwargs)
+
+        if self.lora:
+            peft_config = LoraConfig(
+                peft_type=TaskType.FEATURE_EXTRACTION,
+                inference_mode=infer,
+                r=lora_r,
+                lora_alpha=lora_alpha,
+                lora_dropout=lora_dropout,
+            )
+            self.model = get_peft_model(self.model, peft_config)
+            if adapter_path is not None:
+                self.load_adapter(adapter_path, is_trainable=not infer)
+            else:
+                self.model.print_trainable_parameters()
+        elif adapter_path is not None:
+            self.model = AutoModel.from_pretrained(adapter_path, **model_kwargs)
+
+    def pooling(self, model_output, attention_mask, hidden_states=False):
+        if hidden_states:
+            if self.use_pooling == "average":
+                model_output.masked_fill(~attention_mask[None, ..., None].bool(), 0.0)
+                emb = model_output.sum(dim=2) / attention_mask.sum(dim=1)[..., None]
+            elif self.use_pooling == "max":
+                emb = model_output.masked_fill(~attention_mask[None, ..., None].bool(), float("-inf"))
+                emb, _ = emb.max(dim=2)
+            elif self.use_pooling == "cls":
+                emb = model_output[:, :, 0]
+            else:
+                raise ValueError("Pooling method not supported")
+            emb = emb.permute(1, 0, 2)
+        else:
+            if self.use_pooling == "average":
+                model_output.masked_fill(~attention_mask[..., None].bool(), 0.0)
+                emb = model_output.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+            elif self.use_pooling == "max":
+                emb = model_output.masked_fill(~attention_mask[..., None].bool(), float("-inf"))
+                emb, _ = emb.max(dim=1)
+            elif self.use_pooling == "cls":
+                emb = model_output[:, 0]
+            else:
+                raise ValueError("Pooling method not supported")
+        return emb
+
+    def forward(self, encoded_batch, hidden_states=False, retrun_all_emb=False):
+        if "t5" in self.model_name.lower():
+            input_ids = encoded_batch['input_ids']
+            decoder_input_ids = torch.zeros((input_ids.shape[0], 1), dtype=torch.long, device=input_ids.device)
+            model_output = self.model(**encoded_batch,
+                                  decoder_input_ids=decoder_input_ids)
+        else:
+            model_output = self.model(**encoded_batch)
+        
+        
+        if isinstance(model_output, tuple):
+            model_output = model_output[0]
+        if isinstance(model_output, dict):
+            if hidden_states:
+                model_output = model_output["hidden_states"]
+                model_output = torch.stack(model_output, dim=0)
+            else:
+                model_output = model_output["last_hidden_state"]
+
+        emb = self.pooling(model_output, encoded_batch['attention_mask'], hidden_states)
+        if retrun_all_emb:
+            return emb, model_output
+        return emb
+
+    def save_pretrained(self, save_directory: str, save_tokenizer: bool = True):
+        os.makedirs(save_directory, exist_ok=True)
+        if isinstance(self.model, PeftModel):
+            self.model.save_pretrained(save_directory)
+        else:
+            self.model.save_pretrained(save_directory)
+        if save_tokenizer:
+            self.tokenizer.save_pretrained(save_directory)
+
+    def load_adapter(self, adapter_path: str, is_trainable: bool = False):
+        if not self.lora or not isinstance(self.model, PeftModel):
+            raise ValueError("LoRA is not enabled for this model instance.")
+        self.model = PeftModel.from_pretrained(
+            self.model.base_model,
+            adapter_path,
+            is_trainable=is_trainable,
+        )
+        self.model.print_trainable_parameters()
+
+    def merge_and_unload(self):
+        if not isinstance(self.model, PeftModel):
+            raise ValueError("The current model does not contain a LoRA adapter to merge.")
+        merged_model = self.model.merge_and_unload()
+        return merged_model
+
+
+class ClassificationHead(nn.Module):
+    """Head for sentence-level classification tasks."""
+
+    def __init__(self, hidden_size,num_labels):
+        super().__init__()
+        self.dense = nn.Linear(hidden_size, hidden_size)
+        self.out_proj = nn.Linear(hidden_size, num_labels)
+
+    def forward(self, x):
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.out_proj(x)
+        return x
+
+class TextClassificationModel(nn.Module):
+    def __init__(self, opt,dim=2):
+        super(TextClassificationModel, self).__init__()
+        self.model = TextEmbeddingModel(opt.model_name,lora=True,use_pooling=opt.pooling,\
+                                        lora_r=opt.lora_r,lora_alpha=opt.lora_alpha,infer=True)
+        self.root_classfier = nn.Linear(opt.embedding_dim, dim)
+
+    def forward(self, encoded_batch):
+        q = self.model(encoded_batch)
+        out = self.root_classfier(q)
+        return out
+
+