#### Model-1 Roberta Context_features+SRL+WSD
class model_v1(torch.nn.Module):

    def __init__(self,device: str = "cuda") -> None:
        super().__init__()

        # Prepare the device
        self.device = torch.device(device)
        #### Context Model
        self.Context_model=AutoModel.from_pretrained('roberta-base')
        self.Context_model.to(self.device)
        
        # Freeze the weights of the Context model
        for param in self.Context_model.parameters():
            param.requires_grad = False
#         print(self.Context_model)
        for i in range(1,3+1):
               for j in self.Context_model.encoder.layer[-i].parameters():
                    j.requires_grad=True
        ### WSD embeddings
        self.wsd_embedding=CustomEmbedding(wsd_vocab.wsd_vectors.shape[0],wsd_vocab.wsd_vectors.shape[1],device=self.device)#,wsd_vocab.name_to_index['PAD'])
        self.wsd_embedding.initialize_embeddings(wsd_vocab.wsd_vectors)
        
        
        # Freeze the weights of the WSD embedding
        for param in self.wsd_embedding.parameters():
#             print(param)
            param.requires_grad = False
#             print(param)
            
        self.wsd_model=nn.Sequential(*[self.wsd_embedding,nn.Linear(wsd_vocab.wsd_vectors.shape[1],32,device=self.device) ,nn.LayerNorm((32,), eps=1e-05, elementwise_affine=True,device=self.device)])

        ### Synset embeddings
        self.srl_embedding_size=16
        self.srl_embedding=CustomEmbedding(len(srl_vocab.tags),self.srl_embedding_size,srl_vocab.tag_to_idx['PAD'],self.device)
        self.srl_model=nn.Sequential(*[torch.nn.GRU(self.srl_embedding_size,int( self.srl_embedding_size/2), num_layers=2, bias=True, batch_first=True,dropout=0.0, bidirectional=True, device=self.device, dtype=None)])#,nn.LayerNorm((self.srl_embedding_size), eps=1e-05, elementwise_affine=True,device=self.device)])
        
        ######## combine m1 srl information from premise and m2 label sequeneces from hypothesis

        self.srl_premise_down=nn.Sequential(*[nn.Linear(self.srl_model[0].hidden_size*2*train_dataset.m1,self.srl_model[0].hidden_size*2*train_dataset.m2,device=self.device),nn.LayerNorm((self.srl_model[0].hidden_size*2*train_dataset.m2), eps=1e-05, elementwise_affine=True,device=self.device)])
        self.srl_hypothesis_down=nn.Sequential(*[nn.Linear(self.srl_model[0].hidden_size*2*train_dataset.m2,self.srl_model[0].hidden_size*2*train_dataset.m2,device=self.device),nn.LayerNorm((self.srl_model[0].hidden_size*2*train_dataset.m2), eps=1e-05, elementwise_affine=True,device=self.device)])
        
        
        ####### sequence representation
        self.out_hidden_embedding_size=832
        self.out_hidden=torch.nn.LSTM( self.out_hidden_embedding_size,int(  self.out_hidden_embedding_size/2), num_layers=2, bias=True, batch_first=True,dropout=0.0, bidirectional=True, device=self.device, dtype=None)
        ##### Output layer
        
        self.out_layer=nn.Sequential(*[nn.Linear(832,out_features=3,device=self.device)])#nn.Linear(1024,out_features=3,device=self.device)
        
        ###error log
        self.error=[]
        
    def forward(self, inp_list:list):
        batch = move_to_device(inp_list, self.device)
        
        input_ids,attention_mask,final_word_ids,final_word_offsets,wsd_input_ids,srl_premise_input_ids, \
        srl_hypothesis_input_ids,srl_final_offset,srl_premise_length,srl_hypothesis_length=batch #
    
