from unsloth import FastLanguageModel, FastModel
from trl import SFTTrainer, SFTConfig
import torch

# 🎯 사전 양자화된 모델 로드
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/gemma-2-2b-it",  # 예시: Gemma 2B 모델
    max_seq_length = 2048,                  # 최대 시퀀스 길이
    dtype = None,                           # 자동 감지
    load_in_4bit = True,                    # 4비트 양자화로 메모리 절약!
)

# 🔧 PEFT (Parameter Efficient Fine-Tuning) 설정
model = FastLanguageModel.get_peft_model(
    model,
    r = 16,                    # LoRA attention dimension
    target_modules = [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_alpha = 16,           # LoRA scaling parameter
    lora_dropout = 0,          # LoRA dropout (0으로 설정 권장)
    bias = "none",             # bias 설정
    use_gradient_checkpointing = "unsloth",  # 메모리 효율성을 위해!
    random_state = 3407,       # 재현 가능한 결과를 위한 시드
    use_rslora = False,        # Rank stabilized LoRA 사용 여부
    loftq_config = None,       # LoftQ quantization config
)

from datasets import Dataset

# 📚 예시 대화 데이터셋
conversations = [
    {
        "input": "안녕하세요! 파이썬에 대해 알려주세요.",
        "output": "안녕하세요! 파이썬은 직관적이고 배우기 쉬운 프로그래밍 언어입니다. 데이터 분석, 웹 개발, AI 등 다양한 분야에서 활용됩니다!"
    },
    {
        "input": "머신러닝이란 무엇인가요?",
        "output": "머신러닝은 컴퓨터가 명시적으로 프로그래밍되지 않고도 데이터에서 패턴을 학습하여 예측이나 결정을 내리는 기술입니다."
    }
    # 더 많은 데이터 추가...
]

# 데이터셋 생성
dataset = Dataset.from_list(conversations)

# 🎨 프롬프트 포맷팅 함수
alpaca_prompt = """Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{}

### Response:
{}"""

def formatting_prompts_func(examples):
    inputs       = examples["input"]
    outputs      = examples["output"]
    texts = []
    for input, output in zip(inputs, outputs):
        text = alpaca_prompt.format(input, output) + tokenizer.eos_token
        texts.append(text)
    return { "text" : texts, }

dataset = dataset.map(formatting_prompts_func, batched = True,)

# 🚀 SFT (Supervised Fine-Tuning) 트레이너 설정
trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = 2048,
    dataset_num_proc = 2,
    
    args = SFTConfig(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        warmup_steps = 5,
        max_steps = 60,
        learning_rate = 2e-4,
        fp16 = not torch.cuda.is_bf16_supported(),
        bf16 = torch.cuda.is_bf16_supported(),
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 3407,
        output_dir = "outputs",
    ),
)

# 🎉 학습 시작!
trainer.train()

# 🎯 다양한 양자화 옵션들

# 4비트 양자화 (최대 메모리 절약)
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/llama-3-8b-bnb-4bit",
    max_seq_length = 2048,
    load_in_4bit = True,
)

# 8비트 양자화 (균형잡힌 선택)
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/llama-3-8b",
    max_seq_length = 2048,
    load_in_8bit = True,
)

# 16비트 (높은 정확도)
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/llama-3-8b",
    max_seq_length = 2048,
    dtype = torch.float16,
)

# 🔧 더 세밀한 PEFT 설정
model = FastLanguageModel.get_peft_model(
    model,
    r = 32,                    # 더 높은 rank = 더 많은 매개변수
    target_modules = [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
        "embed_tokens", "lm_head",  # 임베딩 레이어도 포함
    ],
    lora_alpha = 32,
    lora_dropout = 0.1,        # 과적합 방지를 위한 드롭아웃
    bias = "lora_only",        # LoRA에만 bias 적용
    use_rslora = True,         # Rank Stabilized LoRA 사용
    modules_to_save = ["embed_tokens", "lm_head"],  # 저장할 모듈 지정
)

