# Copyright (c) Meta Platforms, Inc. and affiliates.

# This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.

from typing import Optional

import fire

from llama import Llama

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, set_seed

import torch

def api_preload(

    ckpt_dir: str,

    tokenizer_path: str,

    max_seq_len: int = 512,

    max_batch_size: int = 8,

):

    print(">> start loading model")

    generator = Llama.build(

        ckpt_dir=ckpt_dir,

        tokenizer_path=tokenizer_path,

        max_seq_len=max_seq_len,

        max_batch_size=max_batch_size,

        model_parallel_size=1

    )

    print(">> model loaded")

    return generator

def api_generator(instructions,

                  generator,

                  temperature: float = 0.2,

                  top_p: float = 0.95,

                  max_gen_len: Optional[int] = None,):

    results = generator.chat_completion(

        instructions,  # type: ignore

        max_gen_len=max_gen_len,

        temperature=temperature,

        top_p=top_p,

    )

    return results

def api_preload_hf(

    ckpt_dir: str,

    tokenizer_path: str = None,

    max_seq_len: int = 512,

    max_batch_size: int = 8,

):

    print(">> start loading model")

    tokenizer = AutoTokenizer.from_pretrained(ckpt_dir, padding_side="left")

    model = AutoModelForCausalLM.from_pretrained(ckpt_dir)

    model.to('cuda')

    generator = pipeline('text-generation', model=model, tokenizer=tokenizer, device=0)

    print(">> model loaded")

    return generator

def api_generator_hf(instructions,

                  generator):

    _prompt = ''

    results = generator(instructions[0][0]['content'], renormalize_logits=True, do_sample=True, use_cache=True, max_new_tokens=10)

    return results

def main(

    ckpt_dir: str,

    tokenizer_path: str,

    temperature: float = 0.2,

    top_p: float = 0.95,

    max_seq_len: int = 512,

    max_batch_size: int = 8,

    max_gen_len: Optional[int] = None,

):

    generator = Llama.build(

        ckpt_dir=ckpt_dir,

        tokenizer_path=tokenizer_path,

        max_seq_len=max_seq_len,

        max_batch_size=max_batch_size,

    )

    instructions = [

        [

            {

                "role": "user",

                "content": "In Bash, how do I list all text files in the current directory (excluding subdirectories) that have been modified in the last month?",

            }

        ],

        [

            {

                "role": "user",

                "content": "What is the difference between inorder and preorder traversal? Give an example in Python.",

            }

        ],

        [

            {

                "role": "system",

                "content": "Provide answers in JavaScript",

            },

            {

                "role": "user",

                "content": "Write a function that computes the set of sums of all contiguous sublists of a given list.",

            }

        ],

    ]

    results = generator.chat_completion(

        instructions,  # type: ignore

        max_gen_len=max_gen_len,

        temperature=temperature,

        top_p=top_p,

    )

    for instruction, result in zip(instructions, results):

        for msg in instruction:

            print(f"{msg['role'].capitalize()}: {msg['content']}\n")

        print(

            f"> {result['generation']['role'].capitalize()}: {result['generation']['content']}"

        )

        print("\n==================================\n")

def test1():

    ckpt_dir = 'codellama-main/CodeLlama-7b-Instruct/'

    tokenizer_path = 'codellama-main/CodeLlama-7b-Instruct/tokenizer.model'

    generator = api_preload(ckpt_dir=ckpt_dir,

                            tokenizer_path=tokenizer_path)

    instructions = [

        [

            {

                "role": "user",

                "content": "What is the difference between inorder and preorder traversal? Give an example in Python.",

            }

        ],

    ]

    results = api_generator(instructions=instructions, generator=generator)

    for instruction, result in zip(instructions, results):

        for msg in instruction:

            print(f"{msg['role'].capitalize()}: {msg['content']}\n")

        print(

            f"> {result['generation']['role'].capitalize()}: {result['generation']['content']}"

        )

        print("\n==================================\n")

def test2():

    ckpt_dir = 'codellama-main/CodeLlama-7b-Instruct-hf/'

    tokenizer_path = 'codellama-main/CodeLlama-7b-Instruct-hf/'

    generator = api_preload_hf(ckpt_dir=ckpt_dir)

    instructions = [

        [

            {

                "role": "user",

                "content": "What is the difference between inorder and preorder traversal? Give an example in Python.",

            }

        ],

    ]

    results = api_generator_hf(instructions=instructions, generator=generator)

    print(results)

    for instruction, result in zip(instructions, results):

        for msg in instruction:

            print(f"{msg['role'].capitalize()}: {msg['content']}\n")

        print(

            f"> {result['generation']['role'].capitalize()}: {result['generation']['content']}"

        )

        print("\n==================================\n")

def api_preload_deepseek(

    ckpt_dir: str,

    tokenizer_path: str = None,

    cpu = False

):

    print(">> start loading model")

    tokenizer = AutoTokenizer.from_pretrained(

        tokenizer_path,

        trust_remote_code=True)

    if cpu:

        generator = AutoModelForCausalLM.from_pretrained(

            ckpt_dir,

            trust_remote_code=True,

            torch_dtype=torch.bfloat16)

    else:

        if '67b' in ckpt_dir:

            generator = AutoModelForCausalLM.from_pretrained(

                ckpt_dir,

                trust_remote_code=True,

                torch_dtype=torch.bfloat16,

                load_in_8bit=True)

        else:

            generator = AutoModelForCausalLM.from_pretrained(

                ckpt_dir,

                trust_remote_code=True,

                torch_dtype=torch.bfloat16).cuda()

    print(">> model loaded")

    return tokenizer, generator

def api_generator_deepseek(instructions,

                  tokenizer,

                  generator,

                  max_new_tokens = 512,

                  top_k = 50,

                  top_p = 0.95):

    messages = instructions[0]

    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(generator.device)

    # tokenizer.eos_token_id is the id of <|EOT|> token

    outputs = generator.generate(inputs, max_new_tokens=max_new_tokens, do_sample=False, top_k=top_k, top_p=top_p, num_return_sequences=1,

                             eos_token_id=tokenizer.eos_token_id)

    content = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)

    results = [

        {

        'generation': {

            'role': 'AI',

            'content': content

        }

        }

    ]

    return results

def test3():

    ckpt_dir = 'deepseek/deepseek-coder-7b-instruct-v1.5'

    tokenizer_path = 'deepseek/deepseek-coder-7b-instruct-v1.5'

    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)

    model = AutoModelForCausalLM.from_pretrained(ckpt_dir, trust_remote_code=True,

                                                 torch_dtype=torch.bfloat16).cuda()

    messages = [

        {'role': 'user', 'content': "write a quick sort algorithm in python."}

    ]

    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)

    # tokenizer.eos_token_id is the id of <|EOT|> token

    outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, top_k=50, top_p=0.95, num_return_sequences=1,

                             eos_token_id=tokenizer.eos_token_id)

    print(tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True))

if __name__ == "__main__":

    import os

    import torch.distributed as dist

    os.environ['MASTER_ADDR'] = 'localhost'

    os.environ['MASTER_PORT'] = '5678'

    dist.init_process_group(backend='nccl', init_method='env://', rank=0, world_size=1)

    test3()