imdb-movie-review-sentiment-analysis / train_and_save_model.py
MON3EMPASHA's picture
Upload folder using huggingface_hub
f19f2e1 verified
Raw
History Blame Contribute Delete
10.3 kB
import pandas as pd
import nltk
import re
import pickle
import joblib
import json
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.naive_bayes import MultinomialNB
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
from collections import Counter
import os
# Download NLTK resources
try:
nltk.download('punkt')
nltk.download('stopwords')
nltk.download('wordnet')
nltk.download('punkt_tab')
except Exception as e:
print(f"Warning: Could not download NLTK data: {e}")
# Text Preprocessing Function
def preprocess_text(text):
# Lowercase
text = text.lower()
# Remove special characters and digits
text = re.sub(r'[^a-zA-Z\s]', '', text)
# Tokenize
tokens = word_tokenize(text)
# Remove stopwords
stop_words = set(stopwords.words('english'))
tokens = [word for word in tokens if word not in stop_words]
# Lemmatize
lemmatizer = WordNetLemmatizer()
tokens = [lemmatizer.lemmatize(word) for word in tokens]
# Join tokens back to string
return ' '.join(tokens)
# Load and Preprocess Dataset
def load_and_preprocess_data(file_path="IMDB Dataset.csv"):
try:
df = pd.read_csv(file_path)
print(f"Loaded {len(df)} reviews")
# Apply preprocessing to reviews
print("Preprocessing reviews...")
df['cleaned_review'] = df['review'].apply(preprocess_text)
# Convert sentiment to binary (1 for positive, 0 for negative)
df['sentiment'] = df['sentiment'].replace({'positive': 1, 'negative': 0})
return df
except FileNotFoundError:
print(f"Error: Could not find file '{file_path}'")
return None
# Train Models and Save
def train_and_save_models(df, model_dir="saved_models"):
# Create model directory
os.makedirs(model_dir, exist_ok=True)
# Convert text to TF-IDF features
print("Vectorizing text data...")
vectorizer = TfidfVectorizer(max_features=5000)
X = vectorizer.fit_transform(df['cleaned_review'])
y = df['sentiment']
# Split data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# Train Logistic Regression
print("Training Logistic Regression...")
lr_model = LogisticRegression(max_iter=1000, random_state=42)
lr_model.fit(X_train, y_train)
lr_predictions = lr_model.predict(X_test)
lr_accuracy = accuracy_score(y_test, lr_predictions)
# Train Naive Bayes
print("Training Naive Bayes...")
nb_model = MultinomialNB()
nb_model.fit(X_train, y_train)
nb_predictions = nb_model.predict(X_test)
nb_accuracy = accuracy_score(y_test, nb_predictions)
# Save models
print("Saving models...")
# Save vectorizer
joblib.dump(vectorizer, os.path.join(model_dir, 'tfidf_vectorizer.pkl'))
# Save Logistic Regression model
joblib.dump(lr_model, os.path.join(model_dir, 'logistic_regression_model.pkl'))
# Save Naive Bayes model
joblib.dump(nb_model, os.path.join(model_dir, 'naive_bayes_model.pkl'))
# Save model metadata
metadata = {
'vectorizer_features': vectorizer.get_feature_names_out().tolist(),
'max_features': 5000,
'lr_accuracy': float(lr_accuracy),
'nb_accuracy': float(nb_accuracy),
'training_samples': X_train.shape[0],
'test_samples': X_test.shape[0],
'preprocessing_steps': [
'lowercase',
'remove_special_chars',
'tokenization',
'stopword_removal',
'lemmatization'
]
}
with open(os.path.join(model_dir, 'model_metadata.json'), 'w') as f:
json.dump(metadata, f, indent=2)
# Print results
print("\n" + "="*50)
print("MODEL TRAINING RESULTS")
print("="*50)
print(f"Logistic Regression Accuracy: {lr_accuracy:.4f}")
print(f"Naive Bayes Accuracy: {nb_accuracy:.4f}")
print(f"Models saved to: {model_dir}/")
print("="*50)
return vectorizer, lr_model, nb_model, lr_accuracy, nb_accuracy
# Create inference script
def create_inference_script():
inference_code = '''import joblib
import json
import re
import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
# Download NLTK resources
try:
nltk.download('punkt')
nltk.download('stopwords')
nltk.download('wordnet')
except:
pass
class SentimentAnalyzer:
def __init__(self, model_dir="saved_models"):
# Load models
self.vectorizer = joblib.load(f"{model_dir}/tfidf_vectorizer.pkl")
self.lr_model = joblib.load(f"{model_dir}/logistic_regression_model.pkl")
self.nb_model = joblib.load(f"{model_dir}/naive_bayes_model.pkl")
# Load metadata
with open(f"{model_dir}/model_metadata.json", 'r') as f:
self.metadata = json.load(f)
def preprocess_text(self, text):
# Lowercase
text = text.lower()
# Remove special characters and digits
text = re.sub(r'[^a-zA-Z\\s]', '', text)
# Tokenize
tokens = word_tokenize(text)
# Remove stopwords
stop_words = set(stopwords.words('english'))
tokens = [word for word in tokens if word not in stop_words]
# Lemmatize
lemmatizer = WordNetLemmatizer()
tokens = [lemmatizer.lemmatize(word) for word in tokens]
# Join tokens back to string
return ' '.join(tokens)
def predict(self, text, model_type='both'):
# Preprocess text
cleaned_text = self.preprocess_text(text)
# Vectorize
text_vector = self.vectorizer.transform([cleaned_text])
results = {}
if model_type in ['lr', 'both']:
lr_pred = self.lr_model.predict(text_vector)[0]
lr_prob = self.lr_model.predict_proba(text_vector)[0]
results['logistic_regression'] = {
'prediction': 'positive' if lr_pred == 1 else 'negative',
'confidence': float(max(lr_prob)),
'probabilities': {
'negative': float(lr_prob[0]),
'positive': float(lr_prob[1])
}
}
if model_type in ['nb', 'both']:
nb_pred = self.nb_model.predict(text_vector)[0]
nb_prob = self.nb_model.predict_proba(text_vector)[0]
results['naive_bayes'] = {
'prediction': 'positive' if nb_pred == 1 else 'negative',
'confidence': float(max(nb_prob)),
'probabilities': {
'negative': float(nb_prob[0]),
'positive': float(nb_prob[1])
}
}
return results
# Example usage
if __name__ == "__main__":
analyzer = SentimentAnalyzer()
# Test with sample reviews
test_reviews = [
"This movie was absolutely fantastic! I loved every minute of it.",
"Terrible film, waste of time. Don't watch it.",
"It was okay, nothing special but not bad either."
]
for review in test_reviews:
print(f"\\nReview: {review}")
results = analyzer.predict(review)
for model, result in results.items():
print(f"{model}: {result['prediction']} (confidence: {result['confidence']:.2f})")
'''
with open('inference.py', 'w') as f:
f.write(inference_code)
print("Created inference.py for model deployment")
# Create requirements.txt
def create_requirements():
requirements = '''pandas>=1.3.0
nltk>=3.6
scikit-learn>=1.0.0
joblib>=1.1.0
numpy>=1.21.0
streamlit>=1.0.0
matplotlib>=3.5.0
seaborn>=0.11.0'''
with open('requirements.txt', 'w') as f:
f.write(requirements)
print("Created requirements.txt")
# Create README
def create_readme():
readme = '''# IMDb Sentiment Analysis Model
This repository contains a trained sentiment analysis model for IMDb movie reviews.
## Model Performance
- Logistic Regression: ~88.47% accuracy
- Naive Bayes: ~85.2% accuracy
## Files
- `saved_models/`: Directory containing trained models
- `inference.py`: Script for making predictions
- `train_and_save_model.py`: Script to train and save models
- `requirements.txt`: Python dependencies
## Usage
### Load and Use the Model
```python
from inference import SentimentAnalyzer
# Initialize analyzer
analyzer = SentimentAnalyzer()
# Make prediction
result = analyzer.predict("This movie was amazing!")
print(result)
```
### Deploy on Streamlit
```bash
streamlit run streamlit_deployment.py
```
## Model Details
- **Vectorizer**: TF-IDF with 5000 features
- **Preprocessing**: Lowercase, special char removal, tokenization, stopword removal, lemmatization
- **Models**: Logistic Regression and Naive Bayes
'''
with open('README.md', 'w') as f:
f.write(readme)
print("Created README.md")
if __name__ == "__main__":
print("Training and saving sentiment analysis models...")
# Load data
df = load_and_preprocess_data()
if df is not None:
# Train and save models
train_and_save_models(df)
# Create deployment files
create_inference_script()
create_requirements()
create_readme()
print("\n✅ Model training and saving completed!")
print("📁 Models saved in 'saved_models/' directory")
print("🚀 Ready for deployment on Hugging Face, Kaggle, or other platforms")
else:
print("❌ Failed to load data. Please check if 'IMDB Dataset.csv' exists.")