test_prepare.py

import pandas as pd
import os
import glob
from sklearn.model_selection import train_test_split

print("=" * 60)
print("🎬 QUICK TEST - DATA PREPARATION (5,000 reviews)")
print("=" * 60)

# Find dataset
print("\n[1/4] 🔍 Finding dataset...")
dataset_files = glob.glob('data/*.csv')

if not dataset_files:
    print("❌ No CSV found in data/ folder!")
    exit(1)

dataset_path = dataset_files[0]
print(f"   ✅ Found: {os.path.basename(dataset_path)}")

# Load ONLY first 5,000 rows for quick testing
print("\n[2/4] 📥 Loading 5,000 reviews (quick test)...")
try:
    df = pd.read_csv(dataset_path, nrows=5000)
    print(f"   ✅ Loaded {len(df)} reviews")
except Exception as e:
    print(f"❌ ERROR: {e}")
    exit(1)

# Quick inspection
print(f"\n[3/4] 🔍 Data preview:")
print(f"   Columns: {list(df.columns)}")
print(f"   Shape: {df.shape}")
print(df.head(2))

# Standardize columns
review_col = [c for c in df.columns if 'review' in c.lower()][0]
sentiment_col = [c for c in df.columns if 'sentiment' in c.lower()][0]

df = df.rename(columns={review_col: 'review', sentiment_col: 'sentiment'})

# Convert sentiment to binary
sentiment_map = {'positive': 1, 'negative': 0, 'pos': 1, 'neg': 0}
df['sentiment'] = df['sentiment'].map(lambda x: sentiment_map.get(str(x).lower(), x))

# Remove missing values
df = df.dropna()

print(f"\n   Sentiment distribution: {df['sentiment'].value_counts().to_dict()}")

# Split data
print("\n[4/4] ✂️  Splitting data...")
train_df, temp_df = train_test_split(df, test_size=0.2, random_state=42)
val_df, test_df = train_test_split(temp_df, test_size=0.5, random_state=42)

print(f"   Train: {len(train_df)} | Val: {len(val_df)} | Test: {len(test_df)}")

# Save
os.makedirs('data/processed', exist_ok=True)
train_df.to_csv('data/processed/train.csv', index=False)
val_df.to_csv('data/processed/val.csv', index=False)
test_df.to_csv('data/processed/test.csv', index=False)

print("\n✅ SUCCESS! Files saved in data/processed/")
print("\n🎯 Next: Run 'python train_model.py' to train your model")
print("=" * 60)