#         print([i for i in [input_ids,attention_mask,final_word_ids,final_word_offsets,wsd_input_ids,srl_premise_input_ids, \
#         srl_hypothesis_input_ids,srl_final_offset,srl_premise_length,srl_hypothesis_length,label]])
        srl_premise_input_ids=srl_premise_input_ids.reshape(-1,train_dataset.srl_premise_max_length)
        srl_hypothesis_input_ids=srl_hypothesis_input_ids.reshape(-1,train_dataset.srl_premise_max_length)
        srl_premise_length=srl_premise_length.reshape(-1)
        srl_hypothesis_length=srl_hypothesis_length.reshape(-1)
        bpe_info=self.build_sparse_offsets(
        final_word_offsets ,
        attention_mask ,
        512,
        )
        out=self.Context_model(input_ids=input_ids,attention_mask=attention_mask,output_hidden_states=True)
        tmp=torch.cat([out.hidden_states[-i].unsqueeze(0) for i in range(3)],dim=0)
        embeddings=torch.mean(tmp,dim=0)
        
        word_embeddings=self.merge_sparse(embeddings,bpe_info)
        ### word based embedding
        wsd_embedding_down=self.wsd_model(wsd_input_ids)
        word_mask=wsd_input_ids!=117660
        word_mask=word_mask.unsqueeze(-1).expand_as(wsd_embedding_down)
        wsd_embeddings=wsd_embedding_down*word_mask

        
        ###### srl embedding+BI-LSTM

        #### premise
        srl_emb=self.srl_embedding(srl_premise_input_ids)
        srl_inp=pack_padded_sequence(srl_emb,srl_premise_length.cpu(),batch_first=True,enforce_sorted=False)
#         out, (hidden_state, cell_state)=self.srl_model[0](srl_inp)
        out, hidden_state=self.srl_model[0](srl_inp)
        out_premise=pad_packed_sequence(out,batch_first=True)
        

        out_premise=out_premise[0]#self.srl_model[1](out_premise[0])

        ### rearrange
        out_premise=out_premise.reshape(-1,train_dataset.m1,out_premise.shape[-2],out_premise.shape[-1])

        out_premise=out_premise.permute(0,2,1,3).reshape(out_premise.shape[0],out_premise.shape[2],-1)

        ### down project
        out_premise=self.srl_premise_down(out_premise)

        #### zeroing the padding tokens represenation: This is required due to way of combining premise and hypothesis information
        
        srl_premise_mask=self.create_binary_mask(srl_premise_length[0::4])
        srl_premise_mask=srl_premise_mask.unsqueeze(-1).expand_as(out_premise).to(self.device)
        out_premise=out_premise*srl_premise_mask
        
         ### do padding and expand the sequence length till 512
        padding_size = 600 - out_premise.shape[1]
        out_premise = F.pad(out_premise, (0, 0, 0, padding_size, 0, 0), mode='constant', value=0)
       

        ### hypothesis

        srl_emb=self.srl_embedding(srl_hypothesis_input_ids)
        srl_inp=pack_padded_sequence(srl_emb,srl_hypothesis_length.cpu(),batch_first=True,enforce_sorted=False)
#         out, (hidden_state, cell_state)=self.srl_model[0](srl_inp)
        out, hidden_state=self.srl_model[0](srl_inp)
        out_hypothesis=pad_packed_sequence(out,batch_first=True)
        # print(out_premise[0].shape)
