data_augment.py

import argparse
import logging
import math

import numpy as np
import torch

from tqdm import tqdm

from transformers import (
    CTRLLMHeadModel,
    CTRLTokenizer,
    GPT2LMHeadModel,
    GPT2Tokenizer,
    OpenAIGPTLMHeadModel,
    OpenAIGPTTokenizer,
    TransfoXLLMHeadModel,
    TransfoXLTokenizer,
    XLMTokenizer,
    XLMWithLMHeadModel,
    XLNetLMHeadModel,
    XLNetTokenizer, pipeline,
)

from datasets import load_dataset

import utils

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)

MAX_LENGTH = int(10000)  # Hardcoded max length to avoid infinite loop

MODEL_CLASSES = {
    "gpt2": (GPT2LMHeadModel, GPT2Tokenizer),
    "ctrl": (CTRLLMHeadModel, CTRLTokenizer),
    "openai-gpt": (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
    "xlnet": (XLNetLMHeadModel, XLNetTokenizer),
    "transfo-xl": (TransfoXLLMHeadModel, TransfoXLTokenizer),
    "xlm": (XLMWithLMHeadModel, XLMTokenizer),
}


def set_seed(args):
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


def adjust_length_to_model(length, max_sequence_length):
    if length < 0 and max_sequence_length > 0:
        length = max_sequence_length
    elif 0 < max_sequence_length < length:
        length = max_sequence_length  # No generation bigger than model size
    elif length < 0:
        length = MAX_LENGTH  # avoid infinite loop
    return length


def main():
    """# Data augmentation based on conditional text generation using the auto-regressive models of the library: GPT, GPT-2, Transformer-XL, XLNet, CTRL."""
    parser = argparse.ArgumentParser()

    parser.add_argument('--input', type=str, help="input file to augment by sampling prefixes.")
    parser.add_argument('--output', type=str, help="output file to store augmented data")
    parser.add_argument("--model_type", default='gpt2', type=str,
                        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
                        )
    parser.add_argument("--model_name_or_path", default='gpt2', type=str,
                        help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(
                            MODEL_CLASSES.keys()),
                        )
    parser.add_argument("--stop_token", type=str, default=utils.EOS,
                        help="Token at which text generation is stopped")
    parser.add_argument("--num_samples", type=int, default=100, help="number of samples "
                                                                     "generated by each "
                                                                     "conditional text generation")
    parser.add_argument("--temperature", type=float, default=1.0,
                        help="temperature of 1.0 has no effect, lower tend toward greedy sampling",
                        )
    parser.add_argument("--repetition_penalty", type=float, default=1.0,
                        help="primarily useful for CTRL model; in that case, use 1.2"
                        )
    parser.add_argument("--k", type=int, default=0)
    parser.add_argument("--p", type=float, default=0.9)

    parser.add_argument("--xlm_language", type=str, default="",
                        help="Optional language when used with the XLM model.")

    parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
    parser.add_argument("--no_cuda", action="store_true", help="Avoid using CUDA when available")
    parser.add_argument("--num_return_sequences", type=int, default=1,
                        help="The number of samples to generate.")
    parser.add_argument("--fp16", action="store_true",
                        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",
                        )
    args = parser.parse_args()

    args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
    args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()

    logger.warning(
        "device: %s, n_gpu: %s, 16-bits training: %s",
        args.device,
        args.n_gpu,
        args.fp16,
    )

    set_seed(args)

    # Initialize the model and tokenizer
    try:
        args.model_type = args.model_type.lower()
        model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    except KeyError:
        raise KeyError(
            "the model {} you specified is not supported. You are welcome to add it and open a PR :)")

    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
    model = model_class.from_pretrained(args.model_name_or_path)
    model.to(args.device)

    if args.fp16:
        model.half()

    # initialize text generation pipeline
    text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)

    logger.info(args)

    # map between the prefix and its number of occurrences in the input file
    label2count = dict()
    label2count['0'] = 0
    label2count['1'] = 0

    with open(args.input, 'r', encoding='utf8') as file_in:
        for sentence in file_in:
            sentence = sentence.split('\t')
            label2count[sentence[0]] += 1

    total_count = sum(label2count.values())
    factor = int(math.ceil(args.num_return_sequences / args.num_samples))
    p0 = label2count['0'] / total_count
    p1 = 1 - p0

    prefixes = np.random.choice([0, 1], size=(factor,), p=[p0, p1])

    augmented = set()
    with open(args.output, 'w', encoding='utf8') as file:
        for prefix in tqdm(prefixes):
            output_sequences = text_generator(text_inputs=str(prefix) + '\t', early_stopping=True,
                                              temperature=args.temperature,
                                              top_k=args.k,
                                              top_p=args.p,
                                              repetition_penalty=args.repetition_penalty,
                                              do_sample=True,
                                              num_return_sequences=args.num_samples,
                                              clean_up_tokenization_spaces=True,
                                              return_full_text=True)
            for seq in output_sequences:
                text = seq['generated_text']
                text = text[:text.find(args.stop_token) if args.stop_token and text.find(
                    args.stop_token) > -1 else None].strip()
                text = text[:text.find('\n') if text.find('\n') > -1 else None].strip()
                if text not in augmented:
                    file.write(text + '\n')
                    augmented.add(text)
            file.flush()
        # concatenate training set to augmented set
        dataset = load_dataset('glue', 'sst2', split='train')
        for d in dataset:
            file.write(str(d['label']) + '\t' + d['sentence'] + '\n')

    # randomly shuffle the lines of the file
    lines = open(args.output, encoding='utf8').readlines()
    np.random.shuffle(lines)
    open(args.output, 'w', encoding='utf8').writelines(lines)


if __name__ == "__main__":
    main()