E5 Action-Family Classifier for Dacon

Dacon Action 분류 데이터를 이용해 파인튜닝한 계층형 분류 모델입니다.

이 모델은 최종 Action 예측뿐 아니라 E5 embedding과 각 분류 head의 logits·확률을 제공합니다. 사용자는 추출한 특징 위에 LightGBM, XGBoost, RandomForest 등의 분류기를 추가로 학습할 수 있습니다.

모델 구조

Shared multilingual-e5-small encoder
├─ Flat Action Head: 14 classes
├─ Family Head: 4 classes
└─ Specialist Heads
   ├─ file_discovery
   ├─ file_edit
   ├─ execute_validate
   └─ interaction

최종 모델 확률은 다음 Soft Routing으로 계산합니다.

P_hierarchical(action)
= P(family) × P(action | family)

P_final(action)
= α × P_flat(action)
+ (1-α) × P_hierarchical(action)

지원 클래스

저장소 구조

repository/
├─ config.json
├─ inference.py
├─ label_mapping.json
├─ model.safetensors
├─ routing_config.json
├─ sentencepiece.bpe.model
├─ special_tokens_map.json
├─ tokenizer.json
└─ tokenizer_config.json

1. 라이브러리 설치

pip install torch transformers==4.57.6 safetensors sentencepiece huggingface_hub

LightGBM을 추가로 학습하려면 다음 패키지도 설치합니다.

pip install lightgbm scikit-learn pandas joblib

2. 모델 불러오기

import sys

from huggingface_hub import snapshot_download

repo_id = "RiverWon/e5-action-family-dacon"

repo_dir = snapshot_download(
    repo_id=repo_id
)

# 다운로드된 inference.py를 import할 수 있도록 경로 추가
sys.path.insert(0, repo_dir)

from inference import ActionFamilyFeatureExtractor

extractor = ActionFamilyFeatureExtractor.from_pretrained(
    repo_dir
)

print("모델 로드 완료")
print("device:", extractor.device)

snapshot_download()는 모델 파일과 inference.py를 함께 내려받습니다.

3. 단일 샘플 예측

sample = {
    "current_prompt": "수정한 코드가 정상인지 테스트를 실행해줘",
    "history": [
        {
            "role": "assistant_action",
            "name": "edit_file",
            "args": {
                "path": "src/app.py"
            },
            "result_summary": "파일 수정 완료"
        }
    ],
    "session_meta": {
        "user_tier": "pro",
        "language_pref": "ko",
        "turn_index": 3,
        "workspace": {
            "open_files": ["src/app.py"],
            "language_mix": {
                "py": 1.0
            },
            "git_dirty": True,
            "last_ci_status": "none"
        }
    }
}

result = extractor.predict(sample)

print(result["action"])
print(result["family"])
print(result["confidence"])
print(result["top_k"])

출력 예시:

{
    "action": "run_tests",
    "family": "execute_validate",
    "confidence": 0.82,
    "top_k": [
        {
            "action": "run_tests",
            "family": "execute_validate",
            "probability": 0.82
        },
        {
            "action": "run_bash",
            "family": "execute_validate",
            "probability": 0.09
        }
    ]
}

4. 여러 샘플 예측

samples = [
    {
        "current_prompt": "프로젝트에서 UserService를 찾아줘",
        "history": [],
        "session_meta": {"workspace": {}}
    },
    {
        "current_prompt": "테스트를 실행해줘",
        "history": [],
        "session_meta": {"workspace": {}}
    }
]

results = extractor.predict_batch(
    samples,
    batch_size=32,
    top_k=3
)

for result in results:
    print(result["action"], result["confidence"])

5. 모델 특징 추출

후속 분류기를 학습할 수 있도록 모델의 중간 특징을 추출할 수 있습니다.

features = extractor.extract_features(
    samples,
    batch_size=32
)

print(features["embedding"].shape)
print(features["flat_logits"].shape)
print(features["family_logits"].shape)
print(features["specialist_logits"].shape)
print(features["hierarchical_prob"].shape)
print(features["final_prob"].shape)

제공되는 특징:

transform()을 사용하면 후속 모델에 바로 넣을 수 있는 하나의 행렬을 반환합니다.

X = extractor.transform(
    samples,
    batch_size=32
)

