train.py

#!/usr/bin/env python3
"""
DeepSeek-Coder 6.7B QLoRA 파인튜닝 스크립트 (완전 호환성 보장)
- 모든 TRL/Transformers 버전 지원
- T4 GPU 최적화
- save_steps: 10 지원
"""

import os
import sys
import json
import yaml
import logging
import argparse
import tempfile
import inspect
from datetime import datetime
from typing import Dict, Any, Optional

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

# TRL 임포트 (버전별 처리)
try:
    from trl import SFTTrainer
    TRL_AVAILABLE = True
except ImportError:
    TRL_AVAILABLE = False
    print("⚠️ TRL 없음 - 기본 Trainer 사용")

# 로깅 설정
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)s | %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('training.log', mode='a')
    ]
)
log = logging.getLogger(__name__)

class DeepSeekTrainer:
    """DeepSeek-Coder 6.7B QLoRA 파인튜닝 클래스 (완전 호환성)"""
    
    def __init__(self, config_path: str = "config.yaml", mode: str = "prompt"):
        self.mode = mode
        self.cfg = self.load_config(config_path)
        self.model = None
        self.tok = None
        self.train_ds = None
        self.eval_ds = None
        
        # T4 환경 최적화 설정 적용
        self.optimize_for_t4()
        
        log.info(f"🚀 DeepSeek-Coder 6.7B 파인튜닝 초기화 완료 (모드: {mode})")
    
    def get_int_config(self, key, default=0):
        """config에서 값을 가져와 int로 변환하는 헬퍼 함수"""
        value = self.cfg.get(key, default)
        if isinstance(value, str):
            try:
                value = int(value)
            except ValueError:
                log.warning(f"⚠️ 경고: {key} 값 '{value}'을 정수로 변환할 수 없습니다. 기본값 {default} 사용")
                value = default
        return value
    
    def get_float_config(self, key, default=0.0):
        """config에서 값을 가져와 float로 변환하는 헬퍼 함수"""
        value = self.cfg.get(key, default)
        if isinstance(value, str):
            try:
                value = float(value)
            except ValueError:
                log.warning(f"⚠️ 경고: {key} 값 '{value}'을 실수로 변환할 수 없습니다. 기본값 {default} 사용")
                value = default
        return value
        
    def optimize_for_t4(self):
        """T4 GPU 환경에 맞는 최적화 설정"""
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            compute_cap = torch.cuda.get_device_capability(0)
            
            log.info(f"🖥️ GPU: {gpu_name}")
            log.info(f"🔧 Compute Capability: {compute_cap[0]}.{compute_cap[1]}")
            
            # T4는 Compute Capability 7.5이므로 bfloat16 미지원
            if compute_cap[0] < 8:
                log.info("⚠️ bfloat16 미지원 GPU - float16 사용")
                self.cfg["bnb_4bit_compute_dtype"] = "float16"
                self.cfg["torch_dtype"] = "float16"
                self.cfg["fp16"] = True
                self.cfg["bf16"] = False
            
            # T4 메모리 제한 고려 (15GB)
            if "T4" in gpu_name:
                log.info("🔧 T4 환경 최적화 적용")
                self.cfg["batch_size"] = 1
                self.cfg["grad_acc"] = max(8, self.get_int_config("grad_acc", 4))
                self.cfg["max_length"] = min(1024, self.get_int_config("max_length", 2048))
    
    @staticmethod
    def default_cfg():
        """T4 환경 최적화된 기본 설정"""
        return {
            # 모델 설정
            "model_name": "deepseek-ai/deepseek-coder-6.7b-instruct",
            "torch_dtype": "float16",
            
            # 양자화 설정 (T4 최적화)
            "load_in_4bit": True,
            "bnb_4bit_compute_dtype": "float16",
            "bnb_4bit_quant_type": "nf4",
            "bnb_4bit_use_double_quant": True,
            
            # LoRA 설정
            "lora_r": 16,
            "lora_alpha": 32,
            "lora_dropout": 0.1,
            "lora_target_modules": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
            
            # 학습 설정 (T4 메모리 제한 고려)
            "batch_size": 1,
            "grad_acc": 8,
            "learning_rate": 2e-4,
            "num_epochs": 3,
            "max_length": 1024,
            "warmup_ratio": 0.1,
            "lr_scheduler": "cosine",
            "weight_decay": 0.01,
            
            # 최적화 설정
            "fp16": True,
            "bf16": False,
            "gradient_checkpointing": True,
            "dataloader_num_workers": 4,
            "remove_unused_columns": False,
            
            # 저장 설정 (save_steps: 10)
            "save_steps": 10,
            "eval_steps": 50,
            "logging_steps": 5,
            "save_total_limit": 5,
            
            # 데이터 설정
            "data_path": "/home/ubuntu/deepseek-coder/data/train.jsonl"
        }
    
    def load_config(self, config_path: str) -> Dict[str, Any]:
        """설정 파일 로딩"""
        default_config = self.default_cfg()
        
        if os.path.exists(config_path):
            try:
                with open(config_path, 'r', encoding='utf-8') as f:
                    user_config = yaml.safe_load(f)
                default_config.update(user_config)
                log.info(f"✅ 설정 파일 로딩: {config_path}")
            except Exception as e:
                log.warning(f"⚠️ 설정 파일 로딩 실패, 기본값 사용: {e}")
        else:
            log.info("📝 기본 설정 사용")
        
        return default_config
    
    def load_model(self):
        """T4 최적화된 모델 및 토크나이저 로딩"""
        log.info("🔄 모델 로딩 시작")
        log.info("▶️  베이스 모델 4-bit 로드")
        
        # 양자화 설정
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=self.cfg["load_in_4bit"],
            bnb_4bit_compute_dtype=getattr(torch, self.cfg["bnb_4bit_compute_dtype"]),
            bnb_4bit_quant_type=self.cfg["bnb_4bit_quant_type"],
            bnb_4bit_use_double_quant=self.cfg["bnb_4bit_use_double_quant"],
        )
        
        # 토크나이저 로딩
        self.tok = AutoTokenizer.from_pretrained(
            self.cfg["model_name"],
            trust_remote_code=True,
            padding_side="right"
        )
        
        # 패딩 토큰 설정
        if self.tok.pad_token is None:
            self.tok.pad_token = self.tok.eos_token
            self.tok.pad_token_id = self.tok.eos_token_id
        
        # 모델 로딩
        self.model = AutoModelForCausalLM.from_pretrained(
            self.cfg["model_name"],
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True,
            torch_dtype=getattr(torch, self.cfg["torch_dtype"]),
            use_cache=False
        )
        
        # k-bit 학습을 위한 모델 준비
        self.model = prepare_model_for_kbit_training(
            self.model, 
            use_gradient_checkpointing=self.cfg["gradient_checkpointing"]
        )
        log.info("✅ 그래디언트 체크포인팅으로 k-bit 학습을 위한 모델 준비 완료.")
        
        # LoRA 설정 및 적용
        lora_config = LoraConfig(
            r=self.cfg["lora_r"],
            lora_alpha=self.cfg["lora_alpha"],
            target_modules=self.cfg["lora_target_modules"],
            lora_dropout=self.cfg["lora_dropout"],
            bias="none",
            task_type="CAUSAL_LM",
        )
        
        self.model = get_peft_model(self.model, lora_config)
        self.model.print_trainable_parameters()
        
        log.info("✅ 모델 로딩 완료")
    
    def load_data(self):
        """데이터 로딩 및 전처리"""
        log.info("🔄 데이터 로딩 및 전처리 시작")
        log.info(f"▶️  데이터 로딩: {self.cfg['data_path']}")
        
        # JSONL 파일 로딩
        dataset = load_dataset("json", data_files=self.cfg["data_path"], split="train")
        
        log.info(f"📊 데이터셋 컬럼: {dataset.column_names}")
        log.info(f"📊 데이터셋 크기: {len(dataset)}개 샘플")
        
        # 데이터 형식을 감지하고 적절한 변환 함수 선택
        def detect_format(example):
            """데이터 형식을 감지하는 함수"""
            if "messages" in example:
                return "chat"
            elif "prompt" in example and "completion" in example:
                return "prompt_completion"
            elif "error_context" in example and "explanation" in example:
                return "error_explanation"
            elif "error_context" in example and "buggy_code_snippet" in example:
                return "error_fix"
            elif "instruction" in example and "input" in example and "output" in example:
                return "instruction_input_output"
            else:
                return "unknown"
        
        # 샘플을 검사하여 데이터 형식 감지
        sample = dataset[0]
        data_format = detect_format(sample)
        log.info(f"✅ 감지된 데이터 형식: {data_format}")
        
        # 토크나이징 함수
        def tokenize_function(examples):
            # 데이터 형식에 따라 ChatML 형식으로 변환
            texts = []
            
            # 샘플 수 결정 (다양한 형식 지원)
            if "messages" in examples:
                sample_count = len(examples["messages"])
            elif "prompt" in examples:
                sample_count = len(examples["prompt"])
            elif "error_context" in examples:
                sample_count = len(examples["error_context"])
            elif "instruction" in examples:
                sample_count = len(examples["instruction"])
            else:
                log.error("❌ 지원되지 않는 데이터 형식")
                raise ValueError("지원되지 않는 데이터 형식입니다.")
            
            for i in range(sample_count):
                # 데이터 형식별 처리
                if data_format == "chat":
                    # 채팅 형식 (messages 구조)
                    messages = examples["messages"][i]
                    chat_text = ""
                    
                    # 모든 메시지를 ChatML 형식으로 변환
                    for msg in messages:
                        role = msg["role"]
                        content = msg["content"]
                        chat_text += f"<|im_start|>{role}\n{content}<|im_end|>\n"
                    
                    texts.append(chat_text.strip())
                
                elif data_format == "prompt_completion":
                    # 주석 키워드 기반 코드 생성 형식
                    prompt = examples["prompt"][i]
                    completion = examples["completion"][i]
                    
                    # 프롬프트를 사용자 메시지로, 완성을 어시스턴트 메시지로 변환
                    text = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n{completion}<|im_end|>"
                    texts.append(text)
                
                elif data_format == "error_explanation":
                    # 에러 설명 형식
                    error_context = examples["error_context"][i]
                    explanation = examples["explanation"][i]
                    
                    # 에러 컨텍스트 정보 구성
                    error_log = error_context.get("error_log", "")
                    code_snippet = error_context.get("code_snippet", "")
                    language = error_context.get("language", "")
                    
                    # 사용자 입력 구성
                    user_text = f"다음 {language} 코드의 에러를 설명해주세요:\n\n에러 로그:\n{error_log}\n\n코드:\n{code_snippet}"
                    
                    # ChatML 형식으로 변환
                    text = f"<|im_start|>user\n{user_text}<|im_end|>\n<|im_start|>assistant\n{explanation}<|im_end|>"
                    texts.append(text)
                
                elif data_format == "error_fix":
                    # 에러 수정 형식
                    error_context = examples["error_context"][i]
                    buggy_code = examples["buggy_code_snippet"][i]
                    fixed_code = examples["fixed_code_snippet"][i]
                    
                    # 에러 컨텍스트 정보 구성
                    error_log = error_context.get("error_log", "")
                    language = error_context.get("language", "")
                    
                    # 사용자 입력 구성
                    user_text = f"다음 {language} 코드의 에러를 수정해주세요:\n\n에러 로그:\n{error_log}\n\n코드:\n{buggy_code}"
                    
                    # ChatML 형식으로 변환
                    text = f"<|im_start|>user\n{user_text}<|im_end|>\n<|im_start|>assistant\n{fixed_code}<|im_end|>"
                    texts.append(text)
                    
                elif data_format == "instruction_input_output":
                    # 기존 instruction, input, output 형식
                    instruction = examples["instruction"][i].strip()
                    input_text = examples["input"][i].strip()
                    output_text = examples["output"][i].strip()
                    
                    # 태그 정보 추가 (있다면)
                    tags = examples.get("tags", [None] * len(examples["instruction"]))[i]
                    if tags:
                        if isinstance(tags, str):
                            tags = [tags]
                        tag_info = f"Task Type: {', '.join(tags)}\n"
                        if input_text:
                            user_text = f"{tag_info}Instruction: {instruction}\nInput: {input_text}"
                        else:
                            user_text = f"{tag_info}Instruction: {instruction}"
                    else:
                        if input_text:
                            user_text = f"Instruction: {instruction}\nInput: {input_text}"
                        else:
                            user_text = f"Instruction: {instruction}"
                    
                    # ChatML 형식
                    text = f"<|im_start|>user\n{user_text}<|im_end|>\n<|im_start|>assistant\n{output_text}<|im_end|>"
                    texts.append(text)
                
                else:
                    log.error(f"❌ 지원되지 않는 데이터 형식: {data_format}")
                    raise ValueError(f"지원되지 않는 데이터 형식: {data_format}")
            
            # FIM 형식 특수 처리 (prompt에 <|fim begin|>, <|fim hole|>, <|fim end|> 태그가 있는 경우)
            if data_format == "prompt_completion":
                for i in range(len(texts)):
                    prompt = examples["prompt"][i]
                    if "<｜fim begin｜>" in prompt and "<｜fim hole｜>" in prompt and "<｜fim end｜>" in prompt:
                        # FIM 형식의 특수 처리
                        completion = examples["completion"][i]
                        
                        # fim begin과 fim hole 사이의 텍스트를 추출 (prefix)
                        prefix = prompt.split("<｜fim hole｜>")[0].replace("<｜fim begin｜>", "")
                        # fim hole과 fim end 사이의 텍스트를 추출 (suffix)
                        suffix = prompt.split("<｜fim hole｜>")[1].split("<｜fim end｜>")[0]
                        
                        # FIM 형식의 프롬프트로 직접 변환
                        text = f"<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>{completion}"
                        texts[i] = text
            
            # 토크나이징
            model_inputs = self.tok(
                texts,
                truncation=True,
                padding=False,
                max_length=self.cfg["max_length"],
                return_tensors=None
            )
            
            # labels = input_ids (causal LM)
            model_inputs["labels"] = model_inputs["input_ids"].copy()
            
            return model_inputs
        
        # 데이터셋 토크나이징
        dataset = dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=dataset.column_names,
            desc="토크나이징 중"
        )
        
        log.info("✅ 데이터 토크나이징 완료")
        
        # 학습/검증 분할
        split_dataset = dataset.train_test_split(test_size=0.1, seed=42)
        self.train_ds = split_dataset["train"]
        self.eval_ds = split_dataset["test"]
        
        log.info(f"📊 학습 데이터: {len(self.train_ds)}개")
        log.info(f"📊 검증 데이터: {len(self.eval_ds)}개")
    
    def t_args(self) -> TrainingArguments:
        """완전 호환 학습 인자 생성"""
        # 절대 경로 사용으로 변경 (AWS 스팟 인스턴스 중단 시 안전한 체크포인트 저장을 위함)
        # 학습 경로 설정 - 체크포인트는 원래 경로에 저장
        scripts_dir = "/home/ubuntu/deepseek-coder/scripts"
        
        # 학습 체크포인트와 출력 경로
        output_dirs = {
            "complete": f"{scripts_dir}/output/autocomplete-finetuned",
            "prompt": f"{scripts_dir}/output/prompt-finetuned",
            "comment": f"{scripts_dir}/output/comment-finetuned",
            "error_fix": f"{scripts_dir}/output/error-fix-finetuned"
        }
        
        # 최종 모델 경로 (추론용)
        models_dir = "../models"
        final_model_dirs = {
            "complete": f"{models_dir}/autocomplete-finetuned",
            "prompt": f"{models_dir}/prompt-finetuned",
            "comment": f"{models_dir}/comment-finetuned",
            "error_fix": f"{models_dir}/error-fix-finetuned"
        }
        
        # 학습 경로 (체크포인트와 출력 디렉토리)
        output_dir = output_dirs.get(self.mode, f"{scripts_dir}/output/prompt-finetuned")
        
        # 최종 모델 경로 (inference.py에서 사용할 경로) 별도 저장
        self.final_model_dir = final_model_dirs.get(self.mode, f"{models_dir}/prompt-finetuned")
        
        # 최소한의 안전한 인자만 사용 (타입 변환 추가)
        return TrainingArguments(
            output_dir=output_dir,
            per_device_train_batch_size=self.get_int_config("batch_size", 1),
            per_device_eval_batch_size=self.get_int_config("batch_size", 1),
            gradient_accumulation_steps=self.get_int_config("grad_acc", 8),
            learning_rate=self.get_float_config("learning_rate", 2e-4),
            num_train_epochs=self.get_float_config("num_epochs", 3),
            warmup_ratio=self.get_float_config("warmup_ratio", 0.1),
            lr_scheduler_type=self.cfg["lr_scheduler"],
            weight_decay=self.get_float_config("weight_decay", 0.01),
            fp16=self.cfg["fp16"],
            bf16=self.cfg["bf16"],
            gradient_checkpointing=self.cfg["gradient_checkpointing"],
            dataloader_num_workers=self.get_int_config("dataloader_num_workers", 4),
            save_steps=self.get_int_config("save_steps", 500),
            eval_steps=self.get_int_config("eval_steps", 500),
            logging_steps=self.get_int_config("logging_steps", 50),
            save_total_limit=self.get_int_config("save_total_limit", 3),
            remove_unused_columns=self.cfg["remove_unused_columns"],
            report_to=None,
            run_name=f"deepseek-coder-{self.mode}-{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        )
    
    def train(self):
        """모델 학습 실행"""
        log.info(f"🚀 학습 시작: {self.mode} 모드")
        
        # 학습 인자 준비
        args = self.t_args()
        
        # 체크포인트 감지 및 자동 재개 기능
        checkpoint_dir = args.output_dir
        resume_from_checkpoint = None
        
        # 기존 체크포인트 확인
        if os.path.exists(checkpoint_dir):
            checkpoints = [f for f in os.listdir(checkpoint_dir) if f.startswith("checkpoint-")]
            
            if checkpoints:
                # 체크포인트 번호 기준으로 정렬
                sorted_checkpoints = sorted(
                    checkpoints,
                    key=lambda x: int(x.split("-")[1]) if len(x.split("-")) > 1 and x.split("-")[1].isdigit() else 0
                )
                
                if sorted_checkpoints:
                    latest_checkpoint = sorted_checkpoints[-1]
                    resume_from_checkpoint = os.path.join(checkpoint_dir, latest_checkpoint)
                    log.info(f"🔄 체크포인트 감지: {resume_from_checkpoint}에서 학습 재개")
        
        # 데이터 콜레이터
        data_collator = DataCollatorForLanguageModeling(
            tokenizer=self.tok,
            mlm=False,  # Causal LM
            pad_to_multiple_of=8
        )
        
        # Trainer 설정 - resume_from_checkpoint는 train() 메서드에만 전달
        trainer = Trainer(
            model=self.model,
            args=args,
            train_dataset=self.train_ds,
            eval_dataset=self.eval_ds,
            data_collator=data_collator,
            tokenizer=self.tok
        )
        
        log.info("✅ Trainer 초기화 완료, 체크포인트 감지 활성화")
        
        # 학습 실행
        try:
            trainer.train(resume_from_checkpoint=resume_from_checkpoint)
            log.info("✅ 학습 완료")
            
            # 1. 최종 모델 저장 - args.output_dir 내의 final_model 경로에 저장
            output_model_path = os.path.join(args.output_dir, "final_model")
            trainer.save_model(output_model_path)
            self.tok.save_pretrained(output_model_path)
            
            log.info(f"✅ 최종 모델 저장 완료: {output_model_path}")
            
            # 2. inference.py 호환을 위해 모델을 ../models 경로로도 복사
            import shutil
            final_model_path = os.path.join(self.final_model_dir, "final_model")
            
            # 디렉토리가 없으면 생성
            os.makedirs(os.path.dirname(final_model_path), exist_ok=True)
            
            if os.path.exists(output_model_path) and output_model_path != final_model_path:
                try:
                    # 기존 디렉토리 삭제 후 복사
                    if os.path.exists(final_model_path):
                        shutil.rmtree(final_model_path)
                    
                    # 모델 디렉토리 복사
                    shutil.copytree(output_model_path, final_model_path)
                    log.info(f"✅ inference.py 호환을 위해 모델 복사 완료: {final_model_path}")
                except Exception as e:
                    log.error(f"❌ 모델 복사 중 오류 발생: {e}")
            
            return final_model_path
            
        except Exception as e:
            log.error(f"❌ 학습 중 오류 발생: {e}")
            raise
    
    def run(self):
        """전체 학습 파이프라인 실행"""
        try:
            self.load_model()
            self.load_data()
            final_model_path = self.train()
            
            log.info("🎉 학습 파이프라인 완료!")
            return final_model_path
            
        except Exception as e:
            log.error(f"❌ 학습 파이프라인 실패: {e}")
            raise

def parse_arguments():
    """명령줄 인자 파싱"""
    parser = argparse.ArgumentParser(description="DeepSeek-Coder 6.7B QLoRA 파인튜닝 (완전 호환성)")
    parser.add_argument("--config", type=str, default="config.yaml", help="설정 파일 경로")
    parser.add_argument("--mode", type=str, choices=["complete", "prompt", "comment", "error_fix"], default="prompt", help="학습 모드")
    return parser.parse_args()

def main():
    """메인 함수"""
    args = parse_arguments()
    
    mode_descriptions = {
        "complete": "[제1차] 코드 자동완성 전용 (FIM 형식)",
        "prompt": "[제2차] 일반 프롬프트 기반 코드 생성",
        "comment": "[제3차] 주석 기반 코드 생성",
        "error_fix": "[제4차] 오류 코드 설명 및 수정"
    }
    
    log.info("=" * 53)
    log.info("DeepSeek Coder 6.7B 완전 호환성 파인튜닝 시작")
    log.info("=" * 53)
    log.info(f"시작 시간: {datetime.now().strftime('%a %b %d %H:%M:%S %Z %Y')}")
    log.info(f"학습 모드: [{args.mode}] {mode_descriptions.get(args.mode, args.mode)}")
    
    # 환경 정보 출력
    log.info(f"Python 명령어: {sys.executable}")
    if torch.cuda.is_available():
        gpu_info = f"{torch.cuda.get_device_name(0)}, {torch.cuda.get_device_properties(0).total_memory // 1024**2}, {torch.cuda.memory_allocated() // 1024**2}"
        log.info(f"GPU 정보:\n{gpu_info}")
    
    log.info(f"Python 버전: Python {sys.version.split()[0]}")
    
    # 패키지 버전 확인
    log.info("확인 중: 필요한 패키지들...")
    packages = ["torch", "transformers", "datasets", "accelerate", "peft"]
    for pkg in packages:
        try:
            module = __import__(pkg)
            version = getattr(module, "__version__", "unknown")
            log.info(f"✅ {pkg}: {version}")
        except ImportError:
            log.error(f"❌ {pkg}: 설치되지 않음")
    
    try:
        # 학습 실행
        trainer = DeepSeekTrainer(config_path=args.config, mode=args.mode)
        final_model_path = trainer.run()
        
        log.info("=" * 53)
        log.info("학습이 성공적으로 완료되었습니다!")
        log.info("=" * 53)
        log.info(f"모델 저장 위치: {final_model_path}")
        log.info(f"inference.py에서 이 모델을 사용하려면: python inference.py --mode {args.mode} --model {final_model_path}")
        
    except KeyboardInterrupt:
        log.info("⏹️ 사용자에 의해 학습이 중단되었습니다.")
    except Exception as e:
        log.error(f"❌ 학습 실패: {e}")
        sys.exit(1)
    finally:
        # 최종 GPU 상태 출력
        if torch.cuda.is_available():
            log.info("최종 GPU 상태:")
            os.system("nvidia-smi")

if __name__ == "__main__":
    main()

run_training.sh

#!/bin/bash

# DeepSeek 훈련 환경 설정 스크립트

set -e

echo "====================================================="
echo "DeepSeek 훈련 환경 설정 시작"
echo "====================================================="

# 작업 디렉토리 확인
WORK_DIR="$(pwd)"
echo "작업 디렉토리: $WORK_DIR"

# Python 설치 확인 및 설치
echo "1. Python 설치 확인..."
if ! command -v python3 &> /dev/null; then
    echo "Python3가 설치되어 있지 않습니다. 설치를 시작합니다..."
    sudo apt update
    sudo apt install -y python3 python3-pip python3-venv
    echo "✅ Python3 설치 완료"
else
    echo "✅ Python3가 이미 설치되어 있습니다: $(python3 --version)"
fi

# pip 업그레이드
echo "2. pip 업그레이드..."
python3 -m pip install --upgrade pip

# 가상 환경 생성 및 활성화
echo "3. 가상 환경(Hidle) 생성..."
if [ ! -d "Hidle" ]; then
    echo "가상 환경을 생성합니다..."
    python3 -m venv Hidle
    echo "✅ 가상 환경 생성 완료"
else
    echo "✅ 가상 환경이 이미 존재합니다"
fi

echo "4. 가상 환경 활성화..."
source Hidle/bin/activate
echo "✅ 가상 환경 활성화됨: $(which python)"

# 의존성 설치
echo "5. 필요한 패키지 설치..."
pip install --upgrade pip

if [ -f "requirements.txt" ]; then
    echo "requirements.txt에서 패키지 설치 중..."
    pip install -r requirements.txt
    echo "✅ 패키지 설치 완료"
else
    echo "⚠️ requirements.txt 파일이 존재하지 않습니다. 지정된 버전의 기본 패키지를 설치합니다."
    pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 transformers==4.31.0 datasets==2.12.0 accelerate==0.21.0 peft==0.4.0 trl==0.7.1 sentencepiece==0.1.97 protobuf==3.20.0 bitsandbytes==0.41.0 scipy==1.10.0 wandb==0.15.5
    echo "✅ 기본 패키지 설치 완료"
fi

# CUDA 설치 확인
echo "6. CUDA 설치 확인..."
if command -v nvidia-smi &> /dev/null; then
    echo "✅ NVIDIA GPU 감지됨:"
    nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader
    CUDA_VERSION=$(nvcc --version 2>/dev/null | grep "release" | awk '{print $6}' | cut -c2- || echo "미확인")
    echo "CUDA 버전: $CUDA_VERSION"
else
    echo "⚠️ nvidia-smi를 찾을 수 없습니다. CPU 모드로 실행됩니다."
fi

# 올바른 PyTorch 버전 확인
echo "7. PyTorch GPU 지원 확인..."
python -c "import torch; print('CUDA 사용 가능:', torch.cuda.is_available(), '/ 장치 수:', torch.cuda.device_count(), '/ 현재 장치:', torch.cuda.current_device(), '/ 장치 이름:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'None')"

# 데이터 디렉토리 확인
echo "8. 데이터 디렉토리 확인..."
if [ ! -d "data" ]; then
    echo "⚠️ 'data' 디렉토리가 존재하지 않습니다. 데이터셋 파일을 저장할 디렉토리를 생성합니다."
    mkdir -p data
    echo "✅ 'data' 디렉토리 생성 완료"
fi

# 출력 디렉토리 확인
if [ ! -d "output" ]; then
    echo "⚠️ 'output' 디렉토리가 존재하지 않습니다. 모델 체크포인트를 저장할 디렉토리를 생성합니다."
    mkdir -p output
    echo "✅ 'output' 디렉토리 생성 완료"
fi

# 전체 패키지 설치 확인
echo "9. 패키지 설치 확인..."
python -c "
import sys
import torch
import transformers
import datasets
import accelerate
import peft
try:
    import bitsandbytes
    bnb_available = True
except ImportError:
    bnb_available = False

print(f'✅ Python: {sys.version.split()[0]}')
print(f'✅ PyTorch: {torch.__version__}')
print(f'✅ Transformers: {transformers.__version__}')
print(f'✅ Datasets: {datasets.__version__}')
print(f'✅ Accelerate: {accelerate.__version__}')
print(f'✅ PEFT: {peft.__version__}')
print(f'✅ BitsAndBytes: {\"설치됨\" if bnb_available else \"설치되지 않음\"}')
print(f'✅ CUDA 사용 가능: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'✅ CUDA 디바이스: {torch.cuda.get_device_name(0)}')
    print(f'✅ 가용 GPU 메모리: {torch.cuda.get_device_properties(0).total_memory / (1024**3):.2f} GB')
"

# AWS 스팟 인스턴스 중단 감지 설정 확인
echo "10. AWS 스팟 인스턴스 환경 확인..."
if [ -f "test_environment.py" ]; then
    echo "✅ 스팟 인스턴스 중단 감지 스크립트가 존재합니다."
else
    echo "⚠️ 스팟 인스턴스 중단 감지 스크립트가 없습니다."
    # 스크립트 생성
    cat > test_environment.py << 'EOF'
#!/usr/bin/env python3
"""
AWS 스팟 인스턴스 중단 감지 스크립트
"""

import os
import time
import signal
import threading
import requests
import logging
import subprocess

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler()]
)
log = logging.getLogger(__name__)

