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!")