File size: 4,206 Bytes
af6ff2d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """
sentiment.py
모듈화된 감성 분류기
함수:
- load_sentiment_model: 저장된 KcELECTRA 모델/토크나이저 로드
- predict_sentiment_df: 분리된 문장 DataFrame에 감성 추론 후 라벨 및 확률 컬럼 추가
CLI:
--input : 입력 CSV 경로 (ID, 작성시간, cleaned, divided_comment 포함)
--output : 출력 CSV 경로
--model-dir : HuggingFace 포맷 모델 디렉토리 (config.json, tf_model.h5, vocab 등)
--text-col : 분류할 텍스트 컬럼 (기본: 'divided_comment')
--output-col : 예측 라벨 컬럼명 (기본: 'sentiment')
--max-length : 토큰 최대 길이 (기본: 64)
--batch-size : 배치 크기 (기본: 16)
--preview : True일 때 샘플 10개 확인
"""
import os
import time
import argparse
import pandas as pd
import numpy as np
import tensorflow as tf
from transformers import ElectraTokenizer, TFElectraForSequenceClassification
from ace_tools_open import display_dataframe_to_user
def load_sentiment_model(model_dir: str):
"""
저장된 KcELECTRA 모델/토크나이저 로드
model_dir 내부에 config.json, tf_model.h5, vocab.txt, tokenizer_config.json,
special_tokens_map.json 등이 있어야 함
"""
tokenizer = ElectraTokenizer.from_pretrained(model_dir)
model = TFElectraForSequenceClassification.from_pretrained(
model_dir,
num_labels=2
)
return tokenizer, model
def predict_sentiment_df(
df: pd.DataFrame,
tokenizer,
model,
text_col: str = 'divided_comment',
output_col: str = 'sentiment',
max_length: int = 64,
batch_size: int = 16
) -> pd.DataFrame:
"""
분리된 문장 DataFrame(df)에 대해 감성 분류 수행
- output_col: 예측 라벨 컬럼명
- output_col_prob_0, output_col_prob_1: softmax 확률 컬럼 추가
"""
texts = df[text_col].fillna('').astype(str).tolist()
enc = tokenizer(
texts,
return_tensors='tf',
padding='max_length',
truncation=True,
max_length=max_length
)
ds = tf.data.Dataset.from_tensor_slices(dict(enc)).batch(batch_size)
start = time.perf_counter()
preds = model.predict(ds)
elapsed = time.perf_counter() - start
probs = tf.nn.softmax(preds.logits, axis=1).numpy()
labels = np.argmax(probs, axis=1)
df[output_col] = labels
df[f'{output_col}_prob_0'] = probs[:, 0]
df[f'{output_col}_prob_1'] = probs[:, 1]
print(f"[SENTI] {len(texts)} samples → {elapsed:.2f}s total, {elapsed/len(texts):.4f}s per sample")
return df
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='감성 분류 모듈')
parser.add_argument('--input', '-i', required=True, help='입력 CSV 파일 경로')
parser.add_argument('--output', '-o', required=True, help='출력 CSV 파일 경로')
parser.add_argument('--model-dir', '-m', required=True, help='모델 디렉토리 경로')
parser.add_argument('--text-col', default='divided_comment', help='분류할 텍스트 컬럼')
parser.add_argument('--output-col', default='sentiment', help='예측 라벨 컬럼명')
parser.add_argument('--max-length', type=int, default=64, help='토큰 최대 길이')
parser.add_argument('--batch-size', type=int, default=16, help='배치 크기')
parser.add_argument('--preview', action='store_true', help='샘플 확인')
args = parser.parse_args()
df = pd.read_csv(args.input, encoding='utf-8-sig')
tokenizer, model = load_sentiment_model(args.model_dir)
df_out = predict_sentiment_df(
df,
tokenizer,
model,
text_col=args.text_col,
output_col=args.output_col,
max_length=args.max_length,
batch_size=args.batch_size
)
if args.preview:
display_dataframe_to_user('감성 분류 예시', df_out.head(10))
os.makedirs(os.path.dirname(args.output), exist_ok=True)
df_out.to_csv(args.output, index=False, encoding='utf-8-sig')
print(f"[SENTI] 저장 완료: {args.output}")
|