def check_spot_termination():
    """스팟 인스턴스 중단 신호를 확인하고 학습 프로세스에 SIGTERM 신호를 보냔"""
    try:
        # 스팟 인스턴스 중단 신호 업데이트 확인
        url = "http://169.254.169.254/latest/meta-data/spot/instance-action"
        response = requests.get(url, timeout=2)

        if response.status_code == 200:
            # 중단 신호가 있음 - train.py 프로세스에 SIGTERM 전송
            log.warning("스팟 인스턴스 중단 신호 감지: %s", response.text)

            # train.py 프로세스 시그널 전송
            result = subprocess.run(["pgrep", "-f", "train\.py"], capture_output=True, text=True)
            if result.returncode == 0:
                pid = int(result.stdout.strip())
                log.info("train.py 프로세스 (PID: %s)에 SIGTERM 신호 전송", pid)
                os.kill(pid, signal.SIGTERM)
                return True
    except requests.RequestException:
        # EC2 메타데이터 API에 접근할 수 없음 - 일반 인스턴스일 수 있음
        pass
    except Exception as e:
        log.error("스팟 인스턴스 중단 신호 확인 중 오류: %s", e)

    return False

def monitoring_thread_func():
    """지속적으로 스팟 인스턴스 중단을 감시하는 백그라운드 스레드"""
    log.info("스팟 인스턴스 중단 감지 도구 시작")

    while True:
        if check_spot_termination():
            log.warning("스팟 인스턴스 중단 감지 - 학습 중단 및 체크포인트 저장 프로세스 시작")
            break

        # 5초마다 확인
        time.sleep(5)

