import numpy as np
import pandas as pd
import torch
from datasets import Dataset, load_dataset, concatenate_datasets

# 修正：必要なクラスを正しくインポート
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer, 
    Trainer, 
    TrainingArguments,
    DataCollatorForLanguageModeling
)


model_name = "./hourai3-85m-base-ja/"
max_seq_length = 256

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,  # 16bitでロードしてVRAMを節約
    device_map="auto"            # 自動でGPUに配置
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# =================================================================
# 2. データセットの準備と整形
# =================================================================
# データセットの読み込み
# Miwa-Keita/zenz-v2.5-dataset
dataset = load_dataset(
    "json",
    data_files="./dataset/zenz-v2.5-dataset/train_wikipedia.jsonl",
    split="train[:40000]"
) # 60k
 
def format_dataset(sample):
    input_text = sample["input"]
    output_text = sample["output"]    
    left_context = sample.get("left_context") or ""

    prefix = f"\uEE02{left_context}" if left_context else ""
    
    eos = tokenizer.eos_token if tokenizer.eos_token else "</s>"
    text = f"{prefix}\uEE00{input_text}\uEE01{output_text}{eos}"
    
    return {"text": text}

dataset = dataset.map(format_dataset)
print(dataset[0])

# テキストデータをトークン化する関数
def tokenize_function(examples):       
    return tokenizer(
        examples["text"], 
        truncation=True, 
        max_length=max_seq_length,
        padding=False # コレーター側で動的に処理するためここではFalse
    )
dataset = dataset.select_columns(["text"])

tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)

# 2. 完全一致で \uEE01 を探す DataCollator
class ZenzDataCollatorForCompletionOnly(DataCollatorForLanguageModeling):
    def __init__(self, output_start_token="\uEE01", *args, **kwargs):
        super().__init__(*args, **kwargs)
        # \uEE01 のトークンID「配列」を取得
        self.target_ids = torch.tensor(
            self.tokenizer.encode(output_start_token, add_special_tokens=False),
            dtype=torch.long
        )

    def torch_call(self, examples):
        batch = super().torch_call(examples)
        labels = batch["labels"].clone()
        input_ids = batch["input_ids"]

        target_len = len(self.target_ids)

        for i in range(len(labels)):
            found_idx = -1
            # サブシーケンス（\uEE01 のトークン列）を完全一致検索
            for j in range(len(input_ids[i]) - target_len + 1):
                if torch.equal(input_ids[i][j : j + target_len], self.target_ids):
                    found_idx = j
                    break

            if found_idx != -1:
                # \uEE01 の最後のトークンの位置までを -100 でマスクする
                # これにより、\uEE01 の直後の「漢字」から Loss が計算される
                mask_until = found_idx + target_len
                labels[i][:mask_until] = -100
            else:
                # 見つからない場合は全部無視
                labels[i][:] = -100

        batch["labels"] = labels
        return batch
        
data_collator = ZenzDataCollatorForCompletionOnly(
    tokenizer=tokenizer,
    mlm=False,
    output_start_token="\uEE01"
)

# デバッグ用確認（collatorの動作テスト）
test_batch = [tokenized_dataset[0]]
processed_batch = data_collator(test_batch)

print("Input IDs:", processed_batch["input_ids"][0])
print("Labels   :", processed_batch["labels"][0])


training_args = TrainingArguments(
    output_dir="outputs",
    per_device_train_batch_size=16,
    gradient_accumulation_steps=1,  # 1だと勾配が不安定になりやすいため、4程度に増やして実質バッチサイズを確保
    warmup_steps=50,
    max_steps=2500,
    learning_rate=2e-4,
    logging_steps=5,
    optim="adamw_8bit",             # VRAM節約のために8bit AdamWを維持
    weight_decay=0.001,
    lr_scheduler_type="linear",
    seed=3407,
    report_to="none",
    fp16=False,                     # bfloat16推奨
    bf16=True,                      # Ampere世代以降（RTX 30xx/40xx, A100等）ならTrue
    gradient_checkpointing=False,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    remove_unused_columns=False,
    dataloader_num_workers=4,
    

)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator,
)

# 学習開始
trainer.train(resume_from_checkpoint=False)