print(X.shape)  # (샘플 수, 444)

6. Dacon 데이터 준비

다음과 같은 디렉터리 구조를 사용합니다.

project/
├─ data/
│  ├─ train.jsonl
│  ├─ train_labels.csv
│  ├─ test.jsonl
│  └─ sample_submission.csv
└─ train_and_submit.py

데이터 로더:

import json
from pathlib import Path

import pandas as pd

DATA_DIR = Path("./data")


def load_jsonl(path):
    with path.open(encoding="utf-8") as file:
        return [
            json.loads(line)
            for line in file
            if line.strip()
        ]


train_samples = load_jsonl(
    DATA_DIR / "train.jsonl"
)

test_samples = load_jsonl(
    DATA_DIR / "test.jsonl"
)

labels_df = pd.read_csv(
    DATA_DIR / "train_labels.csv"
)

label_by_id = dict(
    zip(labels_df["id"], labels_df["action"])
)

y_text = [
    label_by_id[sample["id"]]
    for sample in train_samples
]

print("train:", len(train_samples))
print("test:", len(test_samples))

7. 모델만 사용해 Dacon 제출

추가 분류기를 학습하지 않고 모델의 기본 예측을 그대로 제출할 수 있습니다.

import pandas as pd

test_results = extractor.predict_batch(
    test_samples,
    batch_size=64
)

prediction_by_id = {
    sample["id"]: result["action"]
    for sample, result in zip(
        test_samples,
        test_results
    )
}

submission = pd.read_csv(
    DATA_DIR / "sample_submission.csv"
)

submission["action"] = submission["id"].map(
    prediction_by_id
)

if submission["action"].isna().any():
    raise ValueError("일부 test ID의 예측값이 없습니다.")

submission.to_csv(
    "submission_e5_actionfamily.csv",
    index=False
)

print(submission.head())
print("saved: submission_e5_actionfamily.csv")

생성된 파일:

submission_e5_actionfamily.csv

이 파일을 Dacon 대회 페이지에 제출합니다.

8. LightGBM으로 성능 강화

특징 추출

import numpy as np

X_train = extractor.transform(
    train_samples,
    batch_size=64
)

X_test = extractor.transform(
    test_samples,
    batch_size=64
)

print("X_train:", X_train.shape)
print("X_test:", X_test.shape)

np.save("X_train_e5_actionfamily.npy", X_train)
np.save("X_test_e5_actionfamily.npy", X_test)

특징 추출에는 GPU 사용을 권장합니다. 저장된 NumPy 파일은 이후 다시 모델을 실행하지 않고 재사용할 수 있습니다.

라벨 인코딩

import joblib

from sklearn.preprocessing import LabelEncoder

label_encoder = LabelEncoder()
y = label_encoder.fit_transform(y_text)

joblib.dump(
    label_encoder,
    "label_encoder.pkl"
)

print(label_encoder.classes_)

간단한 로컬 검증

세션 데이터이므로 ID에서 -step_N을 제거한 값을 group으로 사용합니다.

import re
import numpy as np

from sklearn.model_selection import GroupShuffleSplit
from sklearn.metrics import f1_score


def get_session_id(sample_id):
    return re.sub(
        r"-step_\d+$",
        "",
        sample_id
    )


groups = np.array([
    get_session_id(sample["id"])
    for sample in train_samples
])

splitter = GroupShuffleSplit(
    n_splits=1,
    test_size=0.2,
    random_state=42
)

train_idx, valid_idx = next(
    splitter.split(
        X_train,
        y,
        groups=groups
    )
)

LightGBM 학습

import lightgbm as lgb

classifier = lgb.LGBMClassifier(
    objective="multiclass",
    n_estimators=1500,
    learning_rate=0.03,
    num_leaves=31,
    max_depth=-1,
    subsample=0.9,
    colsample_bytree=0.9,
    class_weight="balanced",
    random_state=42,
    n_jobs=-1
)

classifier.fit(
    X_train[train_idx],
    y[train_idx],
    eval_set=[
        (
            X_train[valid_idx],
            y[valid_idx]
        )
    ],
    callbacks=[
        lgb.early_stopping(100),
        lgb.log_evaluation(100)
    ]
)

valid_prediction = classifier.predict(
    X_train[valid_idx]
)

