File size: 996 Bytes
d8a1c8d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import classification_report
import joblib
from config import MODEL_SAVE_PATH

def get_classifier(model_type):
    if model_type == "logistic":
        return LogisticRegression(max_iter=2000)
    elif model_type == "random_forest":
        return RandomForestClassifier(n_estimators=100)
    elif model_type == "mlp":
        return MLPClassifier(hidden_layer_sizes=(512, 256), max_iter=300)
    else:
        raise ValueError(f"Unknown model type: {model_type}")

def train_classifier(X_train, y_train, model_type):
    clf = get_classifier(model_type)
    clf.fit(X_train, y_train)
    joblib.dump(clf, MODEL_SAVE_PATH)
    return clf

def evaluate_classifier(model, X_test, y_test, label_encoder):
    y_pred = model.predict(X_test)
    print(classification_report(y_test, y_pred, target_names=label_encoder.classes_))