Initial commit

train/classifier_trainer.py

@@ -0,0 +1,215 @@
+"""
+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}")