macro_f1 = f1_score(
    y[valid_idx],
    valid_prediction,
    average="macro"
)

print("validation Macro F1:", macro_f1)

전체 데이터 재학습

Validation에서 선택된 boosting round를 사용합니다.

best_iteration = (
    classifier.best_iteration_
    if classifier.best_iteration_
    else classifier.n_estimators
)

final_classifier = lgb.LGBMClassifier(
    objective="multiclass",
    n_estimators=best_iteration,
    learning_rate=0.03,
    num_leaves=31,
    max_depth=-1,
    subsample=0.9,
    colsample_bytree=0.9,
    class_weight="balanced",
    random_state=42,
    n_jobs=-1
)

final_classifier.fit(
    X_train,
    y
)

joblib.dump(
    final_classifier,
    "lightgbm_action_classifier.pkl"
)

Test 예측 및 제출 파일 생성

test_prediction_ids = final_classifier.predict(
    X_test
)

test_predictions = label_encoder.inverse_transform(
    test_prediction_ids.astype(int)
)

prediction_by_id = {
    sample["id"]: prediction
    for sample, prediction in zip(
        test_samples,
        test_predictions
    )
}

submission = pd.read_csv(
    DATA_DIR / "sample_submission.csv"
)

submission["action"] = submission["id"].map(
    prediction_by_id
)

if submission["action"].isna().any():
    raise ValueError("일부 test ID의 예측값이 없습니다.")

submission.to_csv(
    "submission_e5_lightgbm.csv",
    index=False
)

print(submission.head())
print("saved: submission_e5_lightgbm.csv")

생성된 파일:

submission_e5_lightgbm.csv

이 파일을 Dacon에 제출합니다.

9. 원하는 특징만 선택

Embedding만 사용할 수 있습니다.

X = extractor.transform(
    samples,
    include=["embedding"]
)

모델 logits만 사용할 수 있습니다.

X = extractor.transform(
    samples,
    include=[
        "flat_logits",
        "family_logits",
        "specialist_logits"
    ]
)

확률 특징만 사용할 수 있습니다.

X = extractor.transform(
    samples,
    include=[
        "hierarchical_prob",
        "final_prob"
    ]
)

권장 시작 설정:

X = extractor.transform(
    samples,
    include=[
        "embedding",
        "flat_logits",
        "family_logits",
        "specialist_logits",
        "hierarchical_prob",
        "final_prob"
    ]
)

10. 학습한 분류기 다시 사용

import joblib
import numpy as np

classifier = joblib.load(
    "lightgbm_action_classifier.pkl"
)

label_encoder = joblib.load(
    "label_encoder.pkl"
)

X_new = extractor.transform(
    new_samples,
    batch_size=64
)

prediction_ids = classifier.predict(
    X_new
).astype(int)

predictions = label_encoder.inverse_transform(
    prediction_ids
)

print(predictions)

주의사항

• 이 모델은 일반적인 AutoModelForSequenceClassification 모델이 아닙니다.
• 제공된 inference.py의 ActionFamilyFeatureExtractor를 이용해 로드해야 합니다.
• 특징 추출 시 학습 때와 동일한 입력 전처리를 사용해야 합니다.
• 입력 형식이 학습 데이터와 크게 다르면 성능이 낮아질 수 있습니다.
• 단일 Decision Tree는 쉽게 과적합할 수 있으므로 LightGBM이나 XGBoost를 권장합니다.
• 클래스 불균형이 있으므로 Accuracy뿐 아니라 Macro F1을 확인하세요.
• 세션형 데이터의 로컬 검증은 행 단위가 아니라 session 단위 분할을 권장합니다.
• 공개 모델이 Dacon 전체 학습 데이터로 파인튜닝됐다면, 같은 학습 데이터에서 계산한 로컬 검증 점수는 낙관적일 수 있습니다.
• Dacon 데이터 및 대회 규정에 따라 모델과 외부 데이터 사용 가능 여부를 직접 확인하세요.

모델 정보

인용 및 라이선스

기반 모델 intfloat/multilingual-e5-small의 모델 카드와 라이선스를 확인해주세요.

본 모델을 Dacon 대회에 사용하는 경우 해당 대회의 데이터·외부 모델·제출 규정을 우선적으로 준수해야 합니다.