import torch
import torch.nn as nn




if __name__ == "__main__":
    ems = torch.tensor([
        [[0.3, 0.3, 0], [0, 1, 0], [1, 0.2, 0.2],[1, 0.2, 0.2], [1, 0.2, 0.2]],
        [[1, 0, 0], [0, 0.2, 0], [1, 0.2, 0.2], [1, 0.2, 0.2], [1, 0.2, 0.2]],
        [[1, 0, 0], [0, 1, 1.4], [1, 0.2, 0.2], [1, 0.2, 0.2], [1, 0.2, 0.2]],
        [[1, 0, 0], [0, 1, 0], [1, 0.2, 0.2], [1, 0.2, 0.2], [1, 0.2, 0.2]]
    ])
    labels = torch.tensor([
        [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 1], [0, 0, 1]],
        [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 1], [0, 0, 1]],
        [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 1], [0, 0, 1]],
        [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 1], [0, 0, 1]],
    ])
    crf = CrfLayer(label2idx={"a": 0, "b": 1, "c": 2})
    crf(emission=ems, labels=labels)
