import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM, 
    TrainingArguments, 
    Trainer, 
    DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model

# 1. Tải Dataset
dataset = load_dataset('vntc/wiki-mini-corpus')

# 2. Chia Dataset thành train và validation
split_dataset = dataset['train'].train_test_split(test_size=0.1, seed=42)
train_dataset = split_dataset['train']
validation_dataset = split_dataset['test']

# 3. Tiền Xử Lý Văn Bản
def preprocess_function(examples):
    passages = [passage.lower().strip() for passage in examples['passage']]
    return {'passage': passages}

processed_train = train_dataset.map(preprocess_function, batched=True, remove_columns=['id', 'metadata'])
processed_validation = validation_dataset.map(preprocess_function, batched=True, remove_columns=['id', 'metadata'])

# 4. Tokenization
model_name = "meta-llama/Llama-3.2-3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Đảm bảo tokenizer có pad_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

def tokenize_function(examples):
    return tokenizer(
        examples['passage'],
        padding='max_length',
        truncation=True,
        max_length=512,
    )

tokenized_train = processed_train.map(tokenize_function, batched=True)
tokenized_validation = processed_validation.map(tokenize_function, batched=True)

# 5. Thêm trường 'labels'
def add_labels(examples):
    examples['labels'] = examples['input_ids'].copy()
    return examples

tokenized_train = tokenized_train.map(add_labels, batched=True)
tokenized_validation = tokenized_validation.map(add_labels, batched=True)

# 6. Loại bỏ các cột không cần thiết
tokenized_train = tokenized_train.remove_columns(['passage'])
tokenized_validation = tokenized_validation.remove_columns(['passage'])

# 7. Định dạng dữ liệu cho PyTorch
tokenized_train.set_format('torch')
tokenized_validation.set_format('torch')

# 8. Tạo DatasetDict
final_dataset = {
    'train': tokenized_train,
    'validation': tokenized_validation
}

# 9. Tải và Cấu Hình Mô Hình với LoRA
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16,
    load_in_8bit=False
)

lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
    lora_dropout=0.1,
    bias="none",
)

model = get_peft_model(model, lora_config)
print(model)

# 10. Cấu Hình TrainingArguments
training_args = TrainingArguments(
    output_dir="./lora_finetuned_llama3.2",
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=8,
    num_train_epochs=3,
    learning_rate=3e-4,
    weight_decay=0.01,
    logging_steps=100,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,
    fp16=True,
    report_to="none",
)

# 11. Data Collator Sử Dụng DataCollatorForLanguageModeling
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, 
    mlm=False,  # Vì bạn đang thực hiện Causal LM
    pad_to_multiple_of=8
)

# 12. Tạo Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=final_dataset['train'],
    eval_dataset=final_dataset['validation'],
    tokenizer=tokenizer,
    data_collator=data_collator,
)

# 13. Bắt Đầu Huấn Luyện
trainer.train()