Add changes to save new rag dadataframe

rag_helper.ipynb

@@ -251,6 +251,9 @@
    "outputs": [],
    "source": [
     "def run_evaluation(nba_df):\n",
+    "    team_flags = []\n",
+    "    game_flags = []\n",
+    "    other_stats_flags =[]\n",
     "    for index, row in nba_df.iterrows():\n",
     "        # Create message with sample query and run model\n",
     "        message=[{ 'role': 'user', 'content': input_text + row[\"natural_query\"]}]\n",
@@ -259,9 +262,16 @@
     "\n",
     "        # Obtain output\n",
     "        query_output = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)\n",
-    "\n",
-    "        print(\"Query: \", + row[\"sql_query\"])\n",
-    "        print(\"Response: \",query_output)\n"
+    "        team_flags.append(\"team\" in query_output.lower())\n",
+    "        game_flags.append(\"game\" in query_output.lower())\n",
+    "        other_stats_flags.append(\"other_stats\" in query_output.lower())\n",
+    "        #print(\"Query: \", + row[\"sql_query\"])\n",
+    "        #print(\"Response: \",query_output)\n",
+    "    \n",
+    "    nba_df[\"team_flag\"] = team_flags\n",
+    "    nba_df[\"game_flag\"] = game_flags\n",
+    "    nba_df[\"other_stats_flag\"] = other_stats_flags\n",
+    "    nba_df.to_csv(get_path(\"expanded_dta.tsv\"), sep=\"\\t\", index=False)\n"
    ]
   },
   {