# 🤖 고객 서비스 챗봇 학습 데이터
chatbot_data = [
    {
        "input": "환불은 어떻게 요청하나요?",
        "output": "환불 요청은 마이페이지 > 주문내역에서 해당 상품을 선택하여 환불 신청하실 수 있습니다. 영업일 기준 3-5일 내 처리됩니다."
    },
    # 더 많은 Q&A 데이터...
]

# 🏥 의료 분야 특화 모델
medical_data = [
    {
        "input": "두통의 주요 원인은 무엇인가요?",
        "output": "두통의 주요 원인으로는 스트레스, 수면 부족, 탈수, 눈의 피로, 근육 긴장 등이 있습니다. 지속적인 두통의 경우 전문의 상담을 권합니다."
    },
    # 의료 전문 데이터...
]

# 💻 프로그래밍 도우미 모델
coding_data = [
    {
        "input": "파이썬에서 리스트 컴프리헨션을 설명해주세요",
        "output": "리스트 컴프리헨션은 간결한 방식으로 리스트를 생성하는 방법입니다. 예: [x**2 for x in range(10) if x%2==0]"
    },
    # 코딩 관련 데이터...
]

# 💡 메모리 절약 팁들

# 1. Gradient Checkpointing 활성화
model = FastLanguageModel.get_peft_model(
    model,
    use_gradient_checkpointing = "unsloth",  # 메모리 절약!
)

# 2. 배치 사이즈 줄이기
args = SFTConfig(
    per_device_train_batch_size = 1,  # 기본값보다 작게
    gradient_accumulation_steps = 8,   # 대신 accumulation 늘리기
)

# 3. 시퀀스 길이 조정
max_seq_length = 1024  # 2048에서 1024로 줄이기

# 🚀 학습 속도 최적화

# 1. 적절한 learning rate 설정
learning_rate = 2e-4  # 일반적으로 2e-4가 좋은 시작점

# 2. Warmup steps 조정
warmup_steps = max_steps // 10  # 전체 step의 10%

# 3. 스케줄러 선택
lr_scheduler_type = "cosine"  # cosine 스케줄러 사용

# 🔥 여러 GPU 활용하기
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"  # 4개 GPU 사용

# DDP (Distributed Data Parallel) 설정
args = SFTConfig(
    per_device_train_batch_size = 1,
    gradient_accumulation_steps = 4,
    dataloader_num_workers = 4,
    ddp_find_unused_parameters = False,
)

# 🎯 맞춤형 손실 함수
class CustomTrainer(SFTTrainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.get("logits")
        
        # 커스텀 손실 계산
        loss_fct = torch.nn.CrossEntropyLoss(ignore_index=-100)
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), 
                       shift_labels.view(-1))
        
        return (loss, outputs) if return_outputs else loss

# 🎉 학습 완료 후 모델 저장

# LoRA 어댑터만 저장 (용량 절약)
model.save_pretrained("my-awesome-model-lora")
tokenizer.save_pretrained("my-awesome-model-lora")

# 전체 모델로 병합하여 저장
model = FastLanguageModel.for_inference(model)  # 추론 모드로 전환
model.save_pretrained_merged("my-awesome-model-merged", 
                           tokenizer, 
                           save_method = "merged_16bit")

# 💫 학습된 모델로 추론하기

FastLanguageModel.for_inference(model)  # 추론 속도 2배 향상!

inputs = tokenizer(
    [
        alpaca_prompt.format(
            "파이썬으로 간단한 웹 크롤러를 만드는 방법을 알려주세요",
            ""
        )
    ], return_tensors = "pt").to("cuda")

outputs = model.generate(**inputs, 
                        max_new_tokens = 256, 
                        temperature = 0.7,
                        do_sample = True)

response = tokenizer.batch_decode(outputs)
print(response[0])