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

Hi,
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) https://github.com/luozhouyang/keras-crf/blob/main/keras_crf/crf_model.py which usually tensorflow addons CRF layer.

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