if __name__ == "__main__":
    # 백그라운드 스레드로 스팟 인스턴스 중단 감시 시작
    monitoring_thread = threading.Thread(target=monitoring_thread_func, daemon=True)
    monitoring_thread.start()

    log.info("메인 프로세스 실행 중 (중단하려면 Ctrl+C 를 누르세요)")
    try:
        # 메인 스레드는 계속 실행
        while True:
            time.sleep(60)
    except KeyboardInterrupt:
        log.info("사용자에 의해 중지됨")

    log.info("스팟 인스턴스 중단 감지 도구 종료")
EOF
    chmod +x test_environment.py
    echo "✅ 스팟 인스턴스 중단 감지 스크립트 생성 완료"
fi

echo "====================================================="
echo "환경 설정 완료!"
echo "====================================================="
echo "이제 다음 명령어로 훈련을 시작할 수 있습니다:"
echo "./run_training.sh"
echo ""
echo "AWS 스팟 인스턴스 중단 감지 스크립트를 백그라운드로 실행하려면:"
echo "python test_environment.py &"
echo ""
echo "가상 환경을 사용하는 경우 먼저 활성화하세요:"
echo "source Hidle/bin/activate"

inference.py

#!/usr/bin/env python3
"""
DeepSeek-Coder 실시간 코드 어시스트 인퍼런스 스크립트 (T4 환경 최적화)
지원 기능:
- 사용자 입력 기반 프럼프트 응답
- 오류 수정 및 작성 지원
- 코드 자동완성
- T4 GPU 전용 최적화
"""

