import pandas as pd  
from datasets import Dataset, load_dataset, concatenate_datasets  

### Tagengo  
gpt4_dataset = load_dataset("lightblue/tagengo-gpt4", split="train")  
gpt4_dataset = gpt4_dataset.filter(lambda x: x["response"][1] == "stop")  
####  

### Megagon  
megagon_df = pd.read_json(  
    "https://raw.githubusercontent.com/megagonlabs/instruction_ja/main/data/data.jsonl",  
    lines=True,  
    orient="records"  
    )  
role_map = {"user": "human", "agent": "gpt"}  
megagon_df["conversations"] = megagon_df.utterances.apply(lambda x: [{"from": role_map[y["name"]], "value": y["text"]} for y in x])  
megagon_df["language"] = "Japanese"  
megagon_df = megagon_df[["conversations", "language"]]  
megagon_dataset = Dataset.from_pandas(df)  
###  

### Openchat  
openchat_df = pd.read_json("https://huggingface.co/datasets/openchat/openchat_sharegpt4_dataset/resolve/main/sharegpt_gpt4.json?download=true")  
openchat_df["conversations"] = openchat_df["items"]  
openchat_dataset = Dataset.from_pandas(openchat_df)  
###  

dataset = concatenate_datasets([gpt4_dataset, megagon_dataset, openchat_dataset])  
dataset = dataset.filter(lambda x: not any([y["value"] is None for y in x["conversations"]]))  
dataset.select_columns(["conversations"]).to_json("/workspace/llm_training/axolotl/llama3-multilingual/tagengo_openchat_megagon.json")  

axolotl版本：0.4.0

base_model: meta-llama/Meta-Llama-3-8B-Instruct  
model_type: LlamaForCausalLM  
tokenizer_type: AutoTokenizer  # PreTrainedTokenizerFast  

load_in_8bit: false  
load_in_4bit: false  
strict: false  

datasets:  
  - path: /workspace/llm_training/axolotl/llama3-multilingual/tagengo_openchat_megagon.json  
    ds_type: json # see other options below  
    type: sharegpt  
    conversation: llama-3  
dataset_prepared_path: /workspace/llm_training/axolotl/llama3-multilingual/prepared_tagengo_openchat_megagon  
val_set_size: 0.01  
output_dir: /workspace/llm_training/axolotl/llama3-multilingual/output_tagengo_openchat_megagon_8B_llama3  

sequence_len: 8192  
sample_packing: true  
pad_to_sequence_len: true  

use_wandb: true  
wandb_project: wandb_project  
wandb_entity: wandb_entity  
wandb_name: wandb_name  

gradient_accumulation_steps: 2  
micro_batch_size: 2  
num_epochs: 1  
optimizer: paged_adamw_8bit  
lr_scheduler: cosine  
learning_rate: 1e-5  

train_on_inputs: false  
group_by_length: false  
bf16: auto  
fp16:  
tf32: false  

gradient_checkpointing: true  
gradient_checkpointing_kwargs:  
  use_reentrant: false  
early_stopping_patience:  
resume_from_checkpoint:  
logging_steps: 1  
xformers_attention:  
flash_attention: true  

warmup_steps: 10  
evals_per_epoch: 5  
eval_table_size:  
saves_per_epoch: 1  
debug:  
deepspeed: /workspace/axolotl/deepspeed_configs/zero2.json  
weight_decay: 0.0  
special_tokens:  
  pad_token: <|end_of_text|>  

from fastchat.conversation import Conversation  
from fastchat.conversation import register_conv_template  
from fastchat.conversation import SeparatorStyle  

register_conv_template(  
    Conversation(  
        name="llama-3",  
        system_template="<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_message}",  
        roles=("<|start_header_id|>user<|end_header_id|>\n", "<|start_header_id|>assistant<|end_header_id|>\n"),  
        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,  
        sep="<|eot_id|>",  
        stop_token_ids=[128009],  
        stop_str="<|eot_id|>",  
    )  
)