Spaces:
Sleeping
Sleeping
File size: 2,156 Bytes
492754f | 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 | import os
import mlflow
import mlflow.sklearn
import pandas as pd
import joblib
from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.svm import LinearSVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, f1_score, classification_report
train_df = pd.read_csv("data/splits/train.csv")
val_df = pd.read_csv("data/splits/val.csv")
X_train = train_df["clean_text"]
y_train = train_df["label_text"]
X_val = val_df["clean_text"]
y_val = val_df["label_text"]
os.makedirs("models", exist_ok=True)
models = {
"logistic_regression": LogisticRegression(max_iter=1000),
"linear_svm": LinearSVC(),
"random_forest": RandomForestClassifier(n_estimators=100, random_state=42),
}
mlflow.set_experiment("bbc-document-classification")
best_model = None
best_model_name = None
best_f1 = 0
for model_name, classifier in models.items():
with mlflow.start_run(run_name=model_name):
pipeline = Pipeline([
("tfidf", TfidfVectorizer(max_features=5000)),
("classifier", classifier)
])
pipeline.fit(X_train, y_train)
y_pred = pipeline.predict(X_val)
accuracy = accuracy_score(y_val, y_pred)
f1 = f1_score(y_val, y_pred, average="weighted")
print("\n==============================")
print(f"Model: {model_name}")
print("Accuracy:", accuracy)
print("F1 Score:", f1)
print(classification_report(y_val, y_pred))
mlflow.log_param("model_name", model_name)
mlflow.log_param("vectorizer", "TF-IDF")
mlflow.log_param("max_features", 5000)
mlflow.log_metric("accuracy", accuracy)
mlflow.log_metric("f1_score", f1)
mlflow.sklearn.log_model(pipeline, model_name)
if f1 > best_f1:
best_f1 = f1
best_model = pipeline
best_model_name = model_name
joblib.dump(best_model, "models/best_model.pkl")
print("\nBest model:", best_model_name)
print("Best F1:", best_f1)
print("Saved to models/best_model.pkl") |