salary-prediction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2a33b79",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.linear_model import LinearRegression\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import mean_squared_error, mean_absolute_error\n",
    "import numpy as np\n",
    "from sklearn.tree import DecisionTreeRegressor\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.model_selection import GridSearchCV\n",
    "import pickle\n",
    "\n",
    "df = pd.read_csv('data/survey_results_public.csv')\n",
    "# print(df.columns.sort_values())\n",
    "# Only get the columns that are needed for our model\n",
    "df = df[[\"Employment\", \"Country\", \"EdLevel\", \"ConvertedCompYearly\", \"YearsCodePro\"]]\n",
    "# Rename the columns with more intuitive naming\n",
    "df = df.rename({\"ConvertedCompYearly\" : \"Salary\", \"YearsCodePro\" : \"YearsOfExp\"}, axis = 1)\n",
    "print(df.info())\n",
    "\n",
    "# Drop all the rows that contains missing data and reset the index\n",
    "df = df.dropna()\n",
    "# df = df[df[\"YearsOfExp\"].notna()]\n",
    "# print(df.isnull().sum())\n",
    "# print(df)\n",
    "\n",
    "# Only collect data with people of full time employment\n",
    "df = df[df[\"Employment\"] == \"Employed full-time\"]\n",
    "# After that, drop the column, since we no longer need it\n",
    "df = df.drop(\"Employment\", axis = 1)\n",
    "print(df.info())\n",
    "\n",
    "print(df[\"Country\"].value_counts())\n",
    "\n",
    "def shorten_categories(categories, cutoff):\n",
    "    categorical_map = {}\n",
    "    for i in range(len(categories)):\n",
    "        if categories.values[i] >= cutoff:\n",
    "            categorical_map[categories.index[i]] = categories.index[i]\n",
    "        else:\n",
    "            categorical_map[categories.index[i]] = 'Other'\n",
    "\n",
    "    return categorical_map\n",
    "\n",
    "country_map = shorten_categories(df.Country.value_counts(), 400)\n",
    "print(country_map)\n",
    "df.Country = df.Country.map(country_map)\n",
    "\n",
    "print(df.Country.value_counts())\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize = (12, 7))\n",
    "df.boxplot('Salary', 'Country', ax = ax)\n",
    "plt.suptitle('Salary (US$) v Country')\n",
    "plt.title('')\n",
    "plt.ylabel('Salary')\n",
    "plt.xlabel('Country')\n",
    "plt.xticks(rotation = 90)\n",
    "# plt.show()\n",
    "\n",
    "# Now that we see the outliers for each data, we want to exclude those with salary that is too low or too high that can skew our predictions\n",
    "# Since USA has the most data point, I decide to find the median salary of those from the country to set the upper limit\n",
    "print(df[df.Country == 'United States of America'].Salary.median())\n",
    "df = df[df['Salary'] <= df[df.Country == 'United States of America'].Salary.median()]\n",
    "df = df[df['Salary'] >= 10000]\n",
    "df = df[df['Country'] != 'Other']\n",
    "fig, ax = plt.subplots(1, 1, figsize = (12, 7))\n",
    "df.boxplot('Salary', 'Country', ax = ax)\n",
    "plt.suptitle('Salary (US$) v Country')\n",
    "plt.title('')\n",
    "plt.ylabel('Salary')\n",
    "plt.xlabel('Country')\n",
    "plt.xticks(rotation = 90)\n",
    "# plt.show()\n",
    "# Now we see fewer outliers in our data set\n",
    "\n",
    "# print(df[\"YearsOfExp\"].unique())\n",
    "regressor = []\n",
    "le_country = []\n",
    "le_edu = []\n",
    "try:\n",
    "    with open('saved_model.pickle', 'rb') as file:\n",
    "        data = pickle.load(file)\n",
    "    regressor = data['model']\n",
    "    le_country = data[\"le_country\"]\n",
    "    le_edu = data['le_edu']\n",
    "except:\n",
    "    def convert_exp(num):\n",
    "        if num == 'More than 50 years':\n",
    "            return 50\n",
    "        if num == 'Less than 1 year':\n",
    "            return 0.5\n",
    "        return float(num)\n",
    "    # Convert string to float\n",
    "    df['YearsOfExp'] = df['YearsOfExp'].apply(convert_exp)\n",
    "\n",
    "    print(df.EdLevel.unique())\n",
    "\n",
    "    def clean_edu(x):\n",
    "        if 'Bachelor’s degree' in x:\n",
    "            return 'Bachelor’s degree'\n",
    "        if 'Master’s degree' in x:\n",
    "            return 'Master’s degree'\n",
    "        if 'Professional degree' in x or 'Other doctoral' in x:\n",
    "            return 'Post Grad'\n",
    "        return 'Less than a Bachelor'\n",
    "\n",
    "    # reclassify education level for simplicity\n",
    "    df['EdLevel'] = df['EdLevel'].apply(clean_edu)\n",
    "    print(df.EdLevel.unique())\n",
    "    # Use Label Encoder to transform string data to what computer can interpret which is numbers\n",
    "    le_country = LabelEncoder()\n",
    "    le_edu = LabelEncoder()\n",
    "    df['EdLevel'] = le_edu.fit_transform(df['EdLevel'])\n",
    "    print(df.EdLevel.unique())\n",
    "\n",
    "    df['Country'] = le_country.fit_transform(df['Country'])\n",
    "    # print(df.Country.unique())\n",
    "\n",
    "    # Separate features and label\n",
    "    x = df.drop('Salary', axis = 1)\n",
    "    y = df['Salary']\n",
    "    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size= 0.1)\n",
    "\n",
    "    linear = LinearRegression()\n",
    "    linear.fit(x_train, y_train)\n",
    "    accuracy = linear.score(x_test, y_test)\n",
    "    # We can see that our model accuracy is very low\n",
    "    print(\"Accuracy Percentage for Linear Regression: \", format(accuracy, \"%\"))\n",
    "    # When calculating the error, we can see model predict off by approx $26791\n",
    "    y_pred = linear.predict(x_test)\n",
    "    error = np.sqrt(mean_squared_error(y_test, y_pred))\n",
    "    print(\"${:,.02f}\".format(error))\n",
    "\n",
    "    # Now I choose Decision Tree Regressor as the new model\n",
    "    dec_tree_reg = DecisionTreeRegressor(random_state=0)\n",
    "    dec_tree_reg.fit(x_train, y_train)\n",
    "    accuracy = dec_tree_reg.score(x_test, y_test)\n",
    "    # We can see that our model accuracy is now a bit higher but it's still not enough\n",
    "    print(\"Accuracy Percentage for Decision Tree Regressor: \", format(accuracy, \"%\"))\n",
    "    y_pred = dec_tree_reg.predict(x_test)\n",
    "    error = np.sqrt(mean_squared_error(y_test, y_pred))\n",
    "    print(\"${:,.02f}\".format(error))\n",
    "\n",
    "    max_depth = [None, 2,4,6,8,10,12]\n",
    "    parameters = {\"max_depth\" : max_depth}\n",
    "\n",
    "    regressor = DecisionTreeRegressor(random_state=0)\n",
    "    gs = GridSearchCV(regressor, parameters, scoring = 'neg_mean_squared_error')\n",
    "    gs.fit(x_train, y_train)\n",
    "    regressor = gs.best_estimator_\n",
    "    regressor.fit(x_train, y_train)\n",
    "    accuracy = regressor.score(x_test, y_test)\n",
    "    # We can see that our model accuracy is now a bit higher but it's still not enough\n",
    "    print(\"Accuracy Percentage for Decision Tree Regressor after tuning: \", format(accuracy, \"%\"))\n",
    "    y_pred = regressor.predict(x_test)\n",
    "    error = np.sqrt(mean_squared_error(y_test, y_pred))\n",
    "    print(\"${:,.02f}\".format(error))\n",
    "\n",
    "    # Now I choose Random Forest Regressor\n",
    "    random_forest_reg = RandomForestRegressor(random_state=0)\n",
    "    random_forest_reg.fit(x_train, y_train)\n",
    "    accuracy = random_forest_reg.score(x_test, y_test)\n",
    "    print(\"Accuracy Percentage for Random Forest Regressor: \", format(accuracy, \"%\"))\n",
    "    y_pred = random_forest_reg.predict(x_test)\n",
    "    error = np.sqrt(mean_squared_error(y_test, y_pred))\n",
    "    print(\"${:,.02f}\".format(error))\n",
    "\n",
    "\n",
    "\n",
    "    parameters = {\"max_depth\" : [None, 2,4,6,8,10,12],\n",
    "                  \"max_features\" : [\"sqrt\", \"log2\", None],\n",
    "                  \"n_estimators\" : [100,140,160,180,200,300]}\n",
    "\n",
    "    regressor = RandomForestRegressor(random_state=0)\n",
    "    gs = GridSearchCV(regressor, parameters, scoring = 'neg_mean_squared_error', verbose=True)\n",
    "    gs.fit(x_train, y_train)\n",
    "    regressor = gs.best_estimator_\n",
    "    regressor.fit(x_train, y_train)\n",
    "    accuracy = regressor.score(x_test, y_test)\n",
    "\n",
    "    print(\"Accuracy Percentage for Random Forest Regressor after tuning: \", format(accuracy, \"%\"))\n",
    "    y_pred = regressor.predict(x_test)\n",
    "    error = np.sqrt(mean_squared_error(y_test, y_pred))\n",
    "    print(\"${:,.02f}\".format(error))\n",
    "\n",
    "\n",
    "    data = {\"model\" : regressor, \"le_country\" : le_country, \"le_edu\": le_edu}\n",
    "\n",
    "    # Write data to pickle so we don't have to train the model everytime\n",
    "    with open('saved_model.pickle', 'wb') as file:\n",
    "        pickle.dump(data, file)\n",
    "\n",
    "# country, edLevel, yearsOfExp\n",
    "x = np.array([[\"United States of America\", \"Master's degree\", 15]])\n",
    "x[:, 0] = le_country.fit_transform(x[:, 0])\n",
    "x[:, 1] = le_edu.fit_transform(x[:, 1])\n",
    "x = x.astype(float)\n",
    "\n",
    "y_pred = regressor.predict(x)\n",
    "print(y_pred)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}