import os
import sys
import json
import argparse
import torch
import readline
import argparse
from transformers import AutoTokenizer, AutoModelForCausalLM

# ────────── T4 환경 최적화 설정 ──────────
MODEL_CONFIG = {
    "model_id": "deepseek-ai/deepseek-coder-6.7b-instruct",  # 기본 모델
    "model_path": "../models/prompt-finetuned/final_model",  # train.py [2차] 학습 출력 경로
    "device": "cuda" if torch.cuda.is_available() else "cpu", 
    "dtype": torch.float16,  # T4에서 bfloat16 대신 float16 사용
    "max_tokens": 512,
    "temperature": 0.7,
    "top_p": 0.95,
    "load_in_4bit": True,  # T4 메모리 효율성을 위한 4-bit 양자화
}

def check_gpu_compatibility():
    """GPU 호환성 확인"""
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        compute_capability = torch.cuda.get_device_capability(0)
        print(f"🖥️ GPU: {gpu_name}")
        print(f"🔧 Compute Capability: {compute_capability[0]}.{compute_capability[1]}")
        
        # bfloat16 지원 확인
        if torch.cuda.is_bf16_supported():
            print("✅ bfloat16 지원됨")
        else:
            print("⚠️ bfloat16 미지원 - float16 사용")
            
        return True
    else:
        print("❌ CUDA 사용 불가 - CPU 모드")
        return False

