# Script for obtaining Chinese LLaMA weights from the OpenChineseLLaMA project (https://github.com/OpenLMLab/OpenChineseLLaMA)
# Due to the License of LLaMA, we only provide a delta-version patch
# Adding the patch to the original LLaMA weights makes the Chinese LLaMA weights
import os
import sys
sys.path.append(os.path.abspath(__file__).rsplit('/', 2)[0])
import shutil
import torch
import argparse
from util.misc import download

parser = argparse.ArgumentParser()
parser.add_argument(
    "--llama_dir", default="/path/to/llama", type=str,
    help="Path to official LLaMA weights",
)
args = parser.parse_args()

ori_path = os.path.join(args.llama_dir, '7B')
delta_path = os.path.join(args.llama_dir, '7B_chinese_delta')
new_path = os.path.join(args.llama_dir, '7B_chinese')


import dataclasses
from enum import auto, Enum
from typing import List, Tuple

class SeparatorStyle(Enum):
    """Different separator style."""
    SINGLE = auto()
    TWO = auto()

@dataclasses.dataclass
class Conversation:
    """A class that keeps all conversation history."""
    system: str
    roles: List[str]
    messages: List[List[str]]
    offset: int
    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
    sep: str = "###"
    sep2: str = None


import argparse
import os
import sys

import fairscale.nn.model_parallel.initialize as fs_init
import models_llama_adapter
import torch
import torch.distributed as dist
from conversation import conv_templates, SeparatorStyle
from util import misc

from llama import LLaMA, Tokenizer

def load_model(args, load_8bit=False):
    model = models_llama_adapter.__dict__[args.model_name](args)
    model.eval()
    if args.model_path is None:
        print("Warning: not loading instruct tuned weights.")
    else:
        print("Using instruct tuned weights from:", args.model_path)


import functools
import json
from pathlib import Path

import fairscale.nn.model_parallel.initialize as fs_init
import torch
import torch.distributed as dist

from llama import ModelArgs, Tokenizer, Transformer

def _load_and_redistribute_checkpoint(llama_model_path, model_name):
    with open(Path(llama_model_path) / model_name / "params.json") as f:
        params = json.load(f)
    tokenizer = Tokenizer(model_path=str(Path(llama_model_path) / "tokenizer.model"))
    print("Using model path: %s, model_name: %s" % (llama_model_path, model_name))

    checkpoints = (Path(llama_model_path) / model_name).glob("*.pth")
    checkpoints = sorted(checkpoints)

    mp_world_size = fs_init.get_model_parallel_world_size()


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

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# DeiT: https://github.com/facebookresearch/deit
# BEiT: https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------

import builtins
import datetime
import os
import subprocess
import time
from collections import defaultdict, deque
from pathlib import Path

import torch
import torch.distributed as dist


# Copyright (c) Meta Platforms, Inc. and affiliates.
# This software may be used and distributed according to the terms of the GNU General Public License version 3.

from typing import List

import torch

from llama.model import Transformer
from llama.tokenizer import Tokenizer

class LLaMA:
    def __init__(self, model: Transformer, tokenizer: Tokenizer) -> None:
        self.model = model
        self.tokenizer = tokenizer

    def generate(
        self,
        prompts: List[str],
        max_gen_len: int,
        temperature: float = 0.8,


# Copyright (c) Meta Platforms, Inc. and affiliates.
# This software may be used and distributed according to the terms of the GNU General Public License version 3.

import math
from dataclasses import dataclass
from typing import Optional, Tuple

import fairscale.nn.model_parallel.initialize as fs_init
import torch
import torch.nn.functional as F
from fairscale.nn.model_parallel.layers import ColumnParallelLinear, ParallelEmbedding, RowParallelLinear
from torch import nn

@dataclass
class ModelArgs:
    dim: int = 512
    n_layers: int = 8
    n_heads: int = 8
    vocab_size: int = -1  # defined later by tokenizer
    multiple_of: int = 256  # make SwiGLU hidden layer size multiple of large power of 2


# Copyright (c) Meta Platforms, Inc. and affiliates.
# This software may be used and distributed according to the terms of the GNU General Public License version 3.

from .generation import LLaMA
from .model import ModelArgs, Transformer
from .tokenizer import Tokenizer


# Copyright (c) Meta Platforms, Inc. and affiliates.
# This software may be used and distributed according to the terms of the GNU General Public License version 3.

import os
from logging import getLogger
from typing import List

from sentencepiece import SentencePieceProcessor

logger = getLogger()

class Tokenizer:
    def __init__(self, model_path: str):
        # reload tokenizer
        assert os.path.isfile(model_path), model_path
        self.sp_model = SentencePieceProcessor(model_file=model_path)
        logger.info(f"Reloaded SentencePiece model from {model_path}")

        # BOS / EOS token IDs
        self.n_words: int = self.sp_model.vocab_size()