Files
layoutlmv3_service/train/classifier_trainer.py
2026-06-15 09:54:01 +09:00

215 lines
8.7 KiB
Python

"""
classifier_dataset.load() 로 데이터 로드 후 TF-IDF + SVM 학습
"""
import random
from collections import Counter
from loguru import logger
from sklearn.calibration import CalibratedClassifierCV
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import classification_report
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.svm import LinearSVC
from config import SERVICE
from train import classifier_dataset
from utils import classifier_model_store
_status: dict = {"status": "idle", "message": ""}
def get_status() -> dict:
return _status
def augment_text(text: str, n: int = 25) -> list:
"""텍스트 증강 - 단어 순서 변경, 일부 단어 제거"""
words = text.split()
augmented = []
for _ in range(n):
aug = words.copy()
# 1. 랜덤 단어 10% 제거
drop_count = max(1, int(len(aug) * 0.1))
for _ in range(drop_count):
if aug:
aug.pop(random.randint(0, len(aug) - 1))
# 2. 랜덤 단어 일부 순서 섞기
if len(aug) > 5:
idx = random.randint(0, len(aug) - 3)
aug[idx], aug[idx + 1] = aug[idx + 1], aug[idx]
augmented.append(" ".join(aug))
return augmented
def run_train() -> None:
global _status
try:
_status = {"status": "training", "message": "데이터셋 로딩 중..."}
data = classifier_dataset.load()
if not data:
raise ValueError("데이터셋이 없습니다. /classifier/dataset/build 먼저 실행하세요.")
texts = [d["text"] for d in data]
labels = [d["label"] for d in data]
# ── 특정 클래스 데이터 증강 ──────────────────────
if SERVICE == "amko":
AUGMENT_LABELS = ["AMKOR TECHNOLOGY TAIWAN_B"] # 증강할 클래스
AUGMENT_TARGET = 30 # 목표 샘플 수
augmented_texts, augmented_labels = [], []
label_groups = {}
for t, l in zip(texts, labels):
label_groups.setdefault(l, []).append(t)
for label in AUGMENT_LABELS:
samples = label_groups.get(label, [])
current_count = len(samples)
if current_count < AUGMENT_TARGET:
need = AUGMENT_TARGET - current_count
logger.info(f"{label}: {current_count}개 → {AUGMENT_TARGET}개 증강 ({need}개 추가)")
for i in range(need):
base_text = samples[i % current_count] # 기존 샘플 순환
aug = augment_text(base_text, n=1)[0]
augmented_texts.append(aug)
augmented_labels.append(label)
texts = texts + augmented_texts
labels = labels + augmented_labels
logger.info(f"증강 후 클래스별 샘플: {Counter(labels)}")
# ────────────────────────────────────────────────
# # ── 클래스 언더샘플링 ─────────────────────────
# MAX_SAMPLES_PER_CLASS = 20 # B타입 8개의 2~3배 수준으로 제한
#
# balanced_data = []
# class_counts = Counter(labels)
#
# for label in set(labels):
# class_texts = [t for t, l in zip(texts, labels) if l == label]
# # MAX_SAMPLES_PER_CLASS 초과 시 랜덤 샘플링
# if len(class_texts) > MAX_SAMPLES_PER_CLASS:
# class_texts = random.sample(class_texts, MAX_SAMPLES_PER_CLASS)
# balanced_data.extend([(t, label) for t in class_texts])
#
# texts, labels = zip(*balanced_data)
# texts, labels = list(texts), list(labels)
# logger.info(f"언더샘플링 후: {Counter(labels)}")
# # ── 핵심 키워드 가중치 부여 (특정 클래스만) ──────────────────────
# KEYWORD_BOOST_BY_LABEL = {
# "AMKOR TECHNOLOGY TAIWAN_A": { # A타입에만
# "Process Fee": 5,
# "Total Process Fee": 5,
# "Consigned Value": 5,
# "Total Process Value": 5,
# },
# "AMKOR TECHNOLOGY TAIWAN_B": { # B타입에만
# "PACKING LIST SUMMARY": 5,
# "FOC Y": 3,
# },
# }
#
# def boost_keywords(text: str, label: str) -> str:
# boost_map = KEYWORD_BOOST_BY_LABEL.get(label, {}) # 해당 label 없으면 빈 dict
# boosted = text
# for keyword, repeat in boost_map.items():
# if keyword.lower() in text.lower():
# boosted += f" {keyword}" * repeat
# return boosted
#
# texts = [boost_keywords(t, l) for t, l in zip(texts, labels)]
# logger.info("핵심 키워드 가중치 부여 완료")
# # ────────────────────────────────────────────────
# # ── 샘플 부족 클래스 필터링 ──────────────────────
# counts = Counter(labels)
# MIN_SAMPLES = 5
# before = len(texts)
# filtered = [(t, l) for t, l in zip(texts, labels) if counts[l] >= MIN_SAMPLES]
# if not filtered:
# raise ValueError("모든 클래스가 샘플 부족으로 제외되었습니다.")
# texts, labels = map(list, zip(*filtered))
# excluded = {k: v for k, v in counts.items() if v < MIN_SAMPLES}
# if excluded:
# logger.warning(f"샘플 부족 제외: {excluded}")
# logger.info(f"필터링: {before}개 → {len(texts)}개")
# # ────────────────────────────────────────────────
#
# counts = Counter(labels)
# min_count = min(counts.values())
# cv = 5
# stratify = labels if min_count >= cv else None
MIN_SAMPLES = 5
# ── 필터링 ───────────────────────────────────────
counts = Counter(labels)
before = len(texts)
filtered = [(t, l) for t, l in zip(texts, labels) if counts[l] >= MIN_SAMPLES]
if not filtered:
raise ValueError("모든 클래스가 샘플 부족으로 제외되었습니다.")
texts, labels = map(list, zip(*filtered))
excluded = {k: v for k, v in counts.items() if v < MIN_SAMPLES}
if excluded:
logger.warning(f"샘플 부족 제외: {excluded}")
logger.info(f"필터링: {before}개 → {len(texts)}")
# ── MIN_SAMPLES 기준으로 cv/stratify 자동 결정 ──
counts = Counter(labels)
min_count = min(counts.values())
cv = max(2, min(5, int(MIN_SAMPLES * 0.8))) # MIN_SAMPLES의 80% (train 비율)
stratify = labels if min_count >= cv else None
logger.info(f"MIN_SAMPLES={MIN_SAMPLES} → cv={cv}, stratify={'적용' if stratify else '비적용'}")
_status["message"] = f"{len(texts)}개 샘플 / {len(counts)}개 클래스 학습 중..."
_status["samples"] = len(texts)
_status["classes"] = len(counts)
logger.info(_status["message"])
X_train, X_test, y_train, y_test = train_test_split(
texts, labels,
test_size=0.2,
stratify=stratify,
random_state=42,
)
# train_counts = Counter(y_train)
# min_train_count = min(train_counts.values())
# cv = max(2, min(5, min_train_count))
pipeline = Pipeline([
("tfidf", TfidfVectorizer(max_features=50000, ngram_range=(1, 2))),
("clf", CalibratedClassifierCV(LinearSVC(max_iter=2000, class_weight="balanced"), cv=cv)),
])
pipeline.fit(X_train, y_train)
report = classification_report(y_test, pipeline.predict(X_test), zero_division=0)
logger.info(f"\n{report}")
classifier_model_store.set_model(pipeline)
_status = {
"status": "done",
"message": f"학습 완료 | 샘플={len(texts)}, 클래스={len(counts)}",
"samples": len(texts),
"classes": len(counts),
"report": report, # 필요 시 별도 API로 노출
}
logger.info("classifier 학습 완료")
except Exception as e:
_status = {"status": "error", "message": str(e)}
logger.exception(f"classifier 학습 실패: {e}")