Salary Prediction

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
import xgboost as xgb
from sklearn.metrics import mean_absolute_error
import warnings
from sklearn.model_selection import GridSearchCV, cross_val_score

warnings.filterwarnings('ignore')

BASELINE_FEATURES = ['Rating', 'Size', 'Type of ownership', 'Industry', 'Sector', 'Revenue']
Y_COLUMN = ['Salary Estimate']

!pip install gdown
!gdown 1atCcqIxgKfvmtRUQkqz--wF2FTlEAHal
!gdown 1fA6iChCXW2e6v9a3G68r6b6L36_Em1ki

# ==============================================================================
# === (HELPER FUNCTIONS AND BASELINE MODEL - DO NOT CHANGE THIS CELL) ===
# ==============================================================================

def title_simplifier(title):
    if 'data scientist' in title.lower():
        return 'data scientist'
    elif 'data engineer' in title.lower():
        return 'data engineer'
    elif 'analyst' in title.lower():
        return 'analyst'
    elif 'machine learning' in title.lower():
        return 'mle'
    elif 'manager' in title.lower():
        return 'manager'
    elif 'director' in title.lower():
        return 'director'
    else:
        return 'na'

def seniority(title):
    if 'sr' in title.lower() or 'senior' in title.lower() or 'lead' in title.lower() or 'principal' in title.lower():
        return 'senior'
    elif 'jr' in title.lower() or 'jr.' in title.lower():
        return 'jr'
    else:
        return 'na'

def load_data(path: str) -> pd.DataFrame:
    """
    Loads the data and performs a robust initial cleaning,
    inspired by the detailed analysis.
    """
    df = pd.read_csv(path)
    return df

def clean_train_data(df: pd.DataFrame) -> pd.DataFrame:
    df = df[df['Salary Estimate']!= '-1']

    # Cleaning the salary string
    df['Salary Estimate'] = df['Salary Estimate'].apply(lambda x: x.split('(')[0])
    df['Salary Estimate'] = df['Salary Estimate'].apply(lambda x: x.replace('K','').replace('$',''))

    df['Salary Estimate'] = df['Salary Estimate'].apply(lambda x: x.lower().replace('per hour', ''))
    df['Salary Estimate'] = df['Salary Estimate'].apply(lambda x: x.lower().replace('employer provided salary:', ''))

    df['Rating'] = df['Rating'].apply(lambda x: x if x > 0 else np.nan)
    df['State'] = df.Location.apply(lambda x: x.split(',')[1])

    # Handling both ranges (e.g. "50-100") and single values ("75")
    df['Min_Salary'] = df['Salary Estimate'].apply(lambda x: int(x.split('-')[0]))
    df['Max_Salary'] = df['Salary Estimate'].apply(lambda x: int(x.split('-')[1]))
    df['Salary Estimate']= (df['Min_Salary'] + df['Max_Salary'])/2
    df.drop(['Min_Salary', 'Max_Salary'], axis=1, inplace=True)
    return df

def clean_data_features(df: pd.DataFrame) -> pd.DataFrame:
    df['Rating'] = df['Rating'].apply(lambda x: x if x > 0 else np.nan)
    df['State'] = df.Location.apply(lambda x: x.split(',')[1])
    return df

def train_baseline_model(df: pd.DataFrame) -> float:
    """Trains a baseline model and returns it, together with its MAE and ."""

    # We use several categorical and numerical features
    base_features = BASELINE_FEATURES + Y_COLUMN
    df_base = df[base_features].dropna()

    # Creating dummy variables for the baseline model
    df_base_dum = pd.get_dummies(df_base)
    X = df_base_dum.drop('Salary Estimate', axis=1)
    y = df_base_dum['Salary Estimate'].values

    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

    model = LinearRegression()
    model.fit(X_train, y_train)

    y_pred = model.predict(X_val)
    mae = mean_absolute_error(y_val, y_pred)

    print(f"Baseline Model MAE: {mae:.3f}")
    return model, mae, X, y

# ==============================================================================
# === (MAIN SCRIPT - DO NOT CHANGE THIS CELL) ===
# ==============================================================================
# Loading the data
df_1 = load_data("train.csv")
df = clean_train_data(df_1)

# Training and scoring the baseline model
baseline_model, baseline_mae, X_train_baseline, y_train_baseline = train_baseline_model(df.copy())

df_test = load_data("../../../test_features.csv")
df_test_cleaned = clean_data_features(df_test)
df_test_processed = df_test_cleaned[BASELINE_FEATURES].dropna()
baseline_model_columns = X_train_baseline.columns
# Creating dummy variables for the baseline model
df_test_dum = pd.get_dummies(df_test_processed)
df_test_aligned = df_test_dum.reindex(
    columns=baseline_model_columns, fill_value=0
)
y_pred = baseline_model.predict(df_test_aligned)
y_pred_df = pd.DataFrame(y_pred, columns=Y_COLUMN)
# y_pred_df.to_csv('predictions.csv', index=False)

# ==============================================================================
# ZONE FOR YOUR CODE
# =============================================================================
# You can try to:
# - clean/rework the data further (e.g. )
# - add other features (e.g. the age of the company)
# - experiment with different models
# - combine several models
# - run a grid search to find optimal parameters
# ==============================================================================
student_model = baseline_model

df[['Age','Salary Estimate', 'Rating', 'desc_len']].corr()

_, X_val, _, y_val = train_test_split(X_train_baseline, y_train_baseline, test_size=0.2, random_state=42)
student_predictions = student_model.predict(X_val)
student_mae = mean_absolute_error(y_val, student_predictions)

print("\n--- Final results ---")
print(f"The baseline model has MAE: {baseline_mae:.3f}")
print(f"Your model has MAE:         {student_mae:.3f}")

improvement = baseline_mae - student_mae
if improvement > 0:
    print(f"\nImprovement over the baseline model: {improvement:.3f}!")
else:
    print("\nTry again! Your model is not better than the baseline.")