Spaces:
Running on Zero
Running on Zero
File size: 2,837 Bytes
1ac979a | 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 | import os
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.naive_bayes import MultinomialNB
from sklearn.metrics import classification_report, accuracy_score
import joblib
from src.preprocess import clean_text
def train_and_evaluate(data_path, text_col='text', label_col='label', models_dir='models'):
"""
Trains models for Hate Speech Detection and saves the best one.
Assumes data_path points to a CSV with a text column and a label column.
"""
print(f"Loading data from {data_path}...")
df = pd.read_csv(data_path)
# Preprocessing
print("Cleaning text data... this may take a moment.")
df['clean_text'] = df[text_col].apply(clean_text)
X = df['clean_text']
y = df[label_col]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# Feature extraction using TF-IDF
print("Vectorizing text using TF-IDF...")
vectorizer = TfidfVectorizer(max_features=5000, ngram_range=(1, 2))
X_train_vec = vectorizer.fit_transform(X_train)
X_test_vec = vectorizer.transform(X_test)
# Train Logistic Regression
print("Training Logistic Regression...")
lr_model = LogisticRegression(max_iter=1000)
lr_model.fit(X_train_vec, y_train)
lr_preds = lr_model.predict(X_test_vec)
lr_acc = accuracy_score(y_test, lr_preds)
print(f"Logistic Regression Accuracy: {lr_acc:.4f}")
print(classification_report(y_test, lr_preds))
# Train Naive Bayes
print("Training Naive Bayes...")
nb_model = MultinomialNB()
nb_model.fit(X_train_vec, y_train)
nb_preds = nb_model.predict(X_test_vec)
nb_acc = accuracy_score(y_test, nb_preds)
print(f"Naive Bayes Accuracy: {nb_acc:.4f}")
print(classification_report(y_test, nb_preds))
# Save the best model
os.makedirs(models_dir, exist_ok=True)
# In this case, we prefer Logistic Regression for saving as it usually performs well and is robust
best_model = lr_model if lr_acc >= nb_acc else nb_model
model_name = 'logistic_regression.pkl' if lr_acc >= nb_acc else 'naive_bayes.pkl'
joblib.dump(best_model, os.path.join(models_dir, 'model.pkl'))
joblib.dump(vectorizer, os.path.join(models_dir, 'vectorizer.pkl'))
print(f"Best model ({model_name}) and vectorizer saved to '{models_dir}/'.")
if __name__ == "__main__":
# Example usage (will need dummy data to run directly)
default_data_path = os.path.join('data', 'sample.csv')
if os.path.exists(default_data_path):
train_and_evaluate(default_data_path)
else:
print(f"Data not found at {default_data_path}. Please create a sample dataset first.")
|