import pandas as pd | |
from sklearn.feature_extraction.text import TfidfVectorizer | |
from sklearn.linear_model import LogisticRegression | |
from sklearn.model_selection import train_test_split | |
from sklearn.metrics import roc_auc_score | |
# Load the data | |
train_data = pd.read_csv("./input/train.csv") | |
# Prepare the features and labels | |
X = train_data["comment_text"] | |
y = train_data.iloc[:, 2:] | |
# Split the data into training and validation sets | |
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42) | |
# Create TF-IDF features | |
tfidf_vectorizer = TfidfVectorizer(max_features=10000, stop_words="english") | |
X_train_tfidf = tfidf_vectorizer.fit_transform(X_train) | |
X_val_tfidf = tfidf_vectorizer.transform(X_val) | |
# Train a logistic regression model for each label | |
scores = [] | |
for label in y.columns: | |
lr = LogisticRegression(C=1.0, solver="liblinear") | |
lr.fit(X_train_tfidf, y_train[label]) | |
y_pred = lr.predict_proba(X_val_tfidf)[:, 1] | |
score = roc_auc_score(y_val[label], y_pred) | |
scores.append(score) | |
print(f"ROC AUC for {label}: {score}") | |
# Calculate the mean column-wise ROC AUC | |
mean_auc = sum(scores) / len(scores) | |
print(f"Mean column-wise ROC AUC: {mean_auc}") | |