try_model_generator.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>context</th>\n",
       "      <th>source</th>\n",
       "      <th>cleaned_context</th>\n",
       "      <th>cleaned_context_num_gpt4_tokens</th>\n",
       "      <th>context_num_gpt4_tokens</th>\n",
       "      <th>num_tokens_range</th>\n",
       "      <th>generated_question</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>\\n pyspark.streaming.DStream.groupByKeyAndWind...</td>\n",
       "      <td>https://spark.apache.org/docs/latest/api/pytho...</td>\n",
       "      <td>pyspark.streaming.DStream.groupByKeyAndWindow ...</td>\n",
       "      <td>191</td>\n",
       "      <td>339</td>\n",
       "      <td>200~500</td>\n",
       "      <td>How to use groupByKeyAndWindow in Python to a...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>\\n pyspark.pandas.Series.any\\n [¶](#pyspark-pa...</td>\n",
       "      <td>https://spark.apache.org/docs/latest/api/pytho...</td>\n",
       "      <td>pyspark.pandas.Series.any [](#pysparkpandasser...</td>\n",
       "      <td>204</td>\n",
       "      <td>391</td>\n",
       "      <td>200~500</td>\n",
       "      <td>How do you check if any value is True in a py...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>\\n pyspark.sql.DataFrameNaFunctions.replace\\n ...</td>\n",
       "      <td>https://spark.apache.org/docs/latest/api/pytho...</td>\n",
       "      <td>pyspark.sql.DataFrameNaFunctions.replace [](#p...</td>\n",
       "      <td>477</td>\n",
       "      <td>1007</td>\n",
       "      <td>500~1200</td>\n",
       "      <td>How to replace values in a DataFrame column u...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>\\n pyspark.pandas.Series.dt.is\\_year\\_end\\n [¶...</td>\n",
       "      <td>https://spark.apache.org/docs/latest/api/pytho...</td>\n",
       "      <td>pyspark.pandas.Series.dt.is_year_end [](#pyspa...</td>\n",
       "      <td>230</td>\n",
       "      <td>337</td>\n",
       "      <td>200~500</td>\n",
       "      <td>How do you check if a datetime Series in Pand...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>\\n`dateadd`\\n (days) function\\n===============...</td>\n",
       "      <td>https://docs.databricks.com/sql/language-manua...</td>\n",
       "      <td>dateadd (days) function Returns the date numDa...</td>\n",
       "      <td>173</td>\n",
       "      <td>307</td>\n",
       "      <td>200~500</td>\n",
       "      <td>How do you use the dateadd function in Databr...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>195</th>\n",
       "      <td>\\n`regexp\\_count`\\n function\\n================...</td>\n",
       "      <td>https://docs.databricks.com/sql/language-manua...</td>\n",
       "      <td>regexp_count function Applies to: [check marke...</td>\n",
       "      <td>227</td>\n",
       "      <td>411</td>\n",
       "      <td>200~500</td>\n",
       "      <td>How can I use the regexp_count function in Py...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>196</th>\n",
       "      <td>\\n`array\\_union`\\n function\\n=================...</td>\n",
       "      <td>https://docs.databricks.com/sql/language-manua...</td>\n",
       "      <td>array_union function Applies to: [check marked...</td>\n",
       "      <td>149</td>\n",
       "      <td>264</td>\n",
       "      <td>200~500</td>\n",
       "      <td>How to use the array_union function in Databr...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>197</th>\n",
       "      <td>\\n Genomics guide\\n================\\n Importan...</td>\n",
       "      <td>https://docs.databricks.com/archive/genomics/i...</td>\n",
       "      <td>Genomics guide Important This documentation ha...</td>\n",
       "      <td>363</td>\n",
       "      <td>624</td>\n",
       "      <td>500~1200</td>\n",
       "      <td>How do I perform secondary analysis pipelines...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>198</th>\n",
       "      <td>\\n p\\n \\n[org](../../../../index.html) \\n .\\n ...</td>\n",
       "      <td>https://spark.apache.org/docs/latest/api/scala...</td>\n",
       "      <td>p [org](../../../../index.html) . [apache](../...</td>\n",
       "      <td>421</td>\n",
       "      <td>1051</td>\n",
       "      <td>500~1200</td>\n",
       "      <td>How to generate random data in Spark using th...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>199</th>\n",
       "      <td>\\n Rating\\n [¶](#rating \"Permalink to this hea...</td>\n",
       "      <td>https://spark.apache.org/docs/latest/api/pytho...</td>\n",
       "      <td>Rating [](#rating Permalink to this headline) ...</td>\n",
       "      <td>379</td>\n",
       "      <td>646</td>\n",
       "      <td>500~1200</td>\n",
       "      <td>How to use the Rating class in PySpark for re...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>200 rows × 7 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               context  \\\n",
       "0    \\n pyspark.streaming.DStream.groupByKeyAndWind...   \n",
       "1    \\n pyspark.pandas.Series.any\\n [¶](#pyspark-pa...   \n",
       "2    \\n pyspark.sql.DataFrameNaFunctions.replace\\n ...   \n",
       "3    \\n pyspark.pandas.Series.dt.is\\_year\\_end\\n [¶...   \n",
       "4    \\n`dateadd`\\n (days) function\\n===============...   \n",
       "..                                                 ...   \n",
       "195  \\n`regexp\\_count`\\n function\\n================...   \n",
       "196  \\n`array\\_union`\\n function\\n=================...   \n",
       "197  \\n Genomics guide\\n================\\n Importan...   \n",
       "198  \\n p\\n \\n[org](../../../../index.html) \\n .\\n ...   \n",
       "199  \\n Rating\\n [¶](#rating \"Permalink to this hea...   \n",
       "\n",
       "                                                source  \\\n",
       "0    https://spark.apache.org/docs/latest/api/pytho...   \n",
       "1    https://spark.apache.org/docs/latest/api/pytho...   \n",
       "2    https://spark.apache.org/docs/latest/api/pytho...   \n",
       "3    https://spark.apache.org/docs/latest/api/pytho...   \n",
       "4    https://docs.databricks.com/sql/language-manua...   \n",
       "..                                                 ...   \n",
       "195  https://docs.databricks.com/sql/language-manua...   \n",
       "196  https://docs.databricks.com/sql/language-manua...   \n",
       "197  https://docs.databricks.com/archive/genomics/i...   \n",
       "198  https://spark.apache.org/docs/latest/api/scala...   \n",
       "199  https://spark.apache.org/docs/latest/api/pytho...   \n",
       "\n",
       "                                       cleaned_context  \\\n",
       "0    pyspark.streaming.DStream.groupByKeyAndWindow ...   \n",
       "1    pyspark.pandas.Series.any [](#pysparkpandasser...   \n",
       "2    pyspark.sql.DataFrameNaFunctions.replace [](#p...   \n",
       "3    pyspark.pandas.Series.dt.is_year_end [](#pyspa...   \n",
       "4    dateadd (days) function Returns the date numDa...   \n",
       "..                                                 ...   \n",
       "195  regexp_count function Applies to: [check marke...   \n",
       "196  array_union function Applies to: [check marked...   \n",
       "197  Genomics guide Important This documentation ha...   \n",
       "198  p [org](../../../../index.html) . [apache](../...   \n",
       "199  Rating [](#rating Permalink to this headline) ...   \n",
       "\n",
       "     cleaned_context_num_gpt4_tokens  context_num_gpt4_tokens  \\\n",
       "0                                191                      339   \n",
       "1                                204                      391   \n",
       "2                                477                     1007   \n",
       "3                                230                      337   \n",
       "4                                173                      307   \n",
       "..                               ...                      ...   \n",
       "195                              227                      411   \n",
       "196                              149                      264   \n",
       "197                              363                      624   \n",
       "198                              421                     1051   \n",
       "199                              379                      646   \n",
       "\n",
       "    num_tokens_range                                 generated_question  \n",
       "0            200~500   How to use groupByKeyAndWindow in Python to a...  \n",
       "1            200~500   How do you check if any value is True in a py...  \n",
       "2           500~1200   How to replace values in a DataFrame column u...  \n",
       "3            200~500   How do you check if a datetime Series in Pand...  \n",
       "4            200~500   How do you use the dateadd function in Databr...  \n",
       "..               ...                                                ...  \n",
       "195          200~500   How can I use the regexp_count function in Py...  \n",
       "196          200~500   How to use the array_union function in Databr...  \n",
       "197         500~1200   How do I perform secondary analysis pipelines...  \n",
       "198         500~1200   How to generate random data in Spark using th...  \n",
       "199         500~1200   How to use the Rating class in PySpark for re...  \n",
       "\n",
       "[200 rows x 7 columns]"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load the environment variables\n",
    "import pandas as pd\n",
    "\n",
    "import dotenv\n",
    "dotenv.load_dotenv('../.env')\n",
    "\n",
    "\n",
    "small_dataset = pd.read_csv(\"../datasets/docs_qa_small_context_200.csv\")\n",
    "\n",
    "small_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:databricks.labs.doc_qa.model_generators.model_generator:Generated total number of batches for prompts: 20\n",
      "INFO:databricks.labs.doc_qa.model_generators.model_generator:Generated total number of results: 20\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_rows: 20\n",
      "num_successful_rows: 20\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>is_successful</th>\n",
       "      <th>error_msg</th>\n",
       "      <th>content</th>\n",
       "      <th>temperature</th>\n",
       "      <th>max_tokens</th>\n",
       "      <th>model_name</th>\n",
       "      <th>prompt</th>\n",
       "      <th>generated_question</th>\n",
       "      <th>cleaned_context</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To check if any value is True in a pyspark.pan...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you check if any value is True in a py...</td>\n",
       "      <td>pyspark.pandas.Series.any [](#pysparkpandasser...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To use the dateadd function in Databricks SQL ...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you use the dateadd function in Databr...</td>\n",
       "      <td>dateadd (days) function Returns the date numDa...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To convert strings in a pyspark.pandas Series ...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you convert strings in a pyspark.panda...</td>\n",
       "      <td>pyspark.pandas.Series.str.lower [](#pysparkpan...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To use the getbit function in Spark SQL to get...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you use the getbit function in Spark S...</td>\n",
       "      <td>getbit function Applies to: [check marked yes]...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To use the max_by aggregate function in Databr...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you use the max_by aggregate function ...</td>\n",
       "      <td>max_by aggregate function Applies to: [check m...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To check if a datetime Series in Pandas on Spa...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you check if a datetime Series in Pand...</td>\n",
       "      <td>pyspark.pandas.Series.dt.is_year_end [](#pyspa...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To use the `explode` function in Python to exp...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How to use explode in Python to explode an ar...</td>\n",
       "      <td>pyspark.sql.functions.explode [](#pysparksqlfu...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To register and use a user-defined function (U...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you register and use a user-defined fu...</td>\n",
       "      <td>Functions Applies to: [check marked yes](../.....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To use the groupByKeyAndWindow function in Pyt...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How to use groupByKeyAndWindow in Python to a...</td>\n",
       "      <td>pyspark.streaming.DStream.groupByKeyAndWindow ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To use the any_value aggregate function in Dat...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you use the any_value aggregate functi...</td>\n",
       "      <td>any_value aggregate function Applies to: [chec...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To replace values in a DataFrame column using ...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How to replace values in a DataFrame column u...</td>\n",
       "      <td>pyspark.sql.DataFrameNaFunctions.replace [](#p...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To use the regr_avgx aggregate function in Dat...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you use the regr_avgx aggregate functi...</td>\n",
       "      <td>regr_avgx aggregate function Applies to: [chec...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To use VarcharType in Python, you can import t...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How to use VarcharType in Python?</td>\n",
       "      <td>VarcharType [](#varchartype Permalink to this ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To sort a DataFrame in PySpark by index, you c...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How to sort a DataFrame in PySpark by index?</td>\n",
       "      <td>pyspark.pandas.DataFrame.sort_index [](#pyspar...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To drop a database in Spark SQL, you can use t...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you drop a database in Spark SQL?</td>\n",
       "      <td>### [Spark SQL Guide](sqlprogrammingguide.html...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To load data from CSV files into a Delta Lake ...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you load data from CSV files into a De...</td>\n",
       "      <td>Copy Into (Delta Lake on Databricks) Important...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To resolve the RESOURCE_LIMIT_EXCEEDED error w...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you resolve the RESOURCE_LIMIT_EXCEEDE...</td>\n",
       "      <td>Troubleshoot common sharing issues in Delta Sh...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To use `pyspark.pandas.read_orc` to load an OR...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you use pyspark.pandas.read_orc to loa...</td>\n",
       "      <td>pyspark.pandas.read_orc [](#pysparkpandasreado...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To convert between PySpark and pandas DataFram...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How do you convert between PySpark and pandas...</td>\n",
       "      <td>Convert between PySpark and pandas DataFrames ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>To handle the UNSUPPORTED_SUBQUERY_EXPRESSION_...</td>\n",
       "      <td>0.5</td>\n",
       "      <td>100</td>\n",
       "      <td>gpt-3.5-turbo</td>\n",
       "      <td>You are a helpful assistant good at answering ...</td>\n",
       "      <td>How to handle the UNSUPPORTED_SUBQUERY_EXPRES...</td>\n",
       "      <td>UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY error...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    is_successful error_msg  \\\n",
       "0            True      None   \n",
       "1            True      None   \n",
       "2            True      None   \n",
       "3            True      None   \n",
       "4            True      None   \n",
       "5            True      None   \n",
       "6            True      None   \n",
       "7            True      None   \n",
       "8            True      None   \n",
       "9            True      None   \n",
       "10           True      None   \n",
       "11           True      None   \n",
       "12           True      None   \n",
       "13           True      None   \n",
       "14           True      None   \n",
       "15           True      None   \n",
       "16           True      None   \n",
       "17           True      None   \n",
       "18           True      None   \n",
       "19           True      None   \n",
       "\n",
       "                                              content  temperature  \\\n",
       "0   To check if any value is True in a pyspark.pan...          0.5   \n",
       "1   To use the dateadd function in Databricks SQL ...          0.5   \n",
       "2   To convert strings in a pyspark.pandas Series ...          0.5   \n",
       "3   To use the getbit function in Spark SQL to get...          0.5   \n",
       "4   To use the max_by aggregate function in Databr...          0.5   \n",
       "5   To check if a datetime Series in Pandas on Spa...          0.5   \n",
       "6   To use the `explode` function in Python to exp...          0.5   \n",
       "7   To register and use a user-defined function (U...          0.5   \n",
       "8   To use the groupByKeyAndWindow function in Pyt...          0.5   \n",
       "9   To use the any_value aggregate function in Dat...          0.5   \n",
       "10  To replace values in a DataFrame column using ...          0.5   \n",
       "11  To use the regr_avgx aggregate function in Dat...          0.5   \n",
       "12  To use VarcharType in Python, you can import t...          0.5   \n",
       "13  To sort a DataFrame in PySpark by index, you c...          0.5   \n",
       "14  To drop a database in Spark SQL, you can use t...          0.5   \n",
       "15  To load data from CSV files into a Delta Lake ...          0.5   \n",
       "16  To resolve the RESOURCE_LIMIT_EXCEEDED error w...          0.5   \n",
       "17  To use `pyspark.pandas.read_orc` to load an OR...          0.5   \n",
       "18  To convert between PySpark and pandas DataFram...          0.5   \n",
       "19  To handle the UNSUPPORTED_SUBQUERY_EXPRESSION_...          0.5   \n",
       "\n",
       "    max_tokens     model_name  \\\n",
       "0          100  gpt-3.5-turbo   \n",
       "1          100  gpt-3.5-turbo   \n",
       "2          100  gpt-3.5-turbo   \n",
       "3          100  gpt-3.5-turbo   \n",
       "4          100  gpt-3.5-turbo   \n",
       "5          100  gpt-3.5-turbo   \n",
       "6          100  gpt-3.5-turbo   \n",
       "7          100  gpt-3.5-turbo   \n",
       "8          100  gpt-3.5-turbo   \n",
       "9          100  gpt-3.5-turbo   \n",
       "10         100  gpt-3.5-turbo   \n",
       "11         100  gpt-3.5-turbo   \n",
       "12         100  gpt-3.5-turbo   \n",
       "13         100  gpt-3.5-turbo   \n",
       "14         100  gpt-3.5-turbo   \n",
       "15         100  gpt-3.5-turbo   \n",
       "16         100  gpt-3.5-turbo   \n",
       "17         100  gpt-3.5-turbo   \n",
       "18         100  gpt-3.5-turbo   \n",
       "19         100  gpt-3.5-turbo   \n",
       "\n",
       "                                               prompt  \\\n",
       "0   You are a helpful assistant good at answering ...   \n",
       "1   You are a helpful assistant good at answering ...   \n",
       "2   You are a helpful assistant good at answering ...   \n",
       "3   You are a helpful assistant good at answering ...   \n",
       "4   You are a helpful assistant good at answering ...   \n",
       "5   You are a helpful assistant good at answering ...   \n",
       "6   You are a helpful assistant good at answering ...   \n",
       "7   You are a helpful assistant good at answering ...   \n",
       "8   You are a helpful assistant good at answering ...   \n",
       "9   You are a helpful assistant good at answering ...   \n",
       "10  You are a helpful assistant good at answering ...   \n",
       "11  You are a helpful assistant good at answering ...   \n",
       "12  You are a helpful assistant good at answering ...   \n",
       "13  You are a helpful assistant good at answering ...   \n",
       "14  You are a helpful assistant good at answering ...   \n",
       "15  You are a helpful assistant good at answering ...   \n",
       "16  You are a helpful assistant good at answering ...   \n",
       "17  You are a helpful assistant good at answering ...   \n",
       "18  You are a helpful assistant good at answering ...   \n",
       "19  You are a helpful assistant good at answering ...   \n",
       "\n",
       "                                   generated_question  \\\n",
       "0    How do you check if any value is True in a py...   \n",
       "1    How do you use the dateadd function in Databr...   \n",
       "2    How do you convert strings in a pyspark.panda...   \n",
       "3    How do you use the getbit function in Spark S...   \n",
       "4    How do you use the max_by aggregate function ...   \n",
       "5    How do you check if a datetime Series in Pand...   \n",
       "6    How to use explode in Python to explode an ar...   \n",
       "7    How do you register and use a user-defined fu...   \n",
       "8    How to use groupByKeyAndWindow in Python to a...   \n",
       "9    How do you use the any_value aggregate functi...   \n",
       "10   How to replace values in a DataFrame column u...   \n",
       "11   How do you use the regr_avgx aggregate functi...   \n",
       "12                  How to use VarcharType in Python?   \n",
       "13       How to sort a DataFrame in PySpark by index?   \n",
       "14           How do you drop a database in Spark SQL?   \n",
       "15   How do you load data from CSV files into a De...   \n",
       "16   How do you resolve the RESOURCE_LIMIT_EXCEEDE...   \n",
       "17   How do you use pyspark.pandas.read_orc to loa...   \n",
       "18   How do you convert between PySpark and pandas...   \n",
       "19   How to handle the UNSUPPORTED_SUBQUERY_EXPRES...   \n",
       "\n",
       "                                      cleaned_context  \n",
       "0   pyspark.pandas.Series.any [](#pysparkpandasser...  \n",
       "1   dateadd (days) function Returns the date numDa...  \n",
       "2   pyspark.pandas.Series.str.lower [](#pysparkpan...  \n",
       "3   getbit function Applies to: [check marked yes]...  \n",
       "4   max_by aggregate function Applies to: [check m...  \n",
       "5   pyspark.pandas.Series.dt.is_year_end [](#pyspa...  \n",
       "6   pyspark.sql.functions.explode [](#pysparksqlfu...  \n",
       "7   Functions Applies to: [check marked yes](../.....  \n",
       "8   pyspark.streaming.DStream.groupByKeyAndWindow ...  \n",
       "9   any_value aggregate function Applies to: [chec...  \n",
       "10  pyspark.sql.DataFrameNaFunctions.replace [](#p...  \n",
       "11  regr_avgx aggregate function Applies to: [chec...  \n",
       "12  VarcharType [](#varchartype Permalink to this ...  \n",
       "13  pyspark.pandas.DataFrame.sort_index [](#pyspar...  \n",
       "14  ### [Spark SQL Guide](sqlprogrammingguide.html...  \n",
       "15  Copy Into (Delta Lake on Databricks) Important...  \n",
       "16  Troubleshoot common sharing issues in Delta Sh...  \n",
       "17  pyspark.pandas.read_orc [](#pysparkpandasreado...  \n",
       "18  Convert between PySpark and pandas DataFrames ...  \n",
       "19  UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY error...  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from databricks.labs.doc_qa.llm_utils import PromptTemplate\n",
    "from databricks.labs.doc_qa.model_generators.model_generator import OpenAiModelGenerator\n",
    "\n",
    "small_dataset = small_dataset.head(20)\n",
    "\n",
    "prompt_template = PromptTemplate(\"\"\"You are a helpful assistant good at answering questions based on context. Please make sure the answer is correct, comprehensive and reader-friendly. \n",
    "\n",
    "Here is the question: \n",
    "{generated_question}\n",
    "\n",
    "\n",
    "Here is the context:\n",
    "{cleaned_context}\n",
    "\n",
    "Please give your answer:\n",
    "\"\"\")\n",
    "                                 \n",
    "\n",
    "model_generator = OpenAiModelGenerator(prompt_formatter=prompt_template,model_name=\"gpt-3.5-turbo\", concurrency=10)\n",
    "\n",
    "result = model_generator.run_tasks(input_df=small_dataset,temperature=0.5, max_tokens=100)\n",
    "\n",
    "result_df = result.to_dataframe()\n",
    "print(result.summary())\n",
    "result_df"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "docs3.9",
   "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.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}