def METHOD_NAME(self):
    model_state = {
        **super().METHOD_NAME(),
        "document_embeddings": self.embeddings.save_embeddings(use_state_dict=False),
        "label_dictionary": self.label_dictionary,
        "label_type": self.label_type,
        "multi_label": self.multi_label,
        "multi_label_threshold": self.multi_label_threshold,
        "weight_dict": self.weight_dict,
    }
    return model_state