NLP-Algorithms/Algorithms/Text_Classification/train.py at main · UTSAVS26/NLP-Algorithms · GitHub

34 lines (27 loc) · 1.04 KB

import pandas as pd

from sklearn.feature_extraction.text import TfidfVectorizer

from sklearn.model_selection import train_test_split

from sklearn.naive_bayes import MultinomialNB

from sklearn.pipeline import Pipeline

from sklearn.metrics import accuracy_score, classification_report

import joblib

from preprocess import preprocess_text

# Load dataset

data = pd.read_csv('data/text_data.csv')

# Preprocess text data

data['text'] = data['text'].apply(preprocess_text)

# Split dataset into training and testing sets

X_train, X_test, y_train, y_test = train_test_split(data['text'], data['label'], test_size=0.2, random_state=42)

# Create a pipeline with TfidfVectorizer and MultinomialNB

model = Pipeline([

('tfidf', TfidfVectorizer()),

('nb', MultinomialNB())

])

# Train the model

model.fit(X_train, y_train)

# Evaluate the model

y_pred = model.predict(X_test)

print(f'Accuracy: {accuracy_score(y_test, y_pred)}')

print(f'Classification Report:\n{classification_report(y_test, y_pred)}')

# Save the model

joblib.dump(model, 'models/text_classifier.pkl')

Expand file tree