40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
import pandas as pd
|
|
import numpy as np
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from sklearn.naive_bayes import MultinomialNB
|
|
from sklearn.model_selection import train_test_split
|
|
import joblib
|
|
|
|
# Load the dataset
|
|
print("Loading dataset...")
|
|
df = pd.read_csv('dataset.csv')
|
|
|
|
# Assuming the columns are 'content' and 'Label'
|
|
X = df['content']
|
|
y = df['Label']
|
|
|
|
# Create and fit TF-IDF vectorizer
|
|
print("Creating TF-IDF vectorizer...")
|
|
vectorizer = TfidfVectorizer(max_features=5000)
|
|
X_tfidf = vectorizer.fit_transform(X)
|
|
|
|
# Split the data
|
|
X_train, X_test, y_train, y_test = train_test_split(X_tfidf, y, test_size=0.2, random_state=42)
|
|
|
|
# Train Naive Bayes model
|
|
print("Training Naive Bayes model...")
|
|
model = MultinomialNB()
|
|
model.fit(X_train, y_train)
|
|
|
|
# Evaluate the model
|
|
train_score = model.score(X_train, y_train)
|
|
test_score = model.score(X_test, y_test)
|
|
print(f"Train accuracy: {train_score:.4f}")
|
|
print(f"Test accuracy: {test_score:.4f}")
|
|
|
|
# Save the models
|
|
print("Saving models...")
|
|
joblib.dump(vectorizer, 'tfidf_vectorizer.pkl')
|
|
joblib.dump(model, 'model_nb.pkl')
|
|
|
|
print("Models saved successfully!") |