{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd9ef55b",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Script for getting prediction for a TEST.CSV file with three columns of text.\n",
    "It looks at the columns Q1, Q2, Q3, concatenates them and passes the full string as text into \n",
    "PersonalityClassifier(). \n",
    ">>> The method \"predict_all_traits\" defined in the class will get the predictions by running five \n",
    "separate prediction models (optimized random forests). \n",
    "The predictions are then applies as labels for each of the train columns.\n",
    ">>> Important: All the other columns in the original CSV file will be untouched (Q1, Q2, Q3 and Humility). \n",
    "The CSV input file does not need to have empty values for traits; \n",
    "the script replaces the predictions with annotations.\n",
    "'''\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=FutureWarning)\n",
    "import pandas as pd\n",
    "from personality_model import PersonalityClassifier\n",
    "\n",
    "# ***************** LOAD THE TEST DATA WITH Q1, Q2, Q3 *********************\n",
    "\n",
    "input_path = \"/path/to/test.csv\"  # path to test data\n",
    "output_path = \"/path/to/output/filled_predictions.csv\"  # change PATH and NAME of output\n",
    "\n",
    "df = pd.read_csv(input_path)\n",
    "\n",
    "#  concatenating Q1, Q2, Q3 \n",
    "texts = df[[\"Q1\", \"Q2\", \"Q3\"]].fillna(\"\").agg(\" \".join, axis=1)\n",
    "\n",
    "# model initialization\n",
    "model = PersonalityClassifier()\n",
    "\n",
    "# predicting trait labels for each row\n",
    "predictions = texts.apply(model.predict_all_traits)\n",
    "\n",
    "# applying the predictions and filling the columns\n",
    "for trait in [\"Openness\", \"Conscientiousness\", \"Extraversion\", \"Agreeableness\", \"Emotional stability\"]:\n",
    "    df[trait] = predictions.apply(lambda d: d[trait])\n",
    "\n",
    "df.to_csv(output_path, index=False)\n",
    "\n",
    "print(f\"Trait predictions saved to: {output_path}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "amiv_nlp_2025",
   "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.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}