Delete nlp/pretrain_bert_mlm.py

nlp/pretrain_bert_mlm.py

@@ -1,153 +0,0 @@
-import codecs
-import collections
-from datetime import datetime
-import json
-import numpy as np
-import random
-
-from datasets import Dataset, DatasetDict, load_dataset
-import torch
-from transformers import (
-    AutoTokenizer,
-    AutoModelForMaskedLM,
-    default_data_collator,
-    DataCollatorForLanguageModeling,
-    TrainingArguments,
-    Trainer,
-)
-
-
-model_path = '../models/bert-base-uncased'
-tokenizer = AutoTokenizer.from_pretrained(model_path, padding=True, truncation=True, max_length=512)
-model = AutoModelForMaskedLM.from_pretrained(model_path)
-
-
-# Data
-def chunkize(text, n_words=300, overlap=150):
-    words = text.split()
-    if len(words) < n_words:
-        return [' '.join(words)]
-    else:
-        return [' '.join(words[i: i+n_words]) for i in range(0, len(words)-n_words+1, n_words-overlap)]
-
-def tokenize_function(examples):
-    result = tokenizer(examples["text"])
-    return result
-
-def group_tokens(examples):
-    # Take a batch of tokens and group them into lines with same chunk size
-    chunk_size = 384
-    n = chunk_size - 2
-    input_ids_all = []
-    cls_id = 101
-    sep_id = 102
-    for item in examples['input_ids']:
-        input_ids_each = item[1: -1]  # Get rid of the first [CLS] and the last [SEP] token
-        input_ids_all += input_ids_each
-    result = {
-        'input_ids': [],
-        'token_type_ids': [],
-        'attention_mask': [],
-    }
-    chunk = []
-    for i in range(len(input_ids_all)):
-        chunk.append(input_ids_all[i])
-        if (i+1) % n == 0:  # complete a chunk
-            result['input_ids'].append([101]+chunk.copy()+[102])
-            result['token_type_ids'].append([0 for j in range(len(chunk)+2)])
-            result['attention_mask'].append([1 for j in range(len(chunk)+2)])
-            chunk = []
-    if len(chunk) > 0:
-        result['input_ids'].append([101]+chunk.copy()+[102])
-        result['token_type_ids'].append([0 for j in range(len(chunk)+2)])
-        result['attention_mask'].append([1 for j in range(len(chunk)+2)])
-    return result
-
-
-def load_data(data_path):
-    with codecs.open(f'{data_path}/train.json', 'r', encoding='utf-8') as f:
-        data = json.load(f)
-        train_data = Dataset.from_dict({'text': data})
-    with codecs.open(f'{data_path}/eval.json', 'r', encoding='utf-8') as f:
-        data = json.load(f)
-        test_data = Dataset.from_dict({'text': data})
-    data_hf = DatasetDict({'train': train_data, 'test': test_data})
-    # Tokenize
-    dataset_tokens = data_hf.map(tokenize_function, batched=True, remove_columns=["text"])
-    # Make each text item have the same length
-    dataset_tokens_group= dataset_tokens.map(group_tokens, batched=True)
-    return dataset_tokens_group
-
-def build_train_valid_data(data_path, target_path, test_size=0.2, n_words=300, overlap=150):
-    # Load raw data
-    with codecs.open(data_path, 'r', encoding='utf-8') as f:
-        data_raw = json.load(f)
-    data_clean = []
-    for text in data_raw:
-        # there are float in data_raw
-        try:
-            nouse = len(text)
-            data_clean.append(text)
-        except:
-            continue
-    # Chunkize
-    data = []
-    for text in data_clean:
-        list_text = chunkize(text, n_words=n_words, overlap=overlap)
-        for r in list_text:
-            data.append(r)
-    n_train = int(len(data)*(1-test_size))
-    set_data = set(data)
-    set_train_data = set(random.sample(set_data, n_train))
-    set_eval_data = set_data - set_train_data
-    train_data = list(set_train_data)
-    eval_data = list(set_eval_data)
-    with codecs.open(f'{target_path}/train.json', 'w', encoding='utf-8') as w:
-        json.dump(train_data, w, ensure_ascii=False)
-    with codecs.open(f'{target_path}/eval.json', 'w', encoding='utf-8') as w:
-        json.dump(eval_data, w, ensure_ascii=False)
-
-# Model
-def train(checkpoints_dir, target_dir, dataset_tokens_group):
-    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=0.15, pad_to_multiple_of=32)
-    batch_size = 64
-    # model_name = 'bert-base-uncased-finetune-mlm-hashtag'
-    training_args = TrainingArguments(
-        output_dir=checkpoints_dir,
-        overwrite_output_dir=True,
-        evaluation_strategy = 'steps',
-        eval_steps = 1000,
-        learning_rate=3e-5,
-        weight_decay=0.01,
-        per_device_train_batch_size=batch_size,
-        per_device_eval_batch_size=batch_size,
-        fp16=True,
-        save_strategy = 'steps',
-        save_steps = 20000,
-        logging_strategy = 'steps',
-        report_to="none",
-        num_train_epochs=160,
-    )
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=dataset_tokens_group["train"],
-        eval_dataset=dataset_tokens_group["test"],
-        data_collator=data_collator,
-        tokenizer=tokenizer,
-    )
-    trainer.train()
-    trainer.save_model(target_dir)
-
-
-if __name__ == '__main__':
-    data_path = '../data/hashtags/dataset_hashtag_english_pretrain.json'
-    target_data_hf_path = './data_hf_for_pretrain_bert_mlm'
-    # build_train_valid_data(data_path, target_data_hf_path, test_size=0.2, n_words=300, overlap=150)
-    # exit()
-    dataset_tokens_group = load_data(target_data_hf_path)
-
-    checkpoints_dir = f'../checkpoints/bert-base-uncased-finetune-mlm-hashtag'
-    now = datetime.strftime(datetime.now(), '%Y_%m_%d_%H_%M')
-    target_dir = f'models_pretrain_bert_mlm/bert-base-uncased-finetune-mlm-hashtag-{now}'
-    train(checkpoints_dir, target_dir, dataset_tokens_group)