#         out_hypothesis=self.srl_model[1](out_hypothesis[0])
        out_hypothesis=out_hypothesis[0]
        # print(out_premise[0].shape)
        ### rearrange
        out_hypothesis=out_hypothesis.reshape(-1,train_dataset.m2,out_hypothesis.shape[-2],out_hypothesis.shape[-1])
        out_hypothesis=out_hypothesis.permute(0,2,1,3).reshape(out_hypothesis.shape[0],out_hypothesis.shape[2],-1)
        
        ### down project: As based on current config no need to down project
        out_hypothesis=self.srl_hypothesis_down(out_hypothesis)
        ### zero mask hypothesis for the pad tokens
        srl_hypothesis_mask=self.create_binary_mask(srl_hypothesis_length[0::2])
        srl_hypothesis_mask=srl_hypothesis_mask.unsqueeze(-1).expand_as(out_hypothesis).to(self.device)
        out_hypothesis=out_hypothesis*srl_hypothesis_mask
        
        
        ### do padding and expand the sequence length till 512
        padding_size = 600 - out_hypothesis.shape[1]
        out_hypothesis = F.pad(out_hypothesis, (0, 0, 0, padding_size, 0, 0), mode='constant', value=0)
       
       
        ### shift each row of each hypothesis by length of actual tokens present in the premise
        shifts=torch.tensor([ [j for i in range(out_hypothesis.shape[-1])] for j in srl_premise_length[0::4]]).to(self.device)
        
        out_hypothesis=self.custom_roll(out_hypothesis, shifts, dim=1)
        ### final srl out
        srl_out=out_premise+out_hypothesis
       
        #### get aligned 

        srl_attention_mask=(srl_final_offset!=-2)*1
        bpe_srl_info=self.build_sparse_offsets(
                srl_final_offset ,
                srl_attention_mask ,
                600,
            )
        srl_embeddings=self.merge_sparse(srl_out,bpe_srl_info)[:,:512,:]
        
        ### final embedding contextual+wsd+srl
#         print(srl_embeddings)
        final_embedding=torch.cat((word_embeddings,srl_embeddings,wsd_embeddings),dim=-1)
        
        
    
#         ### average of the embeddings
#         word_mask=wsd_input_ids!=117660
#         word_mask=word_mask.unsqueeze(-1).expand_as(final_embedding)
#         average_embedding=(final_embedding*word_mask).sum(dim=1)/(word_mask.sum(dim=1))
        
    
        ### average of the embeddings
        word_mask=wsd_input_ids!=117660
        merged_emb_seq_len=word_mask.sum(dim=1)
        merge_inp=pack_padded_sequence(final_embedding,merged_emb_seq_len.cpu(),batch_first=True,enforce_sorted=False)
        out, (hidden_state, cell_state)=self.out_hidden(merge_inp)
#         out, hidden_state=self.out_hidden(merge_inp)
        ### get the logits
        hidden_state=hidden_state[:2,:,:].permute(1,0,2).reshape(hidden_state.shape[1],-1)
        
        out=self.out_layer(hidden_state)
        
        return out
        

        
    def build_sparse_offsets(self,
        offsets: torch.Tensor ,
        bpe_mask: torch.Tensor ,
        max_words_per_sentence: int,
    ) :
        """Build tensors used as info for BPE pooling, starting from the BPE offsets.
        """
        if not isinstance(offsets, torch.Tensor):
            offsets: torch.Tensor = torch.as_tensor(offsets)
        if not isinstance(bpe_mask, torch.Tensor):
            bpe_mask: torch.Tensor = torch.as_tensor(bpe_mask)

        sentence_lengths: torch.Tensor = bpe_mask.sum(dim=1).to(self.device)
#         print('sentence_lengths \n:',sentence_lengths)
        # We want to build triplets as coordinates (document, word, bpe)
        # We start by creating the document index for each triplet
        document_indices = torch.arange(offsets.size(0)).to(self.device).repeat_interleave(
            sentence_lengths
        )
#         print('document_indices: \n',document_indices,document_indices.size())
        
        # then the word indices
        word_indices = offsets[offsets != -2]
#         print('word_indices \n:',word_indices,word_indices.size())
#         print('len of word_indices',len(word_indices))
        # lastly the bpe indices
        max_range: torch.Tensor = torch.arange(bpe_mask.shape[1],device=self.device)
        bpe_indices: torch.LongTensor = torch.cat(
            [max_range[:i] for i in bpe_mask.sum(dim=1)], dim=0
        ).long().to(self.device)
