| import transformers | |
| from mtm import MultitaskModel | |
| from config import model_name, drop_out | |
| multitask_model = MultitaskModel.create( | |
| model_name=model_name, | |
| model_type_dict={ | |
| "document": transformers.AutoModelForSequenceClassification, | |
| "paragraph": transformers.AutoModelForSequenceClassification, | |
| "sentence": transformers.AutoModelForSequenceClassification, | |
| }, | |
| model_config_dict={ | |
| "document": transformers.AutoConfig.from_pretrained(model_name, num_labels=3, hidden_dropout_prob=drop_out, attention_probs_dropout_prob=drop_out), | |
| "paragraph": transformers.AutoConfig.from_pretrained(model_name, num_labels=3, hidden_dropout_prob=drop_out, attention_probs_dropout_prob=drop_out), | |
| "sentence": transformers.AutoConfig.from_pretrained(model_name, num_labels=3, hidden_dropout_prob=drop_out, attention_probs_dropout_prob=drop_out), | |
| }, | |
| ) | |
| tokenizer = transformers.AutoTokenizer.from_pretrained(model_name) | |