clip-model / model.py
JonSnow1512's picture
Upload 5 files
d8a1c8d verified
raw
history blame contribute delete
996 Bytes
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_))