train.py

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
import pickle

# Load dataset (Sentiment140)
df = pd.read_csv(
    "training.1600000.processed.noemoticon.csv",
    encoding="latin-1",
    header=None
)

# Keep only label + text columns
df = df[[0, 5]]
df.columns = ["label", "text"]

# Convert labels: 0=negative, 4=positive → 1=positive
df["label"] = df["label"].replace({4: 1})

# 🔥 IMPORTANT: reduce size for faster training
df = df.sample(5000, random_state=42)

# Split data
X_train, X_test, y_train, y_test = train_test_split(
    df["text"], df["label"], test_size=0.2
)

# Vectorization (better than CountVectorizer)
vectorizer = TfidfVectorizer(stop_words="english", max_features=5000)
X_train_vec = vectorizer.fit_transform(X_train)

# Train model
model = MultinomialNB()
model.fit(X_train_vec, y_train)

# Save model
pickle.dump(model, open("model.pkl", "wb"))
pickle.dump(vectorizer, open("vectorizer.pkl", "wb"))

print("✅ Model trained successfully!")