class CodeAssistant:
    """코드 어시스트 추론 클래스 (T4 최적화)"""
    
    def __init__(self, config=None, mode="code-gen"):
        self.config = config or MODEL_CONFIG
        self.mode = mode
        
        # GPU 호환성 확인
        self.cuda_available = check_gpu_compatibility()
        
        # T4 환경에서 dtype 조정
        if self.cuda_available:
            # T4에서는 bfloat16 대신 float16 사용
            self.config["dtype"] = torch.float16
        else:
            self.config["dtype"] = torch.float32
            self.config["device"] = "cpu"
        
        # 모드에 따라 다른 모델 선택
        model_path = self._get_model_path()
        self.config["model_path"] = model_path
        
        self.model = None
        self.tokenizer = None
        self.load_model()
        
    def _get_model_path(self):
        """모드에 따른 모델 경로 결정"""
        if self.mode == "autocomplete":
            model_path = "../models/autocomplete-finetuned/final_model"
            if not os.path.exists(model_path):
                print(f"\n⚠️ [제1차] 자동완성 전용 모델을 찾을 수 없어 기본 모델을 사용합니다.")
                model_path = self.config["model_path"]
        elif self.mode == "comment":
            model_path = "../models/comment-finetuned/final_model"
            if not os.path.exists(model_path):
                print(f"\n⚠️ [제3차] 주석 기반 코드 생성 모델을 찾을 수 없어 기본 모델을 사용합니다.")
                model_path = self.config["model_path"]
        elif self.mode == "error_fix":
            model_path = "../models/error-fix-finetuned/final_model"
            if not os.path.exists(model_path):
                print(f"\n⚠️ [제4차] 오류 코드 수정 전용 모델을 찾을 수 없어 기본 모델을 사용합니다.")
                model_path = self.config["model_path"]
        else:  # code-gen 또는 prompt 모드
            model_path = self.config["model_path"]  # prompt 모델이 기본값
        
        return model_path
        
    def load_model(self):
        """T4 환경 최적화된 모델과 토크나이저 로딩"""
        print(f"\n📂 모델 로딩 시작 (T4 최적화)")
        model_path = self.config.get('model_path')
        model_id = self.config.get('model_id')
        
        # 로컬 파인튜닝 모델이 있으면 먼저 시도
        try_path = model_path if os.path.exists(model_path) else model_id
        print(f"🔍 모델 경로: {try_path}")
        print(f"🖥️ 디바이스: {self.config['device']}")
        print(f"🔢 데이터 타입: {self.config['dtype']}")
        
        try:
            # 토크나이저 로딩
            self.tokenizer = AutoTokenizer.from_pretrained(
                try_path, 
                trust_remote_code=True
            )
            
            # T4 환경 최적화된 모델 로딩
            if self.cuda_available and self.config.get("load_in_4bit", False):
                # 4-bit 양자화로 메모리 효율성 향상
                from transformers import BitsAndBytesConfig
                
                bnb_config = BitsAndBytesConfig(
                    load_in_4bit=True,
                    bnb_4bit_compute_dtype=torch.float16,  # T4 호환
                    bnb_4bit_quant_type="nf4",
                    bnb_4bit_use_double_quant=True,
                )
                
                self.model = AutoModelForCausalLM.from_pretrained(
                    try_path,
                    quantization_config=bnb_config,
                    device_map="auto",
                    trust_remote_code=True,
                    torch_dtype=torch.float16
                )
                print("✅ 4-bit 양자화 모델 로딩 완료")
            else:
                # 일반 로딩
                self.model = AutoModelForCausalLM.from_pretrained(
                    try_path,
                    torch_dtype=self.config['dtype'],
                    device_map="auto" if self.config['device'] == "cuda" else None,
                    trust_remote_code=True
                )
                
                # CPU의 경우 명시적 이동
                if self.config['device'] == "cpu":
                    self.model = self.model.to("cpu")
                    
            self.model.eval()
            
            # 메모리 사용량 확인
            if self.cuda_available:
                memory_used = torch.cuda.memory_allocated() / 1024**3
                memory_total = torch.cuda.get_device_properties(0).total_memory / 1024**3
                print(f"📊 GPU 메모리 사용량: {memory_used:.2f}GB / {memory_total:.2f}GB")
            
            print("✅ 모델 로딩 완료")
            
        except Exception as e:
            print(f"❌ 모델 로딩 실패: {e}")
            print("💡 해결 방안:")
            print("   1. config.yaml에서 bnb_4bit_compute_dtype: float16으로 설정")
            print("   2. torch_dtype: float16으로 설정")
            print("   3. 메모리 부족 시 per_device_train_batch_size 줄이기")
            sys.exit(1)
    
    def get_int_config(self, key, default=0):
        """config에서 값을 가져와 int로 변환하는 헬퍼 함수"""
        value = self.config.get(key, default)
        if isinstance(value, str):
            try:
                value = int(value)
            except ValueError:
                print(f"⚠️ 경고: {key} 값 '{value}'을 정수로 변환할 수 없습니다. 기본값 {default} 사용")
                value = default
        return value
    
    def get_float_config(self, key, default=0.0):
        """config에서 값을 가져와 float로 변환하는 헬퍼 함수"""
        value = self.config.get(key, default)
        if isinstance(value, str):
            try:
                value = float(value)
            except ValueError:
                print(f"⚠️ 경고: {key} 값 '{value}'을 실수로 변환할 수 없습니다. 기본값 {default} 사용")
                value = default
        return value
    
    def build_prompt(self, user_input, mode="code-gen"):
        """ChatML 형식으로 프럼프트 구성 (다양한 데이터 구조 지원)"""
        system_msg = "You are an expert coding assistant."
        
        # 다양한 입력 형식 처리
        if isinstance(user_input, dict):
            # 1. 채팅 형식 (messages 구조)
            if "messages" in user_input:
                messages = user_input["messages"]
                chat_text = ""
                
                # 시스템 메시지 추가
                chat_text += f"<|im_start|>system\n{system_msg}<|im_end|>\n"
                
                # 모든 메시지를 ChatML 형식으로 변환
                for msg in messages:
                    role = msg["role"]
                    content = msg["content"]
                    chat_text += f"<|im_start|>{role}\n{content}<|im_end|>\n"
                
                # 어시스턴트 응답 시작 태그 추가
                chat_text += "<|im_start|>assistant\n"
                return chat_text
            
            # 2. 주석 키워드 기반 코드 생성 (prompt-completion)
            elif "prompt" in user_input:
                prompt = user_input["prompt"]
                
                # FIM 형식 특수 처리
                if "<｜fim begin｜>" in prompt and "<｜fim hole｜>" in prompt and "<｜fim end｜>" in prompt:
                    # fim begin과 fim hole 사이의 텍스트를 추출 (prefix)
                    prefix = prompt.split("<｜fim hole｜>")[0].replace("<｜fim begin｜>", "")
                    # fim hole과 fim end 사이의 텍스트를 추출 (suffix)
                    suffix = prompt.split("<｜fim hole｜>")[1].split("<｜fim end｜>")[0]
                    
                    # FIM 스페셜 토큰 형식으로 구성
                    return f"<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|"
                else:
                    # 일반 프롬프트 처리
                    return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
            
            # 3. 에러 설명 형식
            elif "error_context" in user_input and "explanation" in user_input.get("error_context", {}):
                error_context = user_input["error_context"]
                error_log = error_context.get("error_log", "")
                code_snippet = error_context.get("code_snippet", "")
                language = error_context.get("language", "")
                
                user_text = f"다음 {language} 코드의 에러를 설명해주세요:\n\n에러 로그:\n{error_log}\n\n코드:\n{code_snippet}"
                return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\n{user_text}<|im_end|>\n<|im_start|>assistant\n"
            
            # 4. 에러 수정 형식
            elif "error_context" in user_input and "buggy_code_snippet" in user_input:
                error_context = user_input["error_context"]
                buggy_code = user_input["buggy_code_snippet"]
                error_log = error_context.get("error_log", "")
                language = error_context.get("language", "")
                
                user_text = f"다음 {language} 코드의 에러를 수정해주세요:\n\n에러 로그:\n{error_log}\n\n코드:\n{buggy_code}"
                return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\n{user_text}<|im_end|>\n<|im_start|>assistant\n"
            
            # 5. 기존 instruction/input 형식
            elif "instruction" in user_input:
                instruction = user_input.get("instruction", "").strip()
                input_text = user_input.get("input", "").strip()
                
                # instruction과 input을 합쳐서 user 메시지로 구성
                if input_text:
                    user_text = f"Instruction: {instruction}\nInput: {input_text}"
                else:
                    user_text = f"Instruction: {instruction}"
                
                # ChatML 형식 프럼프트
                if mode == "error_fix":
                    return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\nPlease fix the following code:\n\n{user_text}<|im_end|>\n<|im_start|>assistant\n"
                elif mode == "autocomplete":
                    return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\nComplete the following code:\n\n{user_text}<|im_end|>\n<|im_start|>assistant\n"
                elif mode == "comment":
                    return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\nGenerate code based on this comment:\n\n{user_text}<|im_end|>\n<|im_start|>assistant\n"
                else:  # code-gen
                    return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\n{user_text}<|im_end|>\n<|im_start|>assistant\n"
            
            # 기타 딕셔너리 형태의 입력
            else:
                user_text = json.dumps(user_input, ensure_ascii=False)
        else:
            # 기존 문자열 입력 처리 (하위 호환성)
            user_text = str(user_input)
        
        # 기본 ChatML 형식 프럼프트
        if mode == "error_fix":
            return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\nPlease fix the following code:\n\n{user_text}<|im_end|>\n<|im_start|>assistant\n"
        elif mode == "autocomplete":
            return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\nComplete the following code:\n\n{user_text}<|im_end|>\n<|im_start|>assistant\n"
        elif mode == "comment":
            return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\nGenerate code based on this comment:\n\n{user_text}<|im_end|>\n<|im_start|>assistant\n"
        else:  # code-gen
            return f"<|im_start|>system\n{system_msg}<|im_end|>\n<|im_start|>user\n{user_text}<|im_end|>\n<|im_start|>assistant\n"
    
    def generate(self, user_input, mode="code-gen"):
        """T4 최적화된 텍스트 생성"""
        # 다양한 데이터 형식 확인
        is_fim_format = False
        if isinstance(user_input, dict) and "prompt" in user_input:
            prompt_str = user_input["prompt"]
            if "<｜fim begin｜>" in prompt_str and "<｜fim hole｜>" in prompt_str and "<｜fim end｜>" in prompt_str:
                is_fim_format = True
                
        # 프롬프트 구성
        prompt = self.build_prompt(user_input, mode)
        
        # T4 메모리 제한을 고려한 토큰 길이 조정
        inputs = self.tokenizer(prompt, return_tensors="pt", truncate=True, max_length=2048)
        
        if self.config['device'] == "cuda":
            inputs = {k: v.to("cuda") for k, v in inputs.items()}
        
        # 모드별 생성 파라미터 조정
        if is_fim_format:
            # FIM 용 특수 설정
            max_new_tokens = min(512, self.get_int_config('max_tokens', 512))
            temp = 0.2  # 더 결정적으로
        elif mode == "autocomplete":
            max_new_tokens = min(128, self.get_int_config('max_tokens', 512))  # 자동완성은 짧게
            temp = 0.3  # 더 결정적으로
        elif mode == "comment":
            max_new_tokens = min(256, self.get_int_config('max_tokens', 512))
            temp = self.get_float_config('temperature', 0.6)
        else:  # code-gen, error_fix
            max_new_tokens = self.get_int_config('max_tokens', 512)
            temp = self.get_float_config('temperature', 0.6)
            
        with torch.no_grad():
            # T4 메모리 효율성을 위한 생성 설정
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=(mode != "autocomplete" and not is_fim_format),
                top_p=self.get_float_config('top_p', 0.95),
                temperature=temp,
                top_k=self.get_int_config('top_k', 50),
                repetition_penalty=self.get_float_config('repetition_penalty', 1.1),
                pad_token_id=self.tokenizer.eos_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
                use_cache=True,  # 캐시 사용으로 속도 향상
            )
        
        # 프롬프트 길이 이후 응답 부분만 추출
        input_length = inputs['input_ids'].shape[1]
        generated = self.tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
    
    def _clean_output(self, text):
        """생성된 텍스트에서 ChatML 태그 및 불필요한 부분 제거"""
        # ChatML 태그 제거
        text = text.replace("<|im_start|>assistant\n", "").replace("<|im_end|>", "")
        
        # FIM 태그 처리 (만약 있을 경우)
        if "<|fim_middle|>" in text:
            text = text.replace("<|fim_middle|>", "")
        if "<|fim_prefix|>" in text:
            text = text.replace("<|fim_prefix|>", "")
        if "<|fim_suffix|>" in text:
            text = text.replace("<|fim_suffix|>", "")
        
        # 불필요한 반복 제거
        lines = text.split('\n')
        cleaned_lines = []
        for line in lines:
            if line.strip() and not line.startswith('<|'):
                cleaned_lines.append(line)
        
        return '\n'.join(cleaned_lines).strip()

