{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "gpuType": "T4"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🏥 한국어 의료 LLM - QLoRA 파인튜닝 (T4 최적화)\n",
    "\n",
    "**실행 전 체크리스트**\n",
    "- [ ] 런타임 → 런타임 유형 변경 → **T4 GPU** 선택\n",
    "- [ ] HuggingFace 토큰 준비 (https://huggingface.co/settings/tokens)\n",
    "- [ ] Google Drive에 `medical_llm/data/jsonl/` 폴더 + JSONL 파일 업로드 완료\n",
    "\n",
    "**T4 최적화 핵심**\n",
    "- 모델: `beomi/gemma-ko-2b` (2B, T4에서 안정적)\n",
    "- 4-bit QLoRA + fp16 (bf16 ❌ T4 미지원)\n",
    "- max_seq_length: 512 (768 이상 → OOM 위험)\n",
    "- gradient_checkpointing: True\n",
    "- LoRA rank: 8 (16 이상 → 메모리 부족)\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_install"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 1] 패키지 설치\n",
    "# 처음 한 번만 실행 (런타임 재시작 후 재실행 필요)\n",
    "# ─────────────────────────────────────────────\n",
    "!pip install -q transformers==4.40.0\n",
    "!pip install -q peft==0.10.0\n",
    "!pip install -q trl==0.8.6\n",
    "!pip install -q bitsandbytes==0.43.1\n",
    "!pip install -q accelerate==0.29.3\n",
    "!pip install -q datasets==2.19.0\n",
    "!pip install -q huggingface_hub\n",
    "print('✅ 설치 완료')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_gpu_check"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 2] GPU 환경 확인\n",
    "# ─────────────────────────────────────────────\n",
    "import torch\n",
    "\n",
    "if not torch.cuda.is_available():\n",
    "    raise RuntimeError('❌ GPU가 없습니다! 런타임 → 런타임 유형 변경 → T4 GPU 선택')\n",
    "\n",
    "gpu_name = torch.cuda.get_device_name(0)\n",
    "vram_gb  = torch.cuda.get_device_properties(0).total_memory / 1e9\n",
    "print(f'✅ GPU : {gpu_name}')\n",
    "print(f'   VRAM: {vram_gb:.1f}GB')\n",
    "\n",
    "# T4는 bf16 미지원 → fp16 강제\n",
    "USE_BF16 = vram_gb >= 38   # A100이면 True\n",
    "USE_FP16 = not USE_BF16\n",
    "print(f'   정밀도: {\"bf16\" if USE_BF16 else \"fp16 (T4 최적)\"}')\n",
    "\n",
    "# VRAM 사용량 확인 함수\n",
    "def print_vram():\n",
    "    used  = torch.cuda.memory_allocated() / 1e9\n",
    "    total = torch.cuda.get_device_properties(0).total_memory / 1e9\n",
    "    print(f'   VRAM 사용: {used:.1f}GB / {total:.1f}GB ({used/total*100:.0f}%)')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_drive_mount"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 3] Google Drive 마운트\n",
    "# ─────────────────────────────────────────────\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "import os\n",
    "\n",
    "# ★ 본인 Drive 경로에 맞게 수정\n",
    "DRIVE_BASE   = '/content/drive/MyDrive/medical_llm'\n",
    "TRAIN_FILE   = f'{DRIVE_BASE}/data/jsonl/train.jsonl'\n",
    "VAL_FILE     = f'{DRIVE_BASE}/data/jsonl/val.jsonl'\n",
    "OUTPUT_DIR   = f'{DRIVE_BASE}/checkpoints'   # 체크포인트 Drive에 저장 (세션 끊겨도 안전)\n",
    "\n",
    "# 파일 존재 확인\n",
    "for f in [TRAIN_FILE, VAL_FILE]:\n",
    "    status = '✅' if os.path.exists(f) else '❌ 없음 → Drive에 파일 업로드 필요'\n",
    "    print(f'{status} {f}')\n",
    "\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "print(f'\\n체크포인트 저장 경로: {OUTPUT_DIR}')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_hf_login"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 4] HuggingFace 로그인\n",
    "# 토큰 발급: https://huggingface.co/settings/tokens\n",
    "# Colab 왼쪽 자물쇠 아이콘 → Secrets → HF_TOKEN 추가 권장\n",
    "# ─────────────────────────────────────────────\n",
    "from huggingface_hub import login\n",
    "\n",
    "try:\n",
    "    from google.colab import userdata\n",
    "    login(token=userdata.get('HF_TOKEN'), add_to_git_credential=False)\n",
    "    print('✅ HuggingFace 로그인 성공 (Secrets 사용)')\n",
    "except Exception:\n",
    "    # Secrets 미설정 시 직접 입력\n",
    "    token = input('HuggingFace 토큰을 입력하세요 (hf_...): ')\n",
    "    login(token=token)\n",
    "    print('✅ HuggingFace 로그인 성공')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_config"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 5] T4 최적화 설정값\n",
    "# ─────────────────────────────────────────────\n",
    "\n",
    "# ── 모델 선택 ─────────────────────────────────\n",
    "# T4(16GB) 권장 모델:\n",
    "#   ① beomi/gemma-ko-2b          → 2B, 가장 안정적, 빠름  ← 처음이라면 이걸로 시작\n",
    "#   ② beomi/Llama-3-Open-Ko-8B   → 8B, 성능 좋음, OOM 주의\n",
    "BASE_MODEL = 'beomi/gemma-ko-2b'   # ← 첫 실행은 2B 추천!\n",
    "\n",
    "# ── T4 VRAM 최적화 핵심 설정 ─────────────────\n",
    "MAX_SEQ_LENGTH  = 512    # T4에서 512가 안전선. 768↑ → OOM 위험\n",
    "LORA_R          = 8      # rank 낮을수록 메모리↓ (처음엔 8로 시작)\n",
    "LORA_ALPHA      = 16     # 보통 rank * 2\n",
    "LORA_DROPOUT    = 0.05\n",
    "\n",
    "# ── 학습 하이퍼파라미터 ───────────────────────\n",
    "BATCH_SIZE      = 1      # T4는 1 고정 (2 이상 → OOM)\n",
    "GRAD_ACCUM      = 8      # 실질 배치 크기 = 1 * 8 = 8\n",
    "LEARNING_RATE   = 2e-4\n",
    "NUM_EPOCHS      = 3\n",
    "WARMUP_RATIO    = 0.05\n",
    "\n",
    "# ── 저장 설정 ─────────────────────────────────\n",
    "SAVE_STEPS      = 50     # T4는 끊김 잦음 → 50스텝마다 저장\n",
    "LOGGING_STEPS   = 10\n",
    "\n",
    "# ── HuggingFace Hub ───────────────────────────\n",
    "HUB_MODEL_ID    = 'your-username/medical-llm-ko'  # ★ 본인 아이디로 변경\n",
    "PUSH_TO_HUB     = True\n",
    "\n",
    "print('설정값 확인:')\n",
    "print(f'  모델           : {BASE_MODEL}')\n",
    "print(f'  최대 시퀀스 길이: {MAX_SEQ_LENGTH} tokens')\n",
    "print(f'  LoRA rank      : {LORA_R}')\n",
    "print(f'  배치 크기      : {BATCH_SIZE} × grad_accum {GRAD_ACCUM} = 실질 {BATCH_SIZE*GRAD_ACCUM}')\n",
    "print(f'  학습 에폭      : {NUM_EPOCHS}')\n",
    "print(f'  체크포인트 저장: 매 {SAVE_STEPS} 스텝')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_model_load"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 6] 모델 로드 (4-bit QLoRA)\n",
    "# 약 3~5분 소요 (첫 다운로드)\n",
    "# ─────────────────────────────────────────────\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n",
    "from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
    "\n",
    "# T4 최적화: fp16 사용 (T4는 bf16 미지원)\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type='nf4',\n",
    "    bnb_4bit_compute_dtype=torch.float16,   # ← T4는 float16 (bf16 ❌)\n",
    "    bnb_4bit_use_double_quant=True,         # 이중 양자화로 추가 메모리 절약\n",
    ")\n",
    "\n",
    "print(f'[모델 로드] {BASE_MODEL}')\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)\n",
    "# pad 토큰 없으면 추가\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "tokenizer.padding_side = 'right'\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    BASE_MODEL,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map='auto',\n",
    "    trust_remote_code=True,\n",
    "    torch_dtype=torch.float16,   # ← T4는 float16\n",
    ")\n",
    "\n",
    "model = prepare_model_for_kbit_training(\n",
    "    model,\n",
    "    use_gradient_checkpointing=True   # 메모리 절약 핵심\n",
    ")\n",
    "\n",
    "print('✅ 모델 로드 완료')\n",
    "print(f'   전체 파라미터: {sum(p.numel() for p in model.parameters()):,}')\n",
    "print_vram()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_lora"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 7] LoRA 어댑터 적용\n",
    "# ─────────────────────────────────────────────\n",
    "lora_config = LoraConfig(\n",
    "    r=LORA_R,\n",
    "    lora_alpha=LORA_ALPHA,\n",
    "    # 어텐션 레이어에만 적용 (FFN 제외 → T4 메모리 절약)\n",
    "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj'],\n",
    "    lora_dropout=LORA_DROPOUT,\n",
    "    bias='none',\n",
    "    task_type='CAUSAL_LM',\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, lora_config)\n",
    "model.print_trainable_parameters()\n",
    "# 예: trainable params: 6,815,744 || all params: 2,512,197,632 || trainable%: 0.27\n",
    "# → 전체의 0.27%만 학습 = LoRA의 핵심!\n",
    "\n",
    "print_vram()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_dataset"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 8] 데이터셋 로드 + 샘플 확인\n",
    "# ─────────────────────────────────────────────\n",
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\n",
    "    'json',\n",
    "    data_files={'train': TRAIN_FILE, 'validation': VAL_FILE},\n",
    ")\n",
    "\n",
    "print(f'학습 데이터 : {len(dataset[\"train\"]):,}개')\n",
    "print(f'검증 데이터 : {len(dataset[\"validation\"]):,}개')\n",
    "\n",
    "# 샘플 미리보기\n",
    "sample = dataset['train'][0]\n",
    "print('\\n--- 샘플 미리보기 (앞 400자) ---')\n",
    "print(sample['text'][:400])\n",
    "\n",
    "# 토큰 길이 분포 확인 (OOM 예방)\n",
    "lengths = [len(tokenizer.encode(d['text'])) for d in dataset['train'].select(range(min(200, len(dataset['train']))))]\n",
    "import statistics\n",
    "print(f'\\n토큰 길이 통계 (샘플 200개 기준):')\n",
    "print(f'  평균: {statistics.mean(lengths):.0f} tokens')\n",
    "print(f'  최대: {max(lengths)} tokens')\n",
    "print(f'  max_seq_length 초과 비율: {sum(1 for l in lengths if l > MAX_SEQ_LENGTH)/len(lengths)*100:.1f}%')\n",
    "print('  → 초과 비율이 높으면 MAX_SEQ_LENGTH를 늘리거나 03번에서 필터를 강화하세요')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_train"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 9] 학습 실행\n",
    "# T4 기준 예상 시간:\n",
    "#   2B 모델 1,000샘플 3epoch → 약 20~30분\n",
    "#   2B 모델 10,000샘플 3epoch → 약 3~4시간\n",
    "# ─────────────────────────────────────────────\n",
    "from transformers import TrainingArguments\n",
    "from trl import SFTTrainer\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=OUTPUT_DIR,\n",
    "    num_train_epochs=NUM_EPOCHS,\n",
    "    per_device_train_batch_size=BATCH_SIZE,\n",
    "    gradient_accumulation_steps=GRAD_ACCUM,\n",
    "    gradient_checkpointing=True,       # T4 필수 (메모리 20~30% 절약)\n",
    "    learning_rate=LEARNING_RATE,\n",
    "    lr_scheduler_type='cosine',\n",
    "    warmup_ratio=WARMUP_RATIO,\n",
    "    fp16=USE_FP16,                     # T4: True\n",
    "    bf16=USE_BF16,                     # T4: False\n",
    "    logging_steps=LOGGING_STEPS,\n",
    "    save_steps=SAVE_STEPS,             # Drive에 자주 저장 (세션 끊김 대비)\n",
    "    save_total_limit=2,                # 최근 2개만 보관 (Drive 용량 절약)\n",
    "    evaluation_strategy='steps',\n",
    "    eval_steps=SAVE_STEPS,\n",
    "    load_best_model_at_end=True,\n",
    "    optim='paged_adamw_8bit',          # T4 메모리 최적화 옵티마이저\n",
    "    dataloader_pin_memory=False,       # T4에서 pin_memory=True → 오히려 느림\n",
    "    report_to='none',\n",
    "    push_to_hub=False,                 # 학습 중 Hub 업로드 끄기 (학습 후 수동 업로드)\n",
    ")\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    train_dataset=dataset['train'],\n",
    "    eval_dataset=dataset['validation'],\n",
    "    dataset_text_field='text',\n",
    "    max_seq_length=MAX_SEQ_LENGTH,\n",
    "    args=training_args,\n",
    "    packing=False,\n",
    ")\n",
    "\n",
    "train_count = len(dataset['train'])\n",
    "steps_per_epoch = train_count // (BATCH_SIZE * GRAD_ACCUM)\n",
    "total_steps = steps_per_epoch * NUM_EPOCHS\n",
    "print(f'🚀 학습 시작!')\n",
    "print(f'   총 스텝: {total_steps:,}')\n",
    "print(f'   체크포인트: 매 {SAVE_STEPS}스텝 → {OUTPUT_DIR}')\n",
    "print(f'   ※ Colab이 끊겨도 Drive 체크포인트에서 이어서 가능')\n",
    "\n",
    "trainer.train()\n",
    "print('\\n✅ 학습 완료!')\n",
    "print_vram()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_resume"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 10] 세션 끊긴 후 이어서 학습하기\n",
    "# (정상 완료했으면 이 셀은 건너뛰세요)\n",
    "# ─────────────────────────────────────────────\n",
    "import os, glob\n",
    "\n",
    "# 가장 최근 체크포인트 자동 탐색\n",
    "checkpoints = sorted(\n",
    "    glob.glob(f'{OUTPUT_DIR}/checkpoint-*'),\n",
    "    key=lambda x: int(x.split('-')[-1])\n",
    ")\n",
    "\n",
    "if checkpoints:\n",
    "    latest_ckpt = checkpoints[-1]\n",
    "    print(f'최근 체크포인트: {latest_ckpt}')\n",
    "    print('이어서 학습하려면 아래 코드 실행:')\n",
    "    print(f'  trainer.train(resume_from_checkpoint=\"{latest_ckpt}\")')\n",
    "else:\n",
    "    print('체크포인트 없음. Cell 9부터 다시 실행하세요.')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_save"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 11] 모델 저장 + HuggingFace Hub 업로드\n",
    "# ─────────────────────────────────────────────\n",
    "FINAL_MODEL_DIR = f'{OUTPUT_DIR}/final'\n",
    "\n",
    "# LoRA 어댑터만 저장 (약 50~150MB)\n",
    "trainer.save_model(FINAL_MODEL_DIR)\n",
    "tokenizer.save_pretrained(FINAL_MODEL_DIR)\n",
    "print(f'✅ 로컬 저장 완료: {FINAL_MODEL_DIR}')\n",
    "\n",
    "# HuggingFace Hub 업로드\n",
    "if PUSH_TO_HUB:\n",
    "    if 'your-username' in HUB_MODEL_ID:\n",
    "        print('⚠️  HUB_MODEL_ID를 본인 HuggingFace 아이디로 변경하세요!')\n",
    "        print('   예: HUB_MODEL_ID = \"홍길동/medical-llm-ko\"')\n",
    "    else:\n",
    "        trainer.push_to_hub(HUB_MODEL_ID)\n",
    "        print(f'✅ Hub 업로드 완료!')\n",
    "        print(f'   https://huggingface.co/{HUB_MODEL_ID}')"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "cell_inference"
   },
   "source": [
    "# ─────────────────────────────────────────────\n",
    "# [Cell 12] 추론 테스트 (파인튜닝 전후 비교)\n",
    "# ─────────────────────────────────────────────\n",
    "SYSTEM_PROMPT = '당신은 환자의 증상과 질문을 듣고 의학적 정보를 제공하는 의료 상담 AI입니다. 정확한 의료 정보를 바탕으로 답변하고 전문의 상담을 권장하세요.'\n",
    "\n",
    "def make_prompt(question: str) -> str:\n",
    "    \"\"\"LLaMA3 포맷 프롬프트 생성\"\"\"\n",
    "    return (\n",
    "        '<|begin_of_text|>'\n",
    "        '<|start_header_id|>system<|end_header_id|>\\n\\n'\n",
    "        f'{SYSTEM_PROMPT}'\n",
    "        '<|eot_id|>'\n",
    "        '<|start_header_id|>user<|end_header_id|>\\n\\n'\n",
    "        f'{question}'\n",
    "        '<|eot_id|>'\n",
    "        '<|start_header_id|>assistant<|end_header_id|>\\n\\n'\n",
    "    )\n",
    "\n",
    "def generate(question: str, max_new_tokens: int = 256) -> str:\n",
    "    prompt = make_prompt(question)\n",
    "    inputs = tokenizer(prompt, return_tensors='pt').to(model.device)\n",
    "    with torch.no_grad():\n",
    "        outputs = model.generate(\n",
    "            **inputs,\n",
    "            max_new_tokens=max_new_tokens,\n",
    "            temperature=0.7,\n",
    "            do_sample=True,\n",
    "            top_p=0.9,\n",
    "            repetition_penalty=1.1,\n",
    "            eos_token_id=tokenizer.eos_token_id,\n",
    "            pad_token_id=tokenizer.pad_token_id,\n",
    "        )\n",
    "    return tokenizer.decode(\n",
    "        outputs[0][inputs['input_ids'].shape[1]:],\n",
    "        skip_special_tokens=True\n",
    "    ).strip()\n",
    "\n",
    "# 테스트 질문\n",
    "test_questions = [\n",
    "    '두통이 3일째 계속되고 구역질도 나는데 어떻게 해야 하나요?',\n",
    "    '당뇨 진단을 받았는데 식단 관리는 어떻게 해야 하나요?',\n",
    "    '혈압이 150/95 정도 나오는데 위험한가요?',\n",
    "]\n",
    "\n",
    "print('=' * 60)\n",
    "print('파인튜닝 모델 추론 테스트')\n",
    "print('=' * 60)\n",
    "\n",
    "for q in test_questions:\n",
    "    print(f'\\n질문: {q}')\n",
    "    ans = generate(q)\n",
    "    print(f'답변: {ans}')\n",
    "    print('-' * 40)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ✅ 완료 후 다음 단계\n",
    "\n",
    "1. **평가** : `07_evaluation.py` 로 ROUGE/BLEU 파인튜닝 전후 비교\n",
    "2. **RAG 연결** : `05_rag_indexer.py` → `06_rag_chain.py`\n",
    "3. **API 배포** : `08_api.py` (FastAPI)\n",
    "\n",
    "### OOM 에러가 났다면\n",
    "```\n",
    "MAX_SEQ_LENGTH  줄이기  512 → 384\n",
    "LORA_R          줄이기  8 → 4\n",
    "target_modules  줄이기  q_proj, v_proj 만\n",
    "BASE_MODEL      교체    beomi/gemma-ko-2b (이미 최소)\n",
    "```\n"
   ]
  }
 ]
}