# Copyright (c) 2022, salesforce.com, inc.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause

import os
import re
import time
import random
import argparse

import torch

from transformers import GPT2TokenizerFast
from jaxformer.hf.codegen.modeling_codegen import CodeGenForCausalLM



########################################################################
# util


class print_time:
    def __init__(self, desc):
        self.desc = desc

    def __enter__(self):
        print(self.desc)
        self.t = time.time()

    def __exit__(self, type, value, traceback):
        print(f'{self.desc} took {time.time()-self.t:.02f}s')


def set_env():
    os.environ['TOKENIZERS_PARALLELISM'] = 'false'


def set_seed(seed, deterministic=True):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.backends.cudnn.deterministic = deterministic
        torch.backends.cudnn.benchmark = not deterministic
        # torch.use_deterministic_algorithms(deterministic)


def cast(model, fp16=True):
    if fp16:
        model.half()
    return model



########################################################################
# model


def create_model(ckpt, fp16=True):
    if fp16:
        return CodeGenForCausalLM.from_pretrained(ckpt, revision='float16', torch_dtype=torch.float16, low_cpu_mem_usage=True)
    else:
        return CodeGenForCausalLM.from_pretrained(ckpt)


def create_tokenizer():
    t = GPT2TokenizerFast.from_pretrained('gpt2')
    t.max_model_input_sizes['gpt2'] = 1e20
    return t


def include_whitespace(t, n_min=2, n_max=20, as_special_tokens=False):
    t.add_tokens([' ' * n for n in reversed(range(n_min, n_max))], special_tokens=as_special_tokens)
    return t


def include_tabs(t, n_min=2, n_max=20, as_special_tokens=False):
    t.add_tokens(['\t' * n for n in reversed(range(n_min, n_max))], special_tokens=as_special_tokens)
    return t


def create_custom_gpt2_tokenizer():
    t = create_tokenizer()
    t = include_whitespace(t=t, n_min=2, n_max=32, as_special_tokens=False)
    t = include_tabs(t=t, n_min=2, n_max=10, as_special_tokens=False)
    return t


########################################################################
# sample

def sample(
    device,
    model,
    tokenizer,
    context,
    pad_token_id,
    num_return_sequences=1,
    temp=0.2,
    top_p=0.95,
    max_length_sample=128,
    max_length=2048
):

    input_ids = tokenizer(
        context,
        truncation=True,
        padding=True,
        max_length=max_length,
        return_tensors='pt',
    ).input_ids

    input_ids_len = input_ids.shape[1]
    assert input_ids_len < max_length

    with torch.no_grad():
        input_ids = input_ids.to(device)
        tokens = model.generate(
            input_ids,
            do_sample=True,
            num_return_sequences=num_return_sequences,
            temperature=temp,
            max_length=input_ids_len + max_length_sample,
            top_p=top_p,
            pad_token_id=pad_token_id,
            use_cache=True,
        )
        text = tokenizer.batch_decode(tokens[:, input_ids_len:, ...])

    return text


def truncate(completion):

    def find_re(string, pattern, start_pos):
        m = pattern.search(string, start_pos)
        return m.start() if m else -1

    terminals = [
        re.compile(r, re.MULTILINE)
        for r in
        [
            '^#',
            re.escape('<|endoftext|>'),
            "^'''",
            '^"""',
            '\n\n\n'
        ]
    ]

    prints = list(re.finditer('^print', completion, re.MULTILINE))
    if len(prints) > 1:
        completion = completion[:prints[1].start()]

    defs = list(re.finditer('^def', completion, re.MULTILINE))
    if len(defs) > 1:
        completion = completion[:defs[1].start()]

    start_pos = 0

    terminals_pos = [pos for pos in [find_re(completion, terminal, start_pos) for terminal in terminals] if pos != -1]
    if len(terminals_pos) > 0:
        return completion[:min(terminals_pos)]
    else:
        return completion


def test_truncate():

    assert truncate('\nif len_a > len_b:\n    result = a\nelse:\n    result = b\n\n\n\n#') == '\nif len_a > len_b:\n    result = a\nelse:\n    result = b'



########################################################################
# main


