train_model2.py

import pandas as pd
from konlpy.tag import Okt
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.multioutput import MultiOutputClassifier
from sklearn.preprocessing import LabelEncoder
import joblib
from scipy.sparse import hstack
import re
import warnings

# 경고 메시지 무시
warnings.filterwarnings('ignore')


def train_model():
    print("1. 데이터 로드 및 전처리...")
    try:
        # 윈도우 환경에서 엑셀이 기본으로 사용하는 CP949 인코딩을 먼저 시도
        df = pd.read_csv('hackathon.csv', encoding='cp949')
    except UnicodeDecodeError:
        # CP949로 실패하면 UTF-8로 다시 시도
        df = pd.read_csv('hackathon.csv', encoding='utf-8')

    okt = Okt()
    # 모델 학습 시 사용할 불용어 리스트
    stopwords = ['좀', '요', '가', '이', '은', '는', '으로', '와', '과', '에', '을', '를', '도', '고', '다', '아요', '어요']

    def preprocess_text_internal(text):
        # 한글과 공백을 제외한 모든 문자 제거
        text = re.sub(r'[^가-힣\s]', '', text)
        # 형태소 분석 (어간 추출)
        tokens = okt.morphs(text, stem=True)
        # 불용어와 한 글자 단어 제거
        filtered_tokens = [word for word in tokens if not word in stopwords and len(word) > 1]
        return ' '.join(filtered_tokens)

    def convert_to_days(duration):
        if pd.isna(duration) or duration == '':
            return 0
        if isinstance(duration, str):
            if '주' in duration:
                return int(duration.replace('주', '')) * 7
            elif '개월' in duration:
                return int(duration.replace('개월', '')) * 30
            else:
                try:
                    return int(duration)
                except ValueError:
                    return 0
        else:
            return int(duration)

    # 증상 키워드와 신체 부위를 합쳐서 하나의 텍스트로 만듭니다.
    df['processed_text'] = df['증상 키워드'].fillna('') + ' ' + df['신체 부위'].fillna('')
    df['processed_text'] = df['processed_text'].apply(preprocess_text_internal)
    # '기간' 열을 일(day) 단위 숫자로 변환합니다.
    df['기간'] = df['기간'].apply(convert_to_days)

    print("2. TF-IDF 벡터화기 학습...")
    vectorizer = TfidfVectorizer()
    X_text = vectorizer.fit_transform(df['processed_text'])
    X_duration = df[['기간']]
    # 텍스트 벡터와 숫자형 기간 데이터를 결합합니다.
    X_combined = hstack([X_text, X_duration.values])

    print("3. 모델 학습 및 평가...")

    # === 모든 타겟(target) 변수들을 LabelEncoder를 사용해 숫자로 변환합니다. ===
    le_department = LabelEncoder()
    le_emergency = LabelEncoder()
    le_disease = LabelEncoder()

    y_department_encoded = le_department.fit_transform(df['추천 진료과'])
    y_emergency_encoded = le_emergency.fit_transform(df['응급/비응급'])
    y_disease_encoded = le_disease.fit_transform(df['질병이름'])

    # 다중 출력을 위해 여러 타겟 변수를 데이터프레임으로 만듭니다.
    y = pd.DataFrame({
        'department': y_department_encoded,
        'emergency': y_emergency_encoded,
        'disease': y_disease_encoded
    })

    X_train, X_test, y_train, y_test = train_test_split(
        X_combined, y, test_size=0.2, random_state=42
    )

    base_model = RandomForestClassifier(n_estimators=100, random_state=42)
    multi_output_model = MultiOutputClassifier(base_model)
    multi_output_model.fit(X_train, y_train)

    print("4. 모델 학습 및 저장이 완료되었습니다.")
    # 학습된 모델과 라벨 인코더들을 반환합니다.
    return multi_output_model, vectorizer, le_department, le_emergency, le_disease


if __name__ == '__main__':
    multi_output_model, vectorizer, le_department, le_emergency, le_disease = train_model()

    print("모델 및 벡터화기 저장...")
    joblib.dump(multi_output_model, 'multi_task_model.pkl')
    joblib.dump(vectorizer, 'tfidf_vectorizer.pkl')
    joblib.dump(le_department, 'le_department.pkl')
    joblib.dump(le_emergency, 'le_emergency.pkl')
    joblib.dump(le_disease, 'le_disease.pkl')
    print("모델과 라벨 인코더가 성공적으로 저장되었습니다.")