#         print('bpe_indices \n:',bpe_indices,bpe_indices.size())
#         print('bpe_indices shape', bpe_indices.shape)
        unique_words, word_lengths = torch.unique_consecutive(
            offsets, return_counts=True
        )
        unique_words=unique_words.to(self.device)
        word_lengths= word_lengths.to(self.device)
#         print('unique words \n:',unique_words)
#         print('word lengths \n:',word_lengths)
        unpadded_word_lengths = word_lengths[unique_words != -2]
#         print('unpadded_word_lengths \n:',unpadded_word_lengths)
#         print(unpadded_word_lengths)
        # and their weight to be used as multiplication factors for normalisation
        bpe_weights: torch.FloatTensor = (
            (1 / unpadded_word_lengths).repeat_interleave(unpadded_word_lengths).float()
        ).to(self.device)
#         print('bpe weights \n:',bpe_weights)
        
        try:
            sparse_indices = torch.stack(
                [document_indices, word_indices, bpe_indices], dim=0
            ).to(self.device)
#             print('sparse indices \n:',sparse_indices)
        except:
            self.error.append([document_indices, word_indices, bpe_indices])
            print(document_indices.shape, word_indices.shape, bpe_indices.shape)
            print('error')
#         print(sparse_indices)
        bpe_shape = torch.Size(
            (
                bpe_mask.size(0),  # batch_size
                max_words_per_sentence,  # max number of words per sentence
                bpe_mask.size(1),  # max bpe_number in batch wrt the sentence
            )
        )
        
        return dict(
            sparse_indices=sparse_indices,
            sparse_values=bpe_weights,
            sparse_size=bpe_shape,
        )
    def merge_sparse(self,
        embeddings: torch.Tensor, bpe_info
    ) -> torch.Tensor:
        """
        Merges the subword embeddings into a single tensor, using sparse indices.

        Args:
            embeddings (`torch.Tensor`):
                The embeddings tensor.
            bpe_info (`Mapping[str, Any]`, `optional`):
                The BPE info.

        Returns:
            `torch.Tensor`: The merged embeddings.
        """
        # it is constructed here and not in the tokenizer/collate because pin_memory is not sparse-compatible
        bpe_weights = torch.sparse_coo_tensor(
            indices=bpe_info["sparse_indices"],
            values=bpe_info["sparse_values"],
            size=bpe_info["sparse_size"],
        )
        # (sentence, word, bpe) x (sentence, bpe, transformer_dim) -> (sentence, word, transformer_dim)
        merged = torch.bmm(bpe_weights.to_dense(), embeddings)
        return merged
        
    def custom_roll(self,tensor, shifts, dim):
        # Ensure shifts is a tensor
        if not isinstance(shifts, torch.Tensor):
            shifts = torch.tensor(shifts, device=tensor.device)

        # Create an index tensor
        idx = torch.arange(tensor.size(dim), device=tensor.device)

        # Broadcast the index tensor to match the input tensor's shape
        idx_shape = [1] * tensor.dim()
        idx_shape[dim] = tensor.size(dim)
        idx = idx.view(idx_shape).expand_as(tensor)

        # Apply the shifts
        idx = (idx - shifts.unsqueeze(dim)) % tensor.size(dim)

        # Use gather to reorder the tensor
        return torch.gather(tensor, dim, idx)
    
    def create_binary_mask(self,counts, width=None):
        # Convert counts to tensor if it's not already
        if not isinstance(counts, torch.Tensor):
            counts = torch.tensor(counts, dtype=torch.long)

        # Determine the width of the tensor
        if width is None:
            width = counts.max().item()

        # Create a range tensor
        range_tensor = torch.arange(width, device=counts.device)

        # Create the binary tensor using broadcasting and cumsum
        binary_tensor = (range_tensor < counts.unsqueeze(1)).float()

        return binary_tensor
NM
I'm currently working on build a model for nli nlp task. The file contains my model definition. Upon braod experments, the model achives similar performance to a model which is just using contextualised features instead of contextual+wsd+srl. Please suggest ways through which I can integrate/utilise the information coming from three soruces more effectively