def main():

    # (0) constants

    models_nl = ['codegen-350M-nl', 'codegen-2B-nl', 'codegen-6B-nl', 'codegen-16B-nl']
    models_pl = ['codegen-350M-multi', 'codegen-2B-multi', 'codegen-6B-multi', 'codegen-16B-multi', 'codegen-350M-mono', 'codegen-2B-mono', 'codegen-6B-mono', 'codegen-16B-mono']
    models = models_nl + models_pl


    # (1) params

    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, choices=models, default='codegen-350M-mono')
    parser.add_argument('--device', type=str, default='cuda:0')
    parser.add_argument('--rng-seed', type=int, default=42)
    parser.add_argument('--rng-deterministic', type=bool, default=True)
    parser.add_argument('--p', type=float, default=0.95)
    parser.add_argument('--t', type=float, default=0.2)
    parser.add_argument('--max-length', type=int, default=128)
    parser.add_argument('--batch-size', type=int, default=1)
    parser.add_argument('--no-fp16', action="store_false")
    parser.add_argument('--pad', type=int, default=50256)
    parser.add_argument('--context', type=str, default='def helloworld():')
    args = parser.parse_args()


    # (2) preamble

    set_env()
    set_seed(args.rng_seed, deterministic=args.rng_deterministic)

    device = torch.device(args.device)
    if device.type == "cpu":
        args.no_fp16 = False

    if args.model.startswith("codegen-16B"):
        args.no_fp16 = True

    ckpt = f'./checkpoints/{args.model}'


    # (3) load

    with print_time('loading parameters'):
        model = create_model(ckpt=ckpt, fp16=args.no_fp16).to(device)


    with print_time('loading tokenizer'):
        if args.model in models_pl:
            tokenizer = create_custom_gpt2_tokenizer()
        else:
            tokenizer = create_tokenizer()
        tokenizer.padding_side = 'left'
        tokenizer.pad_token = args.pad


    # (4) sample

    with print_time('sampling'):
        completion = sample(device=device, model=model, tokenizer=tokenizer, context=args.context, pad_token_id=args.pad, num_return_sequences=args.batch_size, temp=args.t, top_p=args.p, max_length_sample=args.max_length)[0]
        truncation = truncate(completion)

        print('=' * 100)
        print(completion)
        print('=' * 100)
        print(args.context+truncation)
        print('=' * 100)



if __name__ == '__main__':
    test_truncate()
    main()
    print('done.')

[((40, 4, 40, 21), 'random.seed', 'random.seed', ({(40, 16, 40, 20): 'seed'}, {}), '(seed)', False, 'import random\n'), ((42, 4, 42, 27), 'torch.manual_seed', 'torch.manual_seed', ({(42, 22, 42, 26): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((43, 7, 43, 32), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((69, 8, 69, 49), 'transformers.GPT2TokenizerFast.from_pretrained', 'GPT2TokenizerFast.from_pretrained', ({(69, 42, 69, 48): '"""gpt2"""'}, {}), "('gpt2')", False, 'from transformers import GPT2TokenizerFast\n'), ((191, 13, 191, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((211, 13, 211, 38), 'torch.device', 'torch.device', ({(211, 26, 211, 37): 'args.device'}, {}), '(args.device)', False, 'import torch\n'), ((29, 17, 29, 28), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((44, 8, 44, 36), 'torch.cuda.manual_seed', 'torch.cuda.manual_seed', ({(44, 31, 44, 35): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((63, 15, 63, 126), 'jaxformer.hf.codegen.modeling_codegen.CodeGenForCausalLM.from_pretrained', 'CodeGenForCausalLM.from_pretrained', (), '', False, 'from jaxformer.hf.codegen.modeling_codegen import CodeGenForCausalLM\n'), ((65, 15, 65, 55), 'jaxformer.hf.codegen.modeling_codegen.CodeGenForCausalLM.from_pretrained', 'CodeGenForCausalLM.from_pretrained', ({(65, 50, 65, 54): 'ckpt'}, {}), '(ckpt)', False, 'from jaxformer.hf.codegen.modeling_codegen import CodeGenForCausalLM\n'), ((118, 9, 118, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((142, 8, 142, 35), 're.compile', 're.compile', ({(142, 19, 142, 20): 'r', (142, 22, 142, 34): 're.MULTILINE'}, {}), '(r, re.MULTILINE)', False, 'import re\n'), ((153, 18, 153, 65), 're.finditer', 're.finditer', ({(153, 30, 153, 38): '"""^print"""', (153, 40, 153, 50): 'completion', (153, 52, 153, 64): 're.MULTILINE'}, {}), "('^print', completion, re.MULTILINE)", False, 'import re\n'), ((157, 16, 157, 61), 're.finditer', 're.finditer', ({(157, 28, 157, 34): '"""^def"""', (157, 36, 157, 46): 'completion', (157, 48, 157, 60): 're.MULTILINE'}, {}), "('^def', completion, re.MULTILINE)", False, 'import re\n'), ((146, 12, 146, 38), 're.escape', 're.escape', ({(146, 22, 146, 37): '"""<|endoftext|>"""'}, {}), "('<|endoftext|>')", False, 'import re\n'), ((32, 15, 32, 26), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n')]