def parse_arguments():
    """명령줄 인자 파싱"""
    parser = argparse.ArgumentParser(description="DeepSeek-Coder 코드 어시스트 추론 (T4 최적화)")
    parser.add_argument(
        "--model", 
        type=str, 
        default="../models/prompt-finetuned/final_model", 
        help="모델 경로 또는 모델 ID"
    )
    parser.add_argument(
        "--mode", 
        type=str, 
        choices=["code-gen", "prompt", "autocomplete", "comment", "error_fix"], 
        default="prompt",
        help="코드 생성 모드"
    )
    parser.add_argument(
        "--temp", 
        type=float, 
        default=0.7,
        help="온도 설정"
    )
    parser.add_argument(
        "--no-4bit", 
        action="store_true",
        help="4-bit 양자화 비활성화"
    )
    return parser.parse_args()

def main():
    """메인 함수 (Instruction, Input 지원)"""
    args = parse_arguments()
    
    # 설정 구성
    config = MODEL_CONFIG.copy()
    config["temperature"] = args.temp
    config["model_path"] = args.model
    config["load_in_4bit"] = not args.no_4bit
    
    mode_names = {
        "code-gen": "일반 코드 생성",
        "prompt": "[제2차] 일반 프롬프트 기반 코드 생성",
        "autocomplete": "[제1차] 코드 자동완성", 
        "comment": "[제3차] 주석 기반 코드 생성",
        "error_fix": "[제4차] 오류 코드 설명 및 수정"
    }
    
    print(f"\n🤖 DeepSeek-Coder 코드 어시스트 ({mode_names.get(args.mode, args.mode)})")
    print("📝 종료하려면 'exit' 또는 'quit'을 입력하세요.")
    print("💡 T4 환경에 최적화된 설정으로 실행됩니다.")
    print("💡 Instruction과 Input을 분리해서 입력하려면 'instruction:' 형식을 사용하세요.\n")
    
    # 코드 어시스트 초기화
    assistant = CodeAssistant(config, mode=args.mode)
    
    # REPL 루프
    while True:
        try:
            prompts = {
                "comment": "💬 주석 입력 (코드 생성): ",
                "autocomplete": "💬 코드 입력 (자동완성): ",
                "error_fix": "💬 오류 코드 입력 (수정): ",
                "code-gen": "💬 프롬프트 (instruction: 형식 가능): "
            }
            
            user_input = input(prompts.get(args.mode, prompts["code-gen"]))
                
            if user_input.lower() in ["exit", "quit"]:
                print("👋 코드 어시스트를 종료합니다.")
                break
                
            if not user_input.strip():
                continue
            
            # instruction: 형식 파싱
            parsed_input = user_input
            if "instruction:" in user_input.lower():
                try:
                    parts = user_input.split("instruction:", 1)[1].strip()
                    if "input:" in parts.lower():
                        instruction_part, input_part = parts.split("input:", 1)
                        parsed_input = {
                            "instruction": instruction_part.strip(),
                            "input": input_part.strip()
                        }
                    else:
                        parsed_input = {
                            "instruction": parts.strip(),
                            "input": ""
                        }
                except:
                    # 파싱 실패 시 원본 사용
                    parsed_input = user_input
                
            print("\n⏳ 생성 중...")
            output = assistant.generate(parsed_input, mode=args.mode)
            print(f"\n📄 출력:\n{output}\n")
            
        except KeyboardInterrupt:
            print("\n👋 코드 어시스트를 종료합니다.")
            break
        except Exception as e:
            print(f"\n❌ 오류 발생: {e}")
            if "out of memory" in str(e).lower():
                print("💡 메모리 부족 - 더 짧은 입력을 시도하거나 --no-4bit 옵션을 제거하세요.")

