train_model.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.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score
import joblib
import re
import warnings

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

print("1. 데이터 로드 및 전처리...")
try:
    df = pd.read_csv('hackathon_data.csv', encoding='cp949')
except UnicodeDecodeError:
    df = pd.read_csv('hackathon_data.csv', encoding='utf-8')

okt = Okt()
stopwords = ['좀', '요', '가', '이', '은', '는', '으로', '와', '과', '에', '을', '를', '도', '고', '다', '아요', '어요']

def preprocess_text(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)

# 'symptom_keyword'와 'body_parts'를 합쳐서 하나의 텍스트로 만듭니다.
df['processed_text'] = df['symptom_keyword'].fillna('') + ' ' + df['body_parts'].fillna('')
df['processed_text'] = df['processed_text'].apply(preprocess_text)

print("2. TF-IDF 벡터화기 학습...")
vectorizer = TfidfVectorizer()
X = vectorizer.fit_transform(df['processed_text'])

# 타겟(진료과) 데이터 설정
y = df['likely_dept']

print("3. 모델 학습 및 평가...")
# LabelEncoder를 사용해 문자열 타겟을 숫자로 변환합니다.
le = LabelEncoder()
y_encoded = le.fit_transform(y)

X_train, X_test, y_train, y_test = train_test_split(
    X, y_encoded, test_size=0.2, random_state=42
)

model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f"모델 정확도: {accuracy:.2f}")

print("4. 모델 및 벡터화기 저장...")
joblib.dump(model, 'medical_department_model.pkl')
joblib.dump(vectorizer, 'tfidf_vectorizer.pkl')
joblib.dump(le, 'department_label_encoder.pkl')

print("모델 학습 및 저장이 완료되었습니다.")