How to use and implement CRF with sample weights of 2D array to calculate the loss function?

I am currently working on a problem of sequence tagging to give sequence of tags to the sequences of sentences in a document and there are multiple documents. So my input is 3D (batch_size,max_len_sentences,no_of_words). I also want to provide sample weight within the training model during using fit function. And the shape for sample weight is 2D matrix (batch_size,max_len_sentences). However The build in code for calculating loss function with crf_log_liklihood accodmodates 1D weights and not the 2D weights and return a log_liklihood of shape [batch_size] which means one value of logliklihood for each document and not for each sequence of sentences. I also made some padded sentences in some documents to make the length of sentences equal in each document.

I am using keras_crf: (link given) keras-crf/ at main · luozhouyang/keras-crf · GitHub which usually tensorflow addons CRF layer.

Hi, did you find a solution for this ? I am struggling too :frowning: