predict.py

import argparse
import os
import numpy as np
import pandas as pd
from tqdm import tqdm
from sklearn.metrics import f1_score, precision_score, recall_score
import train_utils

# Parse command-line arguments
parser = argparse.ArgumentParser(description='Generate predictions using cross-validation')
parser.add_argument('--target_type', type=int, default=None, choices=[0, 1, 2, 3],
                    help='Target type: 0=numeric, 1=date, 2=categorical, 3=text. If not set, runs all types.')
parser.add_argument('--data_dir', type=str, default='merged_features/',
                    help='Directory containing the feature CSV files (default: merged_features/)')
parser.add_argument('--output', type=str, default='predictions.csv',
                    help='Output CSV file name (default: predictions.csv)')
parser.add_argument('--classifier', type=str, default=None,
                    help='Classifier to use')
parser.add_argument('--balance_type', type=str, default='weight', 
                    help='Balance type for training data (default: weight)')
parser.add_argument('--featureset', type=str, default='small',
                    help='Feature set to use')
parser.add_argument('--fast', action='store_true',
                    help='Evaluate in chunks of five datasets for faster evaluation (default: False)')

args = parser.parse_args()

# Constants
TYPE_MAPPING = {'numeric': 0, 'date': 1, 'categorical': 2, 'text': 3}
TARGET_TYPE_STR_MAP = {v: k for k, v in TYPE_MAPPING.items()}

# Get feature sets
featuresets = train_utils.get_featuresets(args.featureset)
featureset_num = featuresets['num']
featureset_date = featuresets['date']
featureset_cat = featuresets['cat']
featureset_text = featuresets['text']

print("Loading data...")
data_dir = args.data_dir

print("Loading all column CSV files...")
csv_files = sorted([f for f in os.listdir(data_dir) if f.endswith('.csv') and 'lock' not in f])
print(f"Found {len(csv_files)} CSV files")

required_cols = ["label", "dataset", "type", "header"] 
all_cols_needed = set(required_cols)
all_cols_needed.update(featureset_num)
all_cols_needed.update(featureset_date)
all_cols_needed.update(featureset_cat)
all_cols_needed.update(featureset_text)

eval_data = train_utils.load_columnar_data(data_dir, [f.replace('.csv', '') for f in csv_files])

# Determine which types to process
if args.target_type is not None:
    target_types = [args.target_type]
else:
    target_types = list(set(TYPE_MAPPING.values()))

level = 'dataset'
all_eval_preds = []
all_eval_labels = []
all_eval_datasets = []
all_eval_types = []
all_eval_probas = []

print(f"Processing target types: {[TARGET_TYPE_STR_MAP[t] for t in target_types]}")

for TARGET_TYPE in target_types:
    print(f"\nProcessing type: {TARGET_TYPE_STR_MAP[TARGET_TYPE]}")
    
    featureset = train_utils.get_featureset(TARGET_TYPE, featuresets)
    
    # Filter out features that don't exist in the data
    available_features = [f for f in featureset if f in eval_data.columns]
    if len(available_features) < len(featureset):
        missing = set(featureset) - set(available_features)
        print(f"  Warning: Missing features: {missing}")
    featureset = available_features
    
    if not featureset:
         print("  No features available for this type. Skipping.")
         continue

    type_data = eval_data[eval_data['type'] == TARGET_TYPE].copy()
    if len(type_data) == 0:
        print(f"  No data for type {TARGET_TYPE_STR_MAP[TARGET_TYPE]}, skipping...")
        continue
    
    type_data[featureset] = type_data[available_features].fillna(0)
    
    unique_datasets = type_data[level].unique()
    if args.fast:
        rng = np.random.RandomState(42)
        shuffled = unique_datasets.copy()
        rng.shuffle(shuffled)
        dataset_chunks = [shuffled[i:i + 5] for i in range(0, len(shuffled), 5)]
    else:
        dataset_chunks = [[d] for d in unique_datasets]
    
    for dataset_chunk in tqdm(dataset_chunks, desc=f"  Datasets ({TARGET_TYPE_STR_MAP[TARGET_TYPE]})"):
        dataset_mask = type_data[level].isin(dataset_chunk)
        
        train_split = type_data[~dataset_mask]
        eval_split = type_data[dataset_mask]
        
        if len(train_split) == 0 or len(eval_split) == 0:
            continue
        
        # Get initial features and labels
        train_features = train_split[featureset].values
        train_labels = train_split['label'].values.astype(int)
        
        # Balance data
        train_features, train_labels = train_utils.balance_data(
            train_features, 
            train_labels, 
            balance_type=args.balance_type if args.balance_type is not None else 'weight' if TARGET_TYPE in [1,3] else 'stratified'
        )
        
        num_positive = int(train_labels.sum())
        num_negative = len(train_labels) - num_positive
        
        eval_labels = eval_split['label'].values.astype(int)
        threshold = 0.5 

        clf = train_utils.fit_classifier(
            args.classifier if args.classifier is not None else 'xgboost_robust' if TARGET_TYPE in [1,3] else 'xgboost_stratified_categorical', train_features, train_labels
        )

        # ── Predict on eval set ──────────────────────────────────────────────
        eval_features = eval_split[featureset].values

        # Get probabilities
        if hasattr(clf, 'predict_proba'):
            eval_proba = clf.predict_proba(eval_features)[:, 1]
        elif hasattr(clf, 'decision_function'):
            eval_proba = 1 / (1 + np.exp(-clf.decision_function(eval_features)))
        else:
            eval_proba = clf.predict(eval_features).astype(float)

        eval_preds = (eval_proba > threshold).astype(int)
            
        all_eval_preds.extend(eval_preds.tolist())
        all_eval_probas.extend(eval_proba.tolist())
        all_eval_labels.extend(eval_labels.tolist())
        all_eval_datasets.extend(eval_split['dataset'].tolist())
        all_eval_types.extend(eval_split['type'].tolist())

# Save results
evaluated_df = pd.DataFrame({
    'preds': all_eval_preds,
    'proba': all_eval_probas,
    'label': all_eval_labels,
    'dataset': all_eval_datasets,
    'type': all_eval_types,
})

output_file = args.output.split(".csv")[0] + f"_{args.classifier}_{args.featureset}_{args.balance_type}.csv"
evaluated_df.to_csv(output_file, index=False)
print(f"\nPredictions saved to {output_file}")
print(f"Total predictions: {len(evaluated_df)}")

# Quick evaluation
preds = evaluated_df['preds'].values
labels = evaluated_df['label'].values
f1 = f1_score(labels, preds)
precision = precision_score(labels, preds)
recall = recall_score(labels, preds)
print(f"Overall F1: {f1:.4f}, Precision: {precision:.4f}, Recall: {recall:.4f}")