import os
import logging

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM

# -------------------------
# Config
# -------------------------

BASE_DIR = "/home/jovyan"

HF_MODEL_8B_ID = "kakaocorp/kanana-1.5-8b-base"
HF_MODEL_30B_ID = "kakaocorp/kanana-2-30b-a3b-base"

HF_DATASET_ID = "HuggingFaceH4/ultrachat_200k"


MODEL_8B_DIR = f"{BASE_DIR}/kanana-8b"
MODEL_30B_DIR = f"{BASE_DIR}/kanana-30b"

TOKENIZER_DIR = MODEL_8B_DIR
RAW_DATASET_DIR = f"{BASE_DIR}/dataset/ultrachat_raw"
TOKENIZED_DATASET_DIR = f"{BASE_DIR}/dataset/ultrachat_tokenized"

MAX_LENGTH = 2048
NUM_PROC = os.cpu_count()


# -------------------------
# Logging
# -------------------------

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s"
)


# -------------------------
# Step 1. Download tokenizer
# -------------------------

logging.info("Downloading tokenizer from HuggingFace")

tokenizer = AutoTokenizer.from_pretrained(
    HF_MODEL_8B_ID,
    padding_side="left",
    trust_remote_code=True
)

tokenizer.pad_token = tokenizer.eos_token

logging.info(f"Saving tokenizer to {TOKENIZER_DIR}")

os.makedirs(TOKENIZER_DIR, exist_ok=True)

tokenizer.save_pretrained(TOKENIZER_DIR)

# -------------------------
# Step 2. Download model
# -------------------------

logging.info("Downloading models from HuggingFace")

model = AutoModelForCausalLM.from_pretrained(
    HF_MODEL_8B_ID,
    trust_remote_code=True,
    torch_dtype="auto"
)

logging.info(f"Saving model to {MODEL_8B_DIR}")

model.save_pretrained(
    MODEL_8B_DIR,
    safe_serialization=True
)

del model

model = AutoModelForCausalLM.from_pretrained(
    HF_MODEL_30B_ID,
    trust_remote_code=True,
    torch_dtype="auto"
)

logging.info(f"Saving model to {HF_MODEL_30B_ID}")

model.save_pretrained(
    MODEL_30B_DIR,
    safe_serialization=True
)

del model

logging.info("Model saved")


# -------------------------
# Step 3. Download dataset
# -------------------------

logging.info("Downloading ultrachat_200k dataset")

dataset = load_dataset(
    HF_DATASET_ID,
    split="train_sft"
)

logging.info(f"Dataset size: {len(dataset)}")

logging.info(f"Saving raw dataset to {RAW_DATASET_DIR}")

os.makedirs(RAW_DATASET_DIR, exist_ok=True)

dataset.save_to_disk(RAW_DATASET_DIR)

# -------------------------
# Step 4. Tokenize dataset
# -------------------------

logging.info("Tokenizing dataset")

def preprocess(example):

    text = tokenizer.apply_chat_template(
        example["messages"],
        tokenize=False,
        add_generation_prompt=False
    )

    tokens = tokenizer(
        text,
        truncation=True,
        max_length=MAX_LENGTH,
        padding="max_length"
    )

    tokens["labels"] = tokens["input_ids"].copy()

    return tokens


tokenized_dataset = dataset.map(
    preprocess,
    num_proc=NUM_PROC,
    remove_columns=dataset.column_names,
    desc="Tokenizing ultrachat"
)

# -------------------------
# Step 5. Save tokenized dataset
# -------------------------

logging.info(f"Saving tokenized dataset to {TOKENIZED_DATASET_DIR}")

os.makedirs(TOKENIZED_DATASET_DIR, exist_ok=True)

tokenized_dataset.save_to_disk(TOKENIZED_DATASET_DIR)

logging.info("DONE")