from model_mommy import mommy from projects.models import DOCUMENT_CLASSIFICATION, SEQUENCE_LABELING, SEQ2SEQ, SPEECH2TEXT def make_annotation(task, doc, user, **kwargs): annotation_model = { DOCUMENT_CLASSIFICATION: 'Category', SEQUENCE_LABELING: 'Span', SEQ2SEQ: 'TextLabel', SPEECH2TEXT: 'TextLabel' }.get(task) return mommy.make(annotation_model, example=doc, user=user, **kwargs)