# Copyright (c) OpenMMLab. All rights reserved.
from datasets import load_dataset
from mmengine.dataset import DefaultSampler
from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
                            LoggerHook, ParamSchedulerHook)
from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
from torch.optim import AdamW
from transformers import AutoModelForCausalLM, AutoTokenizer

from xtuner.dataset import process_hf_dataset
from xtuner.dataset.collate_fns import default_collate_fn
from xtuner.dataset.map_fns import alpaca_map_fn, template_map_fn_factory
from xtuner.engine.hooks import (DatasetInfoHook, EvaluateChatHook,
                                 VarlenAttnArgsToMessageHubHook)
from xtuner.engine.runner import TrainLoop
from xtuner.model import SupervisedFinetune
from xtuner.parallel.sequence import SequenceParallelSampler
from xtuner.utils import PROMPT_TEMPLATE, SYSTEM_TEMPLATE

#######################################################################
#                          PART 1  Settings                           #
#######################################################################
# Model
pretrained_model_name_or_path = '/newdisk/wuzr/models/Meta-Llama-3-8B/'
use_varlen_attn = False

# Data
alpaca_en_path = '/newdisk/wuzr/xtuner/data/trix_instruct.json'
prompt_template = PROMPT_TEMPLATE.llama3_chat
max_length = 4096
pack_to_max_length = True

# parallel
sequence_parallel_size = 1

# Scheduler & Optimizer
batch_size = 1  # per_device
accumulative_counts = 16
accumulative_counts *= sequence_parallel_size
dataloader_num_workers = 0
max_epochs = 3
optim_type = AdamW
lr = 2e-5
betas = (0.9, 0.999)
weight_decay = 0
max_norm = 1  # grad clip
warmup_ratio = 0.03

# Save
save_steps = 1000
save_total_limit = 2  # Maximum checkpoints to keep (-1 means unlimited)

