{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This notebook illustrates how to use Masked Language Modeling for this competition.\n\nObservation: most of the dataset names consist of only words with uppercased-first-letter and some stopwords like `on`, `in`, `and` (e.g. `Early Childhood Longitudinal Study`, `Trends in International Mathematics and Science Study`). \n\nThus, one approach to find the datasets is: \n- Locate all the sequences of capitalized words (these sequences may contain some stopwords), \n- Replace each sequence with one of 2 special symbols (e.g. `$` and `#`), implying if that sequence represents a dataset name or not.\n- Have the model learn the MLM task.","metadata":{}},{"cell_type":"markdown","source":"The code below shows how to train a model for that purpose with the help of the `huggingface`.","metadata":{}},{"cell_type":"code","source":"MAX_SAMPLE = None # set a small number (e.g. 50) for experimentation, set None for production.\nSEED = 42\n\nmodel_checkpoint = \"bert-base-cased\"\n\nMAX_LENGTH = 64\nOVERLAP = 20\nLENGTH = 1\nFL_TH = 0.75\n\nDATASET_SYMBOL = '$' # this symbol represents a dataset name\nNONDATA_SYMBOL = '#' # this symbol represents a non-dataset name","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:21:39.673118Z","iopub.execute_input":"2021-06-04T23:21:39.673733Z","iopub.status.idle":"2021-06-04T23:21:39.678511Z","shell.execute_reply.started":"2021-06-04T23:21:39.673571Z","shell.execute_reply":"2021-06-04T23:21:39.677830Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Install packages","metadata":{}},{"cell_type":"code","source":"!pip install datasets --no-index --find-links=file:///kaggle/input/coleridge-packages/packages/datasets\n!pip install ../input/coleridge-packages/seqeval-1.2.2-py3-none-any.whl\n!pip install ../input/coleridge-packages/tokenizers-0.10.1-cp37-cp37m-manylinux1_x86_64.whl\n!pip install ../input/coleridge-packages/transformers-4.5.0.dev0-py3-none-any.whl\n\n\nimport os\nimport re\nimport json\nimport time\nimport datetime\nimport random\nimport glob\nimport importlib\n\nimport numpy as np\nimport pandas as pd\n\nfrom tqdm import tqdm\n\nimport torch\nfrom datasets import load_dataset\nfrom transformers import AutoTokenizer, DataCollatorForLanguageModeling, \\\nAutoModelForMaskedLM, Trainer, TrainingArguments, pipeline, AutoConfig\n\nfrom IPython.display import clear_output\n\n\nclear_output()","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:21:39.683266Z","iopub.execute_input":"2021-06-04T23:21:39.683692Z","iopub.status.idle":"2021-06-04T23:22:15.238377Z","shell.execute_reply.started":"2021-06-04T23:21:39.683662Z","shell.execute_reply":"2021-06-04T23:22:15.237099Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# https://huggingface.co/transformers/_modules/transformers/trainer_utils.html\ndef set_seed(seed: int):\n    \"\"\"\n    Helper function for reproducible behavior to set the seed in ``random``, ``numpy``, ``torch`` and/or ``tf`` (if\n    installed).\n\n    Args:\n        seed (:obj:`int`): The seed to set.\n    \"\"\"\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    # ^^ safe to call this function even if cuda is not available\n    \n    print(f'Setted Pipeline SEED = {SEED}')\n\n\nset_seed(SEED)","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:22:15.240021Z","iopub.execute_input":"2021-06-04T23:22:15.240323Z","iopub.status.idle":"2021-06-04T23:22:15.249572Z","shell.execute_reply.started":"2021-06-04T23:22:15.240292Z","shell.execute_reply":"2021-06-04T23:22:15.248493Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def text_cleaning(text):\n    '''\n    Converts all text to lower case, Removes special charecters, emojis and multiple spaces\n    text - Sentence that needs to be cleaned\n    '''\n    text = re.sub('[^A-Za-z0-9]+', ' ', str(text).lower()).strip()\n    text = re.sub(' +', ' ', text)\n    emoji_pattern = re.compile(\"[\"\n                               u\"\\U0001F600-\\U0001F64F\"  # emoticons\n                               u\"\\U0001F300-\\U0001F5FF\"  # symbols & pictographs\n                               u\"\\U0001F680-\\U0001F6FF\"  # transport & map symbols\n                               u\"\\U0001F1E0-\\U0001F1FF\"  # flags (iOS)\n                               \"]+\", flags=re.UNICODE)\n    text = emoji_pattern.sub(r'', text)\n    return text\n\n\ndef jaccard_similarity(s1, s2):\n    l1 = s1.split(\" \")\n    l2 = s2.split(\" \")    \n    intersection = len(list(set(l1).intersection(l2)))\n    union = (len(l1) + len(l2)) - intersection\n    return float(intersection) / union\n\n\ndef find_data_sets_id(publication_id: int) -> list:\n    data_set_ids = []\n    for class_ in data_set_citations:\n        if class_['publication_id'] == publication_id:\n            data_set_ids.append(class_['data_set_id'])\n    return data_set_ids\n\n\ndef find_data_set_citations_mention_list(publication_id: int) -> str:\n    mention_list = []\n    for class_ in data_set_citations:\n        if class_['publication_id'] == publication_id:\n            mention_list.append( '|'.join(class_['mention_list']) )\n    return '|'.join( [label for label in mention_list if label != ''] )\n\n\ndef find_data_sets_title(data_set_id: int) -> str:\n    for class_ in data_sets:\n        if class_['data_set_id'] == data_set_id:\n            return class_['title']\n        \n\ndef RichContextDF(publications: '.json') -> pd.DataFrame:\n    publication_id = []\n    text_file_name = []\n    citations_mention_list = []\n    data_sets_title = []\n    for class_ in publications:\n        publication_id.append(class_['publication_id']) # to get data_set_citations\n        text_file_name.append(class_['text_file_name']) # to get text_file\n        \n        # to get citations_mention_list\n        citations_mention_list.append( find_data_set_citations_mention_list(class_['publication_id']) )\n        \n        # to get data_sets_title\n        data_sets_title_temp = []\n        for data_sets_id in find_data_sets_id( class_['publication_id'] ):\n            data_sets_title_temp.append( find_data_sets_title(data_sets_id) )\n        data_sets_title.append('||'.join(data_sets_title_temp))\n    \n    return pd.DataFrame({\n        'publication_id': publication_id,\n        'text_file_name': text_file_name,\n        'citations_mention_list': citations_mention_list,\n        'data_sets_title': data_sets_title\n    })\n\n\ndef clean_paper_sentence(s):\n    \"\"\"\n    This function is essentially clean_text without lowercasing.\n    \"\"\"\n    s = re.sub('[^A-Za-z0-9]+', ' ', str(s)).strip()\n    s = re.sub(' +', ' ', s)\n    return s\n\n\ndef shorten_sentences(sentences):\n    \"\"\"\n    Sentences that have more than MAX_LENGTH words will be split\n    into multiple sentences with overlappings.\n    \"\"\"\n    short_sentences = []\n    for sentence in sentences:\n        words = sentence.split()\n        if len(words) > MAX_LENGTH:\n            for p in range(0, len(words), MAX_LENGTH - OVERLAP):\n                short_sentences.append(' '.join(words[p:p+MAX_LENGTH]))\n        else:\n            short_sentences.append(sentence)\n    return short_sentences\n\n\ndef find_sublist(big_list, small_list):\n    \"\"\"\n    find all positions of $small_list in $big_list.\n    \"\"\"\n    all_positions = []\n    for i in range(len(big_list) - len(small_list) + 1):\n        if small_list == big_list[i:i+len(small_list)]:\n            all_positions.append(i)\n    \n    return all_positions\n\n\ndef jaccard_similarity_list(l1, l2):\n    \"\"\"\n    Return the Jaccard Similarity score of 2 lists.\n    \"\"\"\n    intersection = len(list(set(l1).intersection(l2)))\n    union = (len(l1) + len(l2)) - intersection\n    return float(intersection) / union\n\n\nconnection_tokens = {'s', 'of', 'and', 'in', 'on', 'for', 'data', 'dataset'}\ndef find_negative_candidates(sentence, labels):\n    \"\"\"\n    Extract negative samples for Masked Dataset Modeling from a given $sentence.\n    A negative candidate should be a continuous sequence of at least 2 words, \n    each of these words either has the first letter in uppercase or is one of\n    the connection words ($connection_tokens). Furthermore, the connection \n    tokens are not allowed to appear at the beginning and the end of the\n    sequence. Lastly, the sequence must be quite different to any of the \n    ground truth labels (measured by Jaccard similarity).\n    \"\"\"\n    def candidate_qualified(words, labels):\n        while len(words) and words[0].lower() in connection_tokens:\n            words = words[1:]\n        while len(words) and words[-1].lower() in connection_tokens:\n            words = words[:-1]\n        \n        return len(words) >= 2 and \\\n               all(jaccard_similarity_list(words, label) < 0.75 for label in labels)\n    \n    candidates = []\n    \n    phrase_start, phrase_end = -1, -1\n    for id in range(1, len(sentence)):\n        word = sentence[id]\n        if word[0].isupper() or word in connection_tokens:\n            if phrase_start == -1:\n                phrase_start = phrase_end = id\n            else:\n                phrase_end = id\n        else:\n            if phrase_start != -1:\n                if candidate_qualified(sentence[phrase_start:phrase_end+1], labels):\n                    candidates.append((phrase_start, phrase_end))\n                phrase_start = phrase_end = -1\n    \n    if phrase_start != -1:\n        if candidate_qualified(sentence[phrase_start:phrase_end+1], labels):\n            candidates.append((phrase_start, phrase_end))\n    \n    return candidates","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:22:15.251341Z","iopub.execute_input":"2021-06-04T23:22:15.251620Z","iopub.status.idle":"2021-06-04T23:22:15.272730Z","shell.execute_reply.started":"2021-06-04T23:22:15.251583Z","shell.execute_reply":"2021-06-04T23:22:15.271501Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load data","metadata":{}},{"cell_type":"code","source":"# train\ntrain_path = '../input/coleridgeinitiative-show-us-the-data/train.csv'\npaper_train_folder = '../input/coleridgeinitiative-show-us-the-data/train'\n\ntrain = pd.read_csv(train_path)\nprint('train size before agg.:', len(train))\n\ntrain = train[:MAX_SAMPLE]\n# Group by publication, training labels should have the same form as expected output.\ntrain = train.groupby('Id').agg({\n    'pub_title': 'first',\n    'dataset_title': '|'.join,\n    'dataset_label': '|'.join,\n    'cleaned_label': '|'.join\n}).reset_index()    \nprint('train size after agg.:', len(train))\n\ntrain = train.sort_values(by=['Id'])\ntrain.head()","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:22:15.274359Z","iopub.execute_input":"2021-06-04T23:22:15.274751Z","iopub.status.idle":"2021-06-04T23:22:15.867422Z","shell.execute_reply.started":"2021-06-04T23:22:15.274711Z","shell.execute_reply":"2021-06-04T23:22:15.866490Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Pseudo","metadata":{}},{"cell_type":"code","source":"# train\npseudo_train_path = '../input/coleridge-pseudolabelsv2-0585/submission.csv'\npaper_train_folder = '../input/coleridgeinitiative-show-us-the-data/train'\n\npseudo_train = pd.read_csv(pseudo_train_path)\nprint('pseudo_train size before drop_duplicates.:', len(pseudo_train))\n\npseudo_train = pseudo_train[:MAX_SAMPLE]\npseudo_train = pseudo_train.drop_duplicates(subset='Id')\nprint('pseudo_train size after drop_duplicates:', len(pseudo_train), '\\n')\n\npseudo_train = pseudo_train.sort_values(by=['Id']).reset_index(drop=True)\npseudo_train.head()","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:22:15.870998Z","iopub.execute_input":"2021-06-04T23:22:15.871292Z","iopub.status.idle":"2021-06-04T23:22:15.963154Z","shell.execute_reply.started":"2021-06-04T23:22:15.871264Z","shell.execute_reply":"2021-06-04T23:22:15.962197Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Original Add Pseudo","metadata":{}},{"cell_type":"code","source":"train['dataset_label'] = train['dataset_label'] + '|' + pseudo_train['PredictionString']\ntrain.head()","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:22:15.965160Z","iopub.execute_input":"2021-06-04T23:22:15.965420Z","iopub.status.idle":"2021-06-04T23:22:15.998721Z","shell.execute_reply.started":"2021-06-04T23:22:15.965394Z","shell.execute_reply":"2021-06-04T23:22:15.997978Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_temp = []\nfor labels in tqdm(train['dataset_label']):\n    labels = labels.split('|')\n    filtered_labels = []\n    for label in labels:\n        if len(filtered_labels) == 0 or all(jaccard_similarity(text_cleaning(label), text_cleaning(got_label)) < FL_TH for got_label in filtered_labels):\n            filtered_labels.append(label)\n    train_temp.append('|'.join(filtered_labels))\ntrain['dataset_label'] = train_temp\ntrain.head()","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:22:15.999673Z","iopub.execute_input":"2021-06-04T23:22:16.000020Z","iopub.status.idle":"2021-06-04T23:22:16.769576Z","shell.execute_reply.started":"2021-06-04T23:22:15.999993Z","shell.execute_reply":"2021-06-04T23:22:16.768532Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## External","metadata":{}},{"cell_type":"code","source":"with open(f'../input/rich-context-competition-train-testtargz/train_test/publications.json', 'r') as f:\n    publications = json.load(f)\n    \nwith open(f'../input/rich-context-competition-train-testtargz/train_test/data_set_citations.json', 'r') as f:\n    data_set_citations = json.load(f)\n    \nwith open(f'../input/rich-context-competition-train-testtargz/train_test/data_sets.json', 'r') as f:\n    data_sets = json.load(f)","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:22:16.771027Z","iopub.execute_input":"2021-06-04T23:22:16.771425Z","iopub.status.idle":"2021-06-04T23:22:17.313908Z","shell.execute_reply.started":"2021-06-04T23:22:16.771384Z","shell.execute_reply":"2021-06-04T23:22:17.312882Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"RichContext_train = RichContextDF(publications)\nRichContext_train = RichContext_train[ RichContext_train['citations_mention_list'] != '' ]\nRichContext_train","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:22:17.314914Z","iopub.execute_input":"2021-06-04T23:22:17.315148Z","iopub.status.idle":"2021-06-04T23:22:20.767496Z","shell.execute_reply.started":"2021-06-04T23:22:17.315124Z","shell.execute_reply":"2021-06-04T23:22:20.766479Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Prepare data for train MLM","metadata":{}},{"cell_type":"markdown","source":"### Extract positive and negative samples","metadata":{}},{"cell_type":"code","source":"corpus = []\ncnt_pos = 0\ncnt_neg = 0","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:22:20.768542Z","iopub.execute_input":"2021-06-04T23:22:20.768791Z","iopub.status.idle":"2021-06-04T23:22:20.773020Z","shell.execute_reply.started":"2021-06-04T23:22:20.768766Z","shell.execute_reply":"2021-06-04T23:22:20.772300Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pbar = tqdm(total = len(RichContext_train))\nfor paper_id, dataset_labels in RichContext_train[['publication_id', 'citations_mention_list']].itertuples(index=False):\n    labels = [clean_paper_sentence(label).split() for label in dataset_labels.split('|')]\n    # papers preparation\n    with open(f'../input/rich-context-competition-train-testtargz/train_test/files/text/{paper_id}.txt', 'r') as f:\n        paper = f.readlines()\n    paper = [line[:-1] for line in paper]\n    content = ' '.join(paper)\n    sentences = set([clean_paper_sentence(sentence) for sentence in content.split('.')])\n    sentences = shorten_sentences(sentences)\n    sentences = [sentence for sentence in sentences if len(sentence) > LENGTH]\n    sentences = [sentence.split() for sentence in sentences]\n\n    # positive samples\n    for sentence in sentences:\n        for label in labels:\n            for pos in find_sublist(sentence, label):\n                dt_point = sentence[:pos] + [DATASET_SYMBOL] + sentence[pos+len(label):]\n                corpus.append(' '.join(dt_point))\n                cnt_pos += 1\n\n    # negative samples\n    for sentence in sentences:\n        sentence_str = ' '.join(sentence)\n        if all(w not in sentence_str for w in {'data', 'study'}):\n            continue\n        for phrase_start, phrase_end in find_negative_candidates(sentence, labels):\n            dt_point = sentence[:phrase_start] + [NONDATA_SYMBOL] + sentence[phrase_end+1:]\n            corpus.append(' '.join(dt_point))\n            cnt_neg += 1\n\n    # process bar\n    pbar.update(1)\n    pbar.set_description(f'Training data size: {cnt_pos} postives + {cnt_neg} negatives')\npbar.close()","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:22:20.774396Z","iopub.execute_input":"2021-06-04T23:22:20.774661Z","iopub.status.idle":"2021-06-04T23:23:27.066376Z","shell.execute_reply.started":"2021-06-04T23:22:20.774633Z","shell.execute_reply":"2021-06-04T23:23:27.065577Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pbar = tqdm(total = len(train))\nfor paper_id, dataset_labels in train[['Id', 'dataset_label']].itertuples(index=False):\n    labels = [clean_paper_sentence(label).split() for label in dataset_labels.split('|')]\n    # papers preparation\n    with open(f'{paper_train_folder}/{paper_id}.json', 'r') as f:\n        paper = json.load(f)\n    content = '. '.join(section['text'] for section in paper)\n    sentences = set([clean_paper_sentence(sentence) for sentence in content.split('.')])\n    sentences = shorten_sentences(sentences)\n    sentences = [sentence for sentence in sentences if len(sentence) > LENGTH]\n    sentences = [sentence.split() for sentence in sentences]\n    \n    # positive samples\n    for sentence in sentences:\n        for label in labels:\n            for pos in find_sublist(sentence, label):\n                dt_point = sentence[:pos] + [DATASET_SYMBOL] + sentence[pos+len(label):]\n                corpus.append(' '.join(dt_point))\n                cnt_pos += 1\n    \n    # negative samples\n    for sentence in sentences:\n        sentence_str = ' '.join(sentence)\n        if all(w not in sentence_str for w in {'data', 'study'}):\n            continue\n        for phrase_start, phrase_end in find_negative_candidates(sentence, labels):\n            dt_point = sentence[:phrase_start] + [NONDATA_SYMBOL] + sentence[phrase_end+1:]\n            corpus.append(' '.join(dt_point))\n            cnt_neg += 1\n    \n    # process bar\n    pbar.update(1)\n    pbar.set_description(f'Training data size: {cnt_pos} postives + {cnt_neg} negatives')\npbar.close()","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:23:27.067500Z","iopub.execute_input":"2021-06-04T23:23:27.067741Z","iopub.status.idle":"2021-06-04T23:27:28.457674Z","shell.execute_reply.started":"2021-06-04T23:23:27.067717Z","shell.execute_reply":"2021-06-04T23:27:28.456595Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Save data to a file","metadata":{}},{"cell_type":"code","source":"with open('train_mlm.json', 'w') as f:\n    for sentence in corpus:\n        row_json = {'text':sentence}\n        json.dump(row_json, f)\n        f.write('\\n')","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:27:28.459158Z","iopub.execute_input":"2021-06-04T23:27:28.459571Z","iopub.status.idle":"2021-06-04T23:27:31.814852Z","shell.execute_reply.started":"2021-06-04T23:27:28.459529Z","shell.execute_reply":"2021-06-04T23:27:31.813786Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Fine-tune the Transformer","metadata":{}},{"cell_type":"code","source":"datasets = load_dataset('json',\n            data_files={'train' : 'train_mlm.json'},\n            )\n\ndatasets[\"train\"][:5]","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:27:31.816063Z","iopub.execute_input":"2021-06-04T23:27:31.816351Z","iopub.status.idle":"2021-06-04T23:27:32.608265Z","shell.execute_reply.started":"2021-06-04T23:27:31.816320Z","shell.execute_reply":"2021-06-04T23:27:32.607537Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Tokenize and collate data","metadata":{}},{"cell_type":"code","source":"tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:27:32.609306Z","iopub.execute_input":"2021-06-04T23:27:32.609724Z","iopub.status.idle":"2021-06-04T23:27:33.573515Z","shell.execute_reply.started":"2021-06-04T23:27:32.609687Z","shell.execute_reply":"2021-06-04T23:27:33.572400Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def tokenize_function(examples):\n    return tokenizer(examples[\"text\"])\n\ntokenized_datasets = datasets.map(tokenize_function, batched=True, num_proc=1, remove_columns=[\"text\"])","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:27:33.574816Z","iopub.execute_input":"2021-06-04T23:27:33.575082Z","iopub.status.idle":"2021-06-04T23:27:59.927828Z","shell.execute_reply.started":"2021-06-04T23:27:33.575055Z","shell.execute_reply":"2021-06-04T23:27:59.926816Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15)","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:27:59.929205Z","iopub.execute_input":"2021-06-04T23:27:59.929771Z","iopub.status.idle":"2021-06-04T23:27:59.934157Z","shell.execute_reply.started":"2021-06-04T23:27:59.929727Z","shell.execute_reply":"2021-06-04T23:27:59.933278Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Load pre-trained model and fine-tune","metadata":{}},{"cell_type":"code","source":"model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:27:59.935661Z","iopub.execute_input":"2021-06-04T23:27:59.936016Z","iopub.status.idle":"2021-06-04T23:28:18.707331Z","shell.execute_reply.started":"2021-06-04T23:27:59.935976Z","shell.execute_reply":"2021-06-04T23:28:18.706498Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"training_args = TrainingArguments(\n    output_dir=\"output-mlm\",\n    evaluation_strategy = \"no\",\n    learning_rate=2e-5,\n    weight_decay=0.01,\n    save_steps=12000,\n    num_train_epochs=2,\n    report_to=\"none\",\n)\n\ntrainer = Trainer(\n    model=model,\n    args=training_args,\n    train_dataset=tokenized_datasets[\"train\"],\n    data_collator=data_collator,\n)","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:28:18.708546Z","iopub.execute_input":"2021-06-04T23:28:18.708817Z","iopub.status.idle":"2021-06-04T23:28:19.140130Z","shell.execute_reply.started":"2021-06-04T23:28:18.708789Z","shell.execute_reply":"2021-06-04T23:28:19.139335Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.train()","metadata":{"execution":{"iopub.status.busy":"2021-06-04T23:28:19.141212Z","iopub.execute_input":"2021-06-04T23:28:19.141649Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Save model","metadata":{}},{"cell_type":"code","source":"trainer.model.save_pretrained('mlm-model')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Save tokenizer","metadata":{}},{"cell_type":"code","source":"config = AutoConfig.from_pretrained(model_checkpoint)\n\ntokenizer.save_pretrained('model_tokenizer')\nconfig.save_pretrained('model_tokenizer')","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}