if __name__ == "__main__":
    main()

training_config.env(학습 선택 환경)

# DeepSeek Coder 학습 모드 설정
# 지원 모드: 'complete', 'prompt', 'comment', 'error_fix'
TRAINING_MODE=prompt

config.yaml

# DeepSeek-Coder 6.7B Instruct T4 환경 최적화 설정
model_name: "deepseek-ai/deepseek-coder-6.7b-instruct"  # instruct 모델 사용
data_path: "/home/ubuntu/deepseek-coder/data/train.jsonl"

# T4 최적화 학습 하이퍼파라미터
learning_rate: 2e-4
batch_size: 1
grad_acc: 8
max_length: 1024
num_epochs: 3
lr_scheduler: "cosine"
warmup_ratio: 0.1
weight_decay: 0.01

# T4 호환 최적화 설정
fp16: true
bf16: false
gradient_checkpointing: true
torch_dtype: "float16"

# 출력 및 저장 설정
output_dir: "/home/ubuntu/deepseek-coder/scripts/output/prompt-finetuned/"
save_steps: 10
eval_steps: 50
logging_steps: 5
save_total_limit: 3
evaluation_strategy: "steps"
save_strategy: "steps"
load_best_model_at_end: true
metric_for_best_model: "eval_loss"
greater_is_better: false

# 데이터 처리 설정
dataloader_num_workers: 4
remove_unused_columns: false

# SFT 트레이너 설정
use_sft_trainer: true

# LoRA 설정
use_lora: true
lora_r: 16
lora_alpha: 32
lora_dropout: 0.1
lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]

# 4비트 양자화 설정 (T4 최적화)
load_in_4bit: true
bnb_4bit_compute_dtype: "float16"
bnb_4bit_quant_type: "nf4"
bnb_4bit_use_double_quant: true

# 추가 설정
max_steps: -1
report_to: null

현 학습 상황

2차 학습 완료 기간까지 대기 중