# Evaluate the generation performance during the training
evaluation_freq = 1000
SYSTEM = SYSTEM_TEMPLATE.alpaca
evaluation_inputs = [
    '## Question\ngive the number of ships that were launched in 1878.\n\n## Table\ncolumn 0 | Country | Builder | Location | Ship | Class / type\n6 March | United States | John Roach and Son | Chester, Pennsylvania | City of Rio de Janeiro | Passenger ship\n13 May | Germany | Kaiserliche Werft Wilhelmshaven | Wilhelmshaven | Bayern | Sachsen-class ironclad\n13 June | United Kingdom | Royal Dockyard | Devonport, Devon | Pegasus | Doterel-class sloop\n31 August | United Kingdom | Royal Dockyard | Sheerness | Gannet | Doterel-class sloop\n23 October | Norway | Karljohansverns Verft | Horten | Nor | Vale-class gunboat\n1 November | Norway | Karljohansverns Verft | Horten | Brage | Vale-class gunboat\n9 November | Germany | A. G. Vulcan | Stettin | Württemberg | Sachsen-class ironclad\n\n## Task:\nYou will answer the question based on the given context. You should reach a short-form answer after reasoning.\nYou are asked to answer the question in three steps.\n1. Analyze the question and the given context. Make up a plan to answer the question.\n2. Write one or more SQL to query the table for necessary information and output expected execution result.\n3. Reason step-by-step to reach the final answer.\n\n## Answer:',
    "## Question\nwhat is the difference in height between key tower and 55 public square\n\n## Table\nRank | Name | Image | Height ft (m) | Floors | Year | Notes\n1 | Key Tower |  | 947 (289) | 57 | 1991 | 104th-tallest building in the world 20th-tallest building in the United States Has been the tallest building in the city and state since 1991 Stood as the tallest building in the United States between New York City and Chicago from its completion until 2007, when Comcast Center in Philadelphia was completed Tallest building constructed in Cleveland in the 1990s\n2 | Terminal Tower |  | 723 (220) | 52 | 1930 | 114th-tallest building in the United States Stood as the tallest building in the world outside of New York City until 1964 Tallest building constructed in the city in the 1930s\n3 | 200 Public Square |  | 658 (201) | 45 | 1985 | Also known as the BP Building Tallest building constructed in the city in the 1980s\n4 | Tower at Erieview |  | 529 (161) | 40 | 1964 | Tallest building constructed in Cleveland in the 1960s\n5 | One Cleveland Center |  | 450 (137) | 31 | 1983 | \n6 | Fifth Third Center |  | 446 (136) | 27 | 1992 | \n7 | Federal Court House Tower |  | 430 (131) | 23 | 2002 | Tallest building constructed in the city in the 2000s Most recently completed skyscraper in the city\n8 | Justice Center Complex |  | 420 (128) | 26 | 1977 | Tallest building constructed in the city in the 1970s\n9 | Anthony J. Celebrezze Federal Building |  | 419 (128) | 31 | 1967 | \n10 | PNC Center |  | 410 (125) | 35 | 1980 | Originally known as the National City Center; building was renamed in 2009\n11 | AT Tower |  | 390 (119) | 28 | 1971 | Previously known as Cleveland Trust Tower Currently being redeveloped as a mixed use hotel, retail, and residential building attached to the new Cuyahoga County Headquarters Also known as 900 Euclid Tower\n12 | AT&T Huron Road Building |  | 365 (111) | 24 | 1927 | Commonly known as Ohio Bell Buildinh Previously known as the Ameritech Building Tallest building constructed in Cleveland in the 1920s\n13 | Rhodes Tower |  | 363 (111) | 20 | 1971 | Originally known as the University Tower\n14 | Eaton Center |  | 356 (109) | 28 | 1983 | \n15 | Ernst & Young Tower |  | 330 (101) | 21 | 2013 | Phase I of the Flats East Bank redevelopment project First downtown private office building constructed since 1992\n16 | Marriott at Key Center |  | 320 (98) | 28 | 1991 | Tallest all-hotel building in the city\n17 | McDonald Investment Center |  | 308 (94) | 23 | 1968 | Also known as Key Center Formerly known as the Central National Bank Building\n18 | 55 Public Square |  | 300 (91) | 22 | 1958 | Tallest building constructed in the city the 1950s Originally known as the Illuminating Building\n19 | Huntington Bank Building | — | 289 (88) | 21 | 1924 | \n20 | North Point Tower |  | 285 (87) | 20 | 1990 | \n21= | Diamond Building |  | 282 (86) | 23 | 1972 | \n21= | Standard Building |  | 282 (86) | 21 | 1925 | \n23 | 1717 East Ninth Building | — | 275 (84) | 21 | 1959 | Also known as the East Ohio Building\n24 | Keith Building |  | 272 (83) | 21 | 1922 | \n25= | East Tower |  | 266 (81) | 25 | 1973 | Tallest all-residential building in the city Also known as the Reserve Square Apartments\n25= | Embassy Suites Tower |  | 266 (81) | 26 | 1969 | Also known as Embassy Suites at Reserve Square\n27 | Superior Building |  | 265 (81) | 22 | 1922 | \n28 | Fenn Tower |  | 265 (81) | 21 | 1930 | \n29 | Landmark Office Towers |  | 260 (79) | 22 | 1930 | \n30= | Penton Media Building | — | 253 (77) | 21 | 1972 | \n30= | Ohio Savings Plaza | — | 253 (77) | 17 | 1969 | \n30= | Ameritech Center |  | 253 (77) | 16 | 1983 | \n\n## Task:\nYou will answer the question based on the given context. You should reach a short-form answer after reasoning.\nYou are asked to answer the question in three steps.\n1. Analyze the question and the given context. Make up a plan to answer the question.\n2. Write one or more SQL to query the table for necessary information and output expected execution result.\n3. Reason step-by-step to reach the final answer.\n\n## Answer:"
]

#######################################################################
#                      PART 2  Model & Tokenizer                      #
#######################################################################
tokenizer = dict(
    type=AutoTokenizer.from_pretrained,
    pretrained_model_name_or_path=pretrained_model_name_or_path,
    trust_remote_code=True,
    padding_side='right')

