예제 소스
한국투자증권 API와 LLM 주식 응용 프로그램 · 파일 19개
train_surge_lora.py
# -*- coding: utf-8 -*-
"""급등 손익 분석 JSONL 로 LoRA(SFT) 미세조정.
사전 준비
pip install -r training/requirements_lora.txt
데이터
python training/build_surge_dataset.py --csv surge_paper_trades_YYYY-MM-DD.csv
학습
python training/train_surge_lora.py --data training/data/surge_analysis.jsonl
베이스 모델
- 기본: Hugging Face `Qwen/Qwen2.5-0.5B-Instruct` (첫 실행 시 캐시 다운로드).
- 로컬 폴더: `export BASE_MODEL=/절대/경로/모델` (mywork/ri 등에 복사해 둔 HF 스냅샷).
참고: `mywork/ri` 경로가 없으면 위 HF 베이스 또는 직접 클론한 모델 경로를 쓰면 됩니다.
maLora 쪽 Qwen-Image LoRA 는 이미지용이라 이 스크립트와 호환되지 않습니다.
"""
from __future__ import annotations
import argparse
import os
import sys
from pathlib import Path
HERE = Path(__file__).resolve().parent
DEFAULT_BASE = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
def _peek_model_type(base_model: str) -> str:
"""로컬 경로면 config.json 에서 model_type 을 읽고, 아니면 빈 문자열."""
try:
p = Path(base_model)
if not p.exists():
return ""
cfg = p / "config.json"
if not cfg.is_file():
return ""
import json
data = json.loads(cfg.read_text(encoding="utf-8"))
return str(data.get("model_type") or "")
except Exception:
return ""
def _pick_device() -> str:
try:
import torch
if torch.cuda.is_available():
return "cuda"
if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
return "mps"
except Exception:
pass
return "cpu"
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--data", type=str, default=str(HERE / "data" / "surge_analysis.jsonl"))
ap.add_argument("--base-model", type=str, default=DEFAULT_BASE)
ap.add_argument("--out", type=str, default=str(HERE / "outputs" / "surge_lora"))
ap.add_argument("--epochs", type=float, default=4.0)
ap.add_argument("--lr", type=float, default=2e-4)
ap.add_argument("--batch", type=int, default=1)
ap.add_argument("--grad-accum", type=int, default=8)
ap.add_argument("--max-len", type=int, default=1024)
ap.add_argument("--lora-r", type=int, default=16)
ap.add_argument("--lora-alpha", type=int, default=32)
args = ap.parse_args()
data_path = Path(args.data)
if not data_path.is_file():
print(f"데이터 없음: {data_path} — 먼저 training/build_surge_dataset.py 실행.", file=sys.stderr)
return 1
mt = _peek_model_type(args.base_model)
if mt and ("_vl" in mt or "idefics" in mt):
print(
"선택한 베이스 모델이 VLM(vision-language) 계열로 보입니다. "
"현재 스크립트는 텍스트 CausalLM LoRA(SFT) 전용입니다.\n"
f" base_model={args.base_model}\n"
f" model_type={mt}\n"
"텍스트 전용 모델(예: Qwen2.5 Instruct)로 학습하거나, VLM 전용 학습 스크립트로 분리해야 합니다.",
file=sys.stderr,
)
return 2
try:
import inspect
import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import SFTTrainer
try:
from trl import SFTConfig
except ImportError:
SFTConfig = None # type: ignore[misc, assignment]
except ImportError as e:
print(f"패키지 필요: {e}\n pip install -r training/requirements_lora.txt", file=sys.stderr)
return 1
device_t = _pick_device()
dtype = torch.bfloat16 if device_t == "cuda" else torch.float32
if device_t == "mps":
dtype = torch.float32 # MPS 에서 bfloat16 일부 불안정할 수 있음
tokenizer = AutoTokenizer.from_pretrained(args.base_model, trust_remote_code=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# 4bit 는 CUDA + bitsandbytes 설치 시만 (선택)
bnb_config = None
if device_t == "cuda" and os.environ.get("USE_QLORA", "") == "1":
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4")
model = AutoModelForCausalLM.from_pretrained(
args.base_model,
trust_remote_code=True,
quantization_config=bnb_config,
device_map="auto" if bnb_config else None,
torch_dtype=dtype if bnb_config is None else None,
)
if bnb_config is None:
model = model.to(device_t)
model.config.use_cache = False
try:
model.enable_input_require_grads()
except Exception:
pass
peft_config = LoraConfig(
r=args.lora_r,
lora_alpha=args.lora_alpha,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
)
ds = load_dataset("json", data_files=str(data_path), split="train")
def _rows_to_messages(batch: dict) -> dict:
instr = batch["instruction"]
inp = batch["input"]
out = batch["output"]
if isinstance(instr, str):
instr, inp, out = [instr], [inp], [out]
msgs: list[list[dict]] = []
for i in range(len(instr)):
msgs.append(
[
{"role": "user", "content": f"{instr[i]}\n\n{inp[i]}"},
{"role": "assistant", "content": out[i]},
]
)
return {"messages": msgs}
out_dir = Path(args.out)
out_dir.mkdir(parents=True, exist_ok=True)
def formatting_func(examples: dict) -> list[str]:
"""구 TRL: 배치 → 프롬프트 문자열 리스트."""
instr = examples["instruction"]
inp = examples["input"]
out = examples["output"]
if isinstance(instr, str):
instr, inp, out = [instr], [inp], [out]
texts: list[str] = []
for i in range(len(instr)):
messages = [
{"role": "user", "content": f"{instr[i]}\n\n{inp[i]}"},
{"role": "assistant", "content": out[i]},
]
try:
t = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=False,
)
except Exception:
t = f"### 질문\n{instr[i]}\n{inp[i]}\n### 답변\n{out[i]}"
texts.append(t)
return texts
sig = inspect.signature(SFTTrainer.__init__)
sft_params = set(sig.parameters)
# ── TRL ≥ 0.26 계열: SFTConfig + processing_class, max_length 는 Config 쪽
if SFTConfig is not None and "processing_class" in sft_params:
ds_conv = ds.map(
_rows_to_messages,
batched=True,
remove_columns=list(ds.column_names),
)
sft_args = dict(
output_dir=str(out_dir),
num_train_epochs=args.epochs,
per_device_train_batch_size=args.batch,
gradient_accumulation_steps=args.grad_accum,
learning_rate=args.lr,
lr_scheduler_type="cosine",
warmup_ratio=0.06,
logging_steps=1,
save_strategy="epoch",
fp16=False,
bf16=bool(device_t == "cuda" and dtype == torch.bfloat16),
optim="adamw_torch",
report_to="none",
gradient_checkpointing=True,
max_length=args.max_len,
assistant_only_loss=True,
)
sft_config = SFTConfig(**sft_args)
trainer = SFTTrainer(
model=model,
args=sft_config,
train_dataset=ds_conv,
processing_class=tokenizer,
peft_config=peft_config,
)
else:
# ── 구 TRL: TrainingArguments + tokenizer / max_seq_length 조합
targs = TrainingArguments(
output_dir=str(out_dir),
num_train_epochs=args.epochs,
per_device_train_batch_size=args.batch,
gradient_accumulation_steps=args.grad_accum,
learning_rate=args.lr,
lr_scheduler_type="cosine",
warmup_ratio=0.06,
logging_steps=1,
save_strategy="epoch",
fp16=False,
bf16=device_t == "cuda" and dtype == torch.bfloat16,
optim="adamw_torch",
report_to="none",
gradient_checkpointing=True,
)
kwargs: dict = dict(
model=model,
args=targs,
train_dataset=ds,
peft_config=peft_config,
formatting_func=formatting_func,
)
if "processing_class" in sft_params:
kwargs["processing_class"] = tokenizer
elif "tokenizer" in sft_params:
kwargs["tokenizer"] = tokenizer
else:
print("TRL SFTTrainer 시그니처를 알 수 없습니다. trl/transformers 버전을 확인하세요.", file=sys.stderr)
return 1
trainer = None
for extra in (
{"max_seq_length": args.max_len},
{"max_length": args.max_len},
{},
):
try:
trainer = SFTTrainer(**kwargs, **extra)
break
except TypeError:
continue
if trainer is None:
print(
"SFTTrainer 생성에 실패했습니다. `pip install -U trl` 또는 training/requirements_lora.txt 버전을 맞춰 주세요.",
file=sys.stderr,
)
return 1
trainer.train()
trainer.save_model(str(out_dir))
tokenizer.save_pretrained(str(out_dir))
print(f"[ok] LoRA 저장: {out_dir.resolve()}")
return 0
if __name__ == "__main__":
raise SystemExit(main())