model = dict(
    type=SupervisedFinetune,
    use_varlen_attn=use_varlen_attn,
    llm=dict(
        type=AutoModelForCausalLM.from_pretrained,
        pretrained_model_name_or_path=pretrained_model_name_or_path,
        trust_remote_code=True))

#######################################################################
#                      PART 3  Dataset & Dataloader                   #
#######################################################################
alpaca_en = dict(
    type=process_hf_dataset,
    #dataset=dict(type=load_dataset, path=alpaca_en_path),
    dataset=dict(
        type=load_dataset, path = 'json', data_files = dict(train = alpaca_en_path)
    ),
    tokenizer=tokenizer,
    max_length=max_length,
    dataset_map_fn=alpaca_map_fn,
    template_map_fn=dict(
        type=template_map_fn_factory, template=prompt_template),
    remove_unused_columns=True,
    shuffle_before_pack=True,
    pack_to_max_length=pack_to_max_length,
    use_varlen_attn=use_varlen_attn)

sampler = SequenceParallelSampler \
    if sequence_parallel_size > 1 else DefaultSampler
train_dataloader = dict(
    batch_size=batch_size,
    num_workers=dataloader_num_workers,
    dataset=alpaca_en,
    sampler=dict(type=sampler, shuffle=True),
    collate_fn=dict(type=default_collate_fn, use_varlen_attn=use_varlen_attn))

#######################################################################
#                    PART 4  Scheduler & Optimizer                    #
#######################################################################
# optimizer
optim_wrapper = dict(
    type=AmpOptimWrapper,
    optimizer=dict(
        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
    accumulative_counts=accumulative_counts,
    loss_scale='dynamic',
    dtype='float16')

# learning policy
# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
param_scheduler = [
    dict(
        type=LinearLR,
        start_factor=1e-5,
        by_epoch=True,
        begin=0,
        end=warmup_ratio * max_epochs,
        convert_to_iter_based=True),
    dict(
        type=CosineAnnealingLR,
        eta_min=0.0,
        by_epoch=True,
        begin=warmup_ratio * max_epochs,
        end=max_epochs,
        convert_to_iter_based=True)
]

# train, val, test setting
train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)

#######################################################################
#                           PART 5  Runtime                           #
#######################################################################
# Log the dialogue periodically during the training process, optional
custom_hooks = [
    dict(type=DatasetInfoHook, tokenizer=tokenizer),
    dict(
        type=EvaluateChatHook,
        tokenizer=tokenizer,
        every_n_iters=evaluation_freq,
        evaluation_inputs=evaluation_inputs,
        system=SYSTEM,
        prompt_template=prompt_template)
]

if use_varlen_attn:
    custom_hooks += [dict(type=VarlenAttnArgsToMessageHubHook)]

# configure default hooks
default_hooks = dict(
    # record the time of every iteration.
    timer=dict(type=IterTimerHook),
    # print log every 10 iterations.
    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
    # enable the parameter scheduler.
    param_scheduler=dict(type=ParamSchedulerHook),
    # save checkpoint per `save_steps`.
    checkpoint=dict(
        type=CheckpointHook,
        by_epoch=False,
        interval=save_steps,
        max_keep_ckpts=save_total_limit),
    # set sampler seed in distributed evrionment.
    sampler_seed=dict(type=DistSamplerSeedHook),
)

# configure environment
env_cfg = dict(
    # whether to enable cudnn benchmark
    cudnn_benchmark=False,
    # set multi process parameters
    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
    # set distributed parameters
    dist_cfg=dict(backend='nccl'),
)

# set visualizer
visualizer = None

# set log level
log_level = 'INFO'

# load from which checkpoint
load_from = None

# whether to resume training from the loaded checkpoint
resume = False

# Defaults to use random seed and disable `deterministic`
randomness = dict(seed=None, deterministic=False)

# set log processor
log_processor = dict(by_epoch=False)