sql_gen.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. 导入模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "\n",
    "# 选择模型 （llama3_8b, llama3_70b, qwen2_7b, qwen2_72b, gpt_4o_mini）\n",
    "model = \"llama3_8b\"\n",
    "\n",
    "module = importlib.import_module('core.models')\n",
    "LLM = getattr(module, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. 构建 输入 Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8034\n",
      "A database 'real_estate_properties' has 5 tables named 'Ref_Feature_Types, Ref_Property_Types, Other_Available_Features, Properties, Other_Property_Features'.\n",
      "Ref_Feature_Types table has columns: 'feature_type_code, feature_type_name'.\n",
      "Ref_Property_Types table has columns: 'property_type_code, property_type_description'.\n",
      "Other_Available_Features table has columns: 'feature_id, feature_type_code, feature_name, feature_description'.\n",
      "Properties table has columns: 'property_id, property_type_code, date_on_market, date_sold, property_name, property_address, room_count, vendor_requested_price, buyer_offered_price, agreed_selling_price, apt_feature_1, apt_feature_2, apt_feature_3, fld_feature_1, fld_feature_2, fld_feature_3, hse_feature_1, hse_feature_2, hse_feature_3, oth_feature_1, oth_feature_2, oth_feature_3, shp_feature_1, shp_feature_2, shp_feature_3, other_property_details'.\n",
      "Other_Property_Features table has columns: 'property_id, feature_id, property_feature_description'.\n",
      "\n",
      "Gave me the SQL query: \"What are the names of properties that are either houses or apartments with more than 1 room?\"\n",
      "No need explanation.\n",
      "Please output the a brief SQL in the following format:\n",
      "```sql\n",
      "...\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "from core.utils import *\n",
    "\n",
    "# 选择数据库 （spider, bird）\n",
    "dataset = 'spider'\n",
    "\n",
    "prompt_list = prompt_construction(dataset)\n",
    "\n",
    "print(len(prompt_list))\n",
    "print(prompt_list[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. 生成 SQL，存储"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------0 sample_num\n",
      "```sql\n",
      "SELECT COUNT(h.head_ID)\n",
      "FROM head h\n",
      "JOIN management m ON h.head_ID = m.head_ID\n",
      "WHERE h.age > 56;\n",
      "```\n",
      "-------------------------1 sample_num\n",
      "```sql\n",
      "SELECT h.name, h.born_state, h.age\n",
      "FROM head h\n",
      "JOIN management m ON h.head_ID = m.head_ID\n",
      "JOIN department d ON m.department_ID = d.Department_ID\n",
      "ORDER BY h.age;\n",
      "```\n",
      "-------------------------2 sample_num\n",
      "```sql\n",
      "SELECT EXTRACT(YEAR FROM Creation) AS Creation_Year, Name, Budget_in_Billions\n",
      "FROM department;\n",
      "```\n",
      "-------------------------3 sample_num\n",
      "```sql\n",
      "SELECT MAX(Budget_in_Billions) AS Max_Budget, MIN(Budget_in_Billions) AS Min_Budget\n",
      "FROM department;\n",
      "```\n",
      "-------------------------4 sample_num\n",
      "```sql\n",
      "SELECT AVG(Num_Employees)\n",
      "FROM department\n",
      "WHERE Ranking BETWEEN 10 AND 15;\n",
      "```\n",
      "-------------------------5 sample_num\n",
      "```sql\n",
      "SELECT h.name\n",
      "FROM head h\n",
      "WHERE h.born_state <> 'California';\n",
      "```\n",
      "-------------------------6 sample_num\n",
      "```sql\n",
      "SELECT DISTINCT YEAR(d.Creation)\n",
      "FROM department d\n",
      "JOIN management m ON d.Department_ID = m.department_ID\n",
      "JOIN head h ON m.head_ID = h.head_ID\n",
      "WHERE h.born_state = 'Alabama';\n",
      "```\n",
      "-------------------------7 sample_num\n",
      "```sql\n",
      "SELECT born_state\n",
      "FROM head\n",
      "GROUP BY born_state\n",
      "HAVING COUNT(*) >= 3;\n",
      "```\n",
      "-------------------------8 sample_num\n",
      "```sql\n",
      "SELECT EXTRACT(YEAR FROM Creation) AS year, COUNT(*) AS num_departments\n",
      "FROM department\n",
      "GROUP BY year\n",
      "ORDER BY num_departments DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------9 sample_num\n",
      "```sql\n",
      "SELECT d.Name, d.Num_Employees\n",
      "FROM department d\n",
      "JOIN management m ON d.Department_ID = m.department_ID\n",
      "JOIN head h ON m.head_ID = h.head_ID\n",
      "WHERE m.temporary_acting = 'Yes';\n",
      "```\n",
      "-------------------------10 sample_num\n",
      "```sql\n",
      "SELECT COUNT(temporary_acting) \n",
      "FROM management;\n",
      "```\n",
      "-------------------------11 sample_num\n",
      "```sql\n",
      "SELECT COUNT(*)\n",
      "FROM department d\n",
      "LEFT JOIN management m ON d.Department_ID = m.department_ID\n",
      "WHERE m.head_ID IS NULL;\n",
      "```\n",
      "-------------------------12 sample_num\n",
      "```sql\n",
      "SELECT DISTINCT h.age\n",
      "FROM head h\n",
      "JOIN management m ON h.head_ID = m.head_ID\n",
      "WHERE m.temporary_acting = 'acting';\n",
      "```\n",
      "-------------------------13 sample_num\n",
      "```sql\n",
      "SELECT born_state\n",
      "FROM head\n",
      "WHERE head_ID IN (\n",
      "  SELECT h.head_ID\n",
      "  FROM department d\n",
      "  JOIN management m ON d.Department_ID = m.department_ID\n",
      "  JOIN head h ON m.head_ID = h.head_ID\n",
      "  WHERE d.Name IN ('Treasury', 'Homeland Security')\n",
      ")\n",
      "GROUP BY born_state\n",
      "HAVING COUNT(DISTINCT d.Name) = 2;\n",
      "```\n",
      "-------------------------14 sample_num\n",
      "```sql\n",
      "SELECT d.Department_ID, d.Name, COUNT(h.head_ID) AS num_heads\n",
      "FROM department d\n",
      "JOIN management m ON d.Department_ID = m.department_ID\n",
      "JOIN head h ON m.head_ID = h.head_ID\n",
      "GROUP BY d.Department_ID, d.Name\n",
      "HAVING COUNT(h.head_ID) > 1;\n",
      "```\n",
      "-------------------------15 sample_num\n",
      "```sql\n",
      "SELECT head_ID, name \n",
      "FROM head \n",
      "WHERE name LIKE '%Ha%';\n",
      "```\n",
      "-------------------------16 sample_num\n",
      "```sql\n",
      "SELECT COUNT(*) \n",
      "FROM farm;\n",
      "```\n",
      "-------------------------17 sample_num\n",
      "```sql\n",
      "SELECT COUNT(*) \n",
      "FROM farm;\n",
      "```\n",
      "-------------------------18 sample_num\n",
      "```sql\n",
      "SELECT SUM(Total_Horses) \n",
      "FROM farm \n",
      "GROUP BY Year \n",
      "ORDER BY SUM(Total_Horses) ASC;\n",
      "```\n",
      "-------------------------19 sample_num\n",
      "```sql\n",
      "SELECT Farm_ID, SUM(Total_Horses) AS Total_Horses_Record\n",
      "FROM farm\n",
      "GROUP BY Farm_ID\n",
      "ORDER BY Total_Horses_Record ASC;\n",
      "```\n",
      "-------------------------20 sample_num\n",
      "```sql\n",
      "SELECT Hosts\n",
      "FROM farm_competition\n",
      "WHERE Theme <> 'Aliens';\n",
      "```\n",
      "-------------------------21 sample_num\n",
      "```sql\n",
      "SELECT Hosts\n",
      "FROM farm_competition\n",
      "WHERE Theme != 'Aliens';\n",
      "```\n",
      "-------------------------22 sample_num\n",
      "```sql\n",
      "SELECT Year, Theme\n",
      "FROM farm_competition\n",
      "ORDER BY Year ASC;\n",
      "```\n",
      "-------------------------23 sample_num\n",
      "```sql\n",
      "SELECT Theme, Year\n",
      "FROM farm_competition\n",
      "ORDER BY Year ASC;\n",
      "```\n",
      "-------------------------24 sample_num\n",
      "```sql\n",
      "SELECT AVG(Working_Horses) \n",
      "FROM farm \n",
      "WHERE Total_Horses > 5000;\n",
      "```\n",
      "-------------------------25 sample_num\n",
      "```sql\n",
      "SELECT AVG(Working_Horses) \n",
      "FROM farm \n",
      "WHERE Total_Horses > 5000;\n",
      "```\n",
      "-------------------------26 sample_num\n",
      "```sql\n",
      "SELECT MAX(Cows), MIN(Cows)\n",
      "FROM farm;\n",
      "```\n",
      "-------------------------27 sample_num\n",
      "```sql\n",
      "SELECT MAX(Cows), MIN(Cows) \n",
      "FROM farm;\n",
      "```\n",
      "-------------------------28 sample_num\n",
      "```sql\n",
      "SELECT COUNT(DISTINCT Status) \n",
      "FROM city;\n",
      "```\n",
      "-------------------------29 sample_num\n",
      "```sql\n",
      "SELECT COUNT(DISTINCT Status) \n",
      "FROM city;\n",
      "```\n",
      "-------------------------30 sample_num\n",
      "```sql\n",
      "SELECT Official_Name\n",
      "FROM city\n",
      "ORDER BY Population DESC;\n",
      "```\n",
      "-------------------------31 sample_num\n",
      "```sql\n",
      "SELECT Official_Name\n",
      "FROM city\n",
      "ORDER BY Population DESC;\n",
      "```\n",
      "-------------------------32 sample_num\n",
      "```sql\n",
      "SELECT Official_Name, Status\n",
      "FROM city\n",
      "ORDER BY Population DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------33 sample_num\n",
      "```sql\n",
      "SELECT Official_Name, Status\n",
      "FROM city\n",
      "WHERE Population = (SELECT MAX(Population) FROM city);\n",
      "```\n",
      "-------------------------34 sample_num\n",
      "```sql\n",
      "SELECT fc.Year, c.Official_Name\n",
      "FROM farm_competition fc\n",
      "JOIN city c ON fc.Host_city_ID = c.City_ID;\n",
      "```\n",
      "-------------------------35 sample_num\n",
      "```sql\n",
      "SELECT fc.Year, c.Official_Name\n",
      "FROM farm_competition fc\n",
      "JOIN city c ON fc.Host_city_ID = c.City_ID;\n",
      "```\n",
      "-------------------------36 sample_num\n",
      "```sql\n",
      "SELECT c.Official_Name\n",
      "FROM city c\n",
      "JOIN farm_competition fc ON c.City_ID = fc.Host_city_ID\n",
      "GROUP BY c.Official_Name\n",
      "HAVING COUNT(DISTINCT fc.Competition_ID) > 1;\n",
      "```\n",
      "-------------------------37 sample_num\n",
      "```sql\n",
      "SELECT c.Official_Name\n",
      "FROM city c\n",
      "JOIN farm_competition fc ON c.City_ID = fc.Host_city_ID\n",
      "GROUP BY c.Official_Name\n",
      "HAVING COUNT(DISTINCT fc.Competition_ID) > 1;\n",
      "```\n",
      "-------------------------38 sample_num\n",
      "```sql\n",
      "SELECT c.Status\n",
      "FROM city c\n",
      "JOIN (\n",
      "  SELECT Host_city_ID, COUNT(*) as count\n",
      "  FROM farm_competition\n",
      "  GROUP BY Host_city_ID\n",
      "  ORDER BY count DESC\n",
      "  LIMIT 1\n",
      ") AS sub ON c.City_ID = sub.Host_city_ID;\n",
      "```\n",
      "-------------------------39 sample_num\n",
      "```sql\n",
      "SELECT c.Status\n",
      "FROM city c\n",
      "JOIN (\n",
      "  SELECT Host_city_ID, COUNT(*) as num_competitions\n",
      "  FROM farm_competition\n",
      "  GROUP BY Host_city_ID\n",
      "  ORDER BY num_competitions DESC\n",
      "  LIMIT 1\n",
      ") AS fc ON c.City_ID = fc.Host_city_ID;\n",
      "```\n",
      "-------------------------40 sample_num\n",
      "```sql\n",
      "SELECT fc.Theme\n",
      "FROM farm_competition fc\n",
      "JOIN city c ON fc.Host_city_ID = c.City_ID\n",
      "WHERE c.Population > 1000;\n",
      "```\n",
      "-------------------------41 sample_num\n",
      "```sql\n",
      "SELECT fc.Theme\n",
      "FROM farm_competition fc\n",
      "JOIN city c ON fc.Host_city_ID = c.City_ID\n",
      "WHERE c.Population > 1000;\n",
      "```\n",
      "-------------------------42 sample_num\n",
      "```sql\n",
      "SELECT Status, AVG(Population) AS Average_Population\n",
      "FROM city\n",
      "GROUP BY Status;\n",
      "```\n",
      "-------------------------43 sample_num\n",
      "```sql\n",
      "SELECT Status, AVG(Population) AS Average_Population\n",
      "FROM city\n",
      "GROUP BY Status;\n",
      "```\n",
      "-------------------------44 sample_num\n",
      "```sql\n",
      "SELECT Status, COUNT(*) as Num_Cities\n",
      "FROM city\n",
      "GROUP BY Status\n",
      "ORDER BY Num_Cities DESC;\n",
      "```\n",
      "-------------------------45 sample_num\n",
      "```sql\n",
      "SELECT Status, COUNT(*) as frequency\n",
      "FROM city\n",
      "GROUP BY Status\n",
      "ORDER BY frequency ASC;\n",
      "```\n",
      "-------------------------46 sample_num\n",
      "```sql\n",
      "SELECT Status, COUNT(*) as count\n",
      "FROM city\n",
      "GROUP BY Status\n",
      "ORDER BY count DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------47 sample_num\n",
      "```sql\n",
      "SELECT Status, COUNT(*) as count\n",
      "FROM city\n",
      "GROUP BY Status\n",
      "ORDER BY count DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------48 sample_num\n",
      "```sql\n",
      "SELECT Official_Name\n",
      "FROM city\n",
      "WHERE City_ID NOT IN (SELECT Host_city_ID FROM farm_competition);\n",
      "```\n",
      "-------------------------49 sample_num\n",
      "```sql\n",
      "SELECT Official_Name\n",
      "FROM city\n",
      "WHERE City_ID NOT IN (\n",
      "  SELECT Host_city_ID\n",
      "  FROM farm_competition\n",
      ");\n",
      "```\n",
      "-------------------------50 sample_num\n",
      "```sql\n",
      "SELECT DISTINCT Status \n",
      "FROM city \n",
      "WHERE Population > 1500 AND Population < 5000;\n",
      "```\n",
      "-------------------------51 sample_num\n",
      "```sql\n",
      "SELECT DISTINCT Status\n",
      "FROM city\n",
      "WHERE Status IN (\n",
      "  SELECT Status\n",
      "  FROM city\n",
      "  WHERE Population > 1500\n",
      ") AND Status IN (\n",
      "  SELECT Status\n",
      "  FROM city\n",
      "  WHERE Population < 500\n",
      ");\n",
      "```\n",
      "-------------------------52 sample_num\n",
      "```sql\n",
      "SELECT Official_Name\n",
      "FROM city\n",
      "WHERE Population > 1500 OR Population < 500;\n",
      "```\n",
      "-------------------------53 sample_num\n",
      "```sql\n",
      "SELECT Official_Name \n",
      "FROM city \n",
      "WHERE Population > 1500 OR Population < 500;\n",
      "```\n",
      "-------------------------54 sample_num\n",
      "```sql\n",
      "SELECT Census_Ranking\n",
      "FROM city\n",
      "WHERE Status != 'Village';\n",
      "```\n",
      "-------------------------55 sample_num\n",
      "```sql\n",
      "SELECT Census_Ranking\n",
      "FROM city\n",
      "WHERE Status != 'Village';\n",
      "```\n",
      "-------------------------56 sample_num\n",
      "```sql\n",
      "SELECT c.course_id, c.course_name, COUNT(scr.student_id) as num_registered_students\n",
      "FROM Student_Course_Registrations scr\n",
      "JOIN Courses c ON scr.course_id = c.course_id\n",
      "GROUP BY scr.course_id, c.course_name\n",
      "ORDER BY num_registered_students DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------57 sample_num\n",
      "```sql\n",
      "SELECT c.course_name\n",
      "FROM Courses c\n",
      "JOIN Student_Course_Registrations r ON c.course_id = r.course_id\n",
      "GROUP BY c.course_name\n",
      "ORDER BY COUNT(r.student_id) DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------58 sample_num\n",
      "```sql\n",
      "SELECT student_id\n",
      "FROM Student_Course_Registrations\n",
      "GROUP BY student_id\n",
      "HAVING COUNT(course_id) = (\n",
      "  SELECT MIN(cnt)\n",
      "  FROM (\n",
      "    SELECT student_id, COUNT(course_id) AS cnt\n",
      "    FROM Student_Course_Registrations\n",
      "    GROUP BY student_id\n",
      "  ) AS subquery\n",
      ")\n",
      "```\n",
      "-------------------------59 sample_num\n",
      "```sql\n",
      "SELECT student_id\n",
      "FROM Student_Course_Registrations\n",
      "GROUP BY student_id\n",
      "HAVING COUNT(course_id) = (\n",
      "  SELECT MIN(cnt)\n",
      "  FROM (\n",
      "    SELECT student_id, COUNT(course_id) AS cnt\n",
      "    FROM Student_Course_Registrations\n",
      "    GROUP BY student_id\n",
      "  ) AS subquery\n",
      ");\n",
      "```\n",
      "-------------------------60 sample_num\n",
      "```sql\n",
      "SELECT p.first_name, p.last_name\n",
      "FROM Candidates c\n",
      "JOIN People p ON c.candidate_id = p.person_id;\n",
      "```\n",
      "-------------------------61 sample_num\n",
      "```sql\n",
      "SELECT p.first_name, p.last_name\n",
      "FROM People p\n",
      "JOIN Candidates c ON p.person_id = c.candidate_id;\n",
      "```\n",
      "-------------------------62 sample_num\n",
      "```sql\n",
      "SELECT s.student_id\n",
      "FROM Students s\n",
      "LEFT JOIN Student_Course_Attendance sca ON s.student_id = sca.student_id\n",
      "WHERE sca.student_id IS NULL;\n",
      "```\n",
      "-------------------------63 sample_num\n",
      "```sql\n",
      "SELECT s.student_id\n",
      "FROM Students s\n",
      "LEFT JOIN Student_Course_Attendance sca ON s.student_id = sca.student_id\n",
      "WHERE sca.student_id IS NULL;\n",
      "```\n",
      "-------------------------64 sample_num\n",
      "```sql\n",
      "SELECT student_id\n",
      "FROM Student_Course_Attendance;\n",
      "```\n",
      "-------------------------65 sample_num\n",
      "```sql\n",
      "SELECT student_id\n",
      "FROM Student_Course_Attendance\n",
      "GROUP BY student_id\n",
      "HAVING COUNT(*) > 0;\n",
      "```\n",
      "-------------------------66 sample_num\n",
      "```sql\n",
      "SELECT s.student_id, c.course_name\n",
      "FROM Students s\n",
      "JOIN Student_Course_Registrations scr ON s.student_id = scr.student_id\n",
      "JOIN Courses c ON scr.course_id = c.course_id;\n",
      "```\n",
      "-------------------------67 sample_num\n",
      "```sql\n",
      "SELECT s.student_id, p.first_name, p.last_name, c.course_name, scr.registration_date\n",
      "FROM Students s\n",
      "JOIN People p ON s.student_id = p.person_id\n",
      "JOIN Student_Course_Registrations scr ON s.student_id = scr.student_id\n",
      "JOIN Courses c ON scr.course_id = c.course_id\n",
      "ORDER BY scr.registration_date DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------68 sample_num\n",
      "```sql\n",
      "SELECT s.student_id, p.first_name, p.last_name, c.course_name, scr.registration_date\n",
      "FROM Students s\n",
      "JOIN People p ON s.student_id = p.person_id\n",
      "JOIN Student_Course_Registrations scr ON s.student_id = scr.student_id\n",
      "JOIN Courses c ON scr.course_id = c.course_id\n",
      "ORDER BY scr.registration_date DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------69 sample_num\n",
      "```sql\n",
      "SELECT COUNT(DISTINCT s.student_id)\n",
      "FROM Students s\n",
      "JOIN Student_Course_Registrations scr ON s.student_id = scr.student_id\n",
      "JOIN Courses c ON scr.course_id = c.course_id\n",
      "WHERE c.course_name = 'English';\n",
      "```\n",
      "-------------------------70 sample_num\n",
      "```sql\n",
      "SELECT COUNT(DISTINCT s.student_id)\n",
      "FROM Students s\n",
      "JOIN Student_Course_Registrations scr ON s.student_id = scr.student_id\n",
      "JOIN Courses c ON scr.course_id = c.course_id\n",
      "WHERE c.course_name LIKE '%English%';\n",
      "```\n",
      "-------------------------71 sample_num\n",
      "```sql\n",
      "SELECT COUNT(DISTINCT course_id)\n",
      "FROM Student_Course_Registrations\n",
      "WHERE student_id = 171;\n",
      "```\n",
      "-------------------------72 sample_num\n",
      "```sql\n",
      "SELECT COUNT(DISTINCT course_id)\n",
      "FROM Student_Course_Attendance\n",
      "WHERE student_id = 171;\n",
      "```\n",
      "-------------------------73 sample_num\n",
      "```sql\n",
      "SELECT c.candidate_id\n",
      "FROM Candidates c\n",
      "JOIN People p ON c.candidate_id = p.person_id\n",
      "WHERE p.email_address = 'stanley.monahan@example.org';\n",
      "```\n",
      "-------------------------74 sample_num\n",
      "```sql\n",
      "SELECT c.candidate_id\n",
      "FROM Candidates c\n",
      "JOIN People p ON c.candidate_id = p.person_id\n",
      "WHERE p.email_address = 'stanley.monahan@example.org';\n",
      "```\n",
      "-------------------------75 sample_num\n",
      "```sql\n",
      "SELECT ca.candidate_id\n",
      "FROM Candidate_Assessments ca\n",
      "JOIN Candidates c ON ca.candidate_id = c.candidate_id\n",
      "JOIN Student_Course_Registrations scr ON c.candidate_id = scr.student_id\n",
      "JOIN Courses co ON scr.course_id = co.course_id\n",
      "ORDER BY scr.registration_date DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------76 sample_num\n",
      "```sql\n",
      "SELECT ca.candidate_id\n",
      "FROM Candidate_Assessments ca\n",
      "JOIN Candidates c ON ca.candidate_id = c.candidate_id\n",
      "JOIN Student_Course_Registrations scr ON c.candidate_id = scr.student_id\n",
      "JOIN Courses co ON scr.course_id = co.course_id\n",
      "ORDER BY ca.assessment_date DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------77 sample_num\n",
      "```sql\n",
      "SELECT s.student_id, p.first_name, p.last_name, COUNT(sc.course_id) as total_courses\n",
      "FROM Students s\n",
      "JOIN People p ON s.student_id = p.person_id\n",
      "JOIN Student_Course_Registrations sc ON s.student_id = sc.student_id\n",
      "GROUP BY s.student_id, p.first_name, p.last_name\n",
      "ORDER BY total_courses DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------78 sample_num\n",
      "```sql\n",
      "SELECT s.student_id, p.first_name, p.middle_name, p.last_name, COUNT(cr.course_id) AS num_courses\n",
      "FROM Students s\n",
      "JOIN People p ON s.student_id = p.person_id\n",
      "JOIN Student_Course_Registrations cr ON s.student_id = cr.student_id\n",
      "GROUP BY s.student_id, p.first_name, p.middle_name, p.last_name\n",
      "ORDER BY num_courses DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------79 sample_num\n",
      "```sql\n",
      "SELECT s.student_id, COUNT(cr.course_id) AS num_registered_courses\n",
      "FROM Students s\n",
      "JOIN Student_Course_Registrations cr ON s.student_id = cr.student_id\n",
      "GROUP BY s.student_id;\n",
      "```\n",
      "-------------------------80 sample_num\n",
      "```sql\n",
      "SELECT s.student_id, COUNT(cr.course_id) AS num_courses_registered\n",
      "FROM Students s\n",
      "JOIN Student_Course_Registrations cr ON s.student_id = cr.student_id\n",
      "GROUP BY s.student_id;\n",
      "```\n",
      "-------------------------81 sample_num\n",
      "```sql\n",
      "SELECT c.course_name, COUNT(scr.student_id) AS num_registered_students\n",
      "FROM Courses c\n",
      "JOIN Student_Course_Registrations scr ON c.course_id = scr.course_id\n",
      "GROUP BY c.course_name;\n",
      "```\n",
      "-------------------------82 sample_num\n",
      "```sql\n",
      "SELECT c.course_id, c.course_name, COUNT(scr.student_id) AS num_students_registered\n",
      "FROM Courses c\n",
      "JOIN Student_Course_Registrations scr ON c.course_id = scr.course_id\n",
      "GROUP BY c.course_id, c.course_name;\n",
      "```\n",
      "-------------------------83 sample_num\n",
      "```sql\n",
      "SELECT candidate_id\n",
      "FROM Candidate_Assessments\n",
      "WHERE assessment_outcome_code = 'Pass';\n",
      "```\n",
      "-------------------------84 sample_num\n",
      "```sql\n",
      "SELECT candidate_id \n",
      "FROM Candidate_Assessments \n",
      "WHERE asessment_outcome_code = 'Pass';\n",
      "```\n",
      "-------------------------85 sample_num\n",
      "```sql\n",
      "SELECT p.cell_mobile_number\n",
      "FROM People p\n",
      "JOIN Candidates c ON p.person_id = c.candidate_id\n",
      "JOIN Candidate_Assessments ca ON c.candidate_id = ca.candidate_id\n",
      "WHERE ca.assessment_outcome_code = 'Fail';\n",
      "```\n",
      "-------------------------86 sample_num\n",
      "```sql\n",
      "SELECT p.cell_mobile_number\n",
      "FROM People p\n",
      "JOIN Candidates c ON p.person_id = c.person_id\n",
      "JOIN Candidate_Assessments ca ON c.candidate_id = ca.candidate_id\n",
      "WHERE ca.assessment_outcome_code = 'Fail';\n",
      "```\n",
      "-------------------------87 sample_num\n",
      "```sql\n",
      "SELECT student_id \n",
      "FROM Student_Course_Registrations \n",
      "WHERE course_id = 301;\n",
      "```\n",
      "-------------------------88 sample_num\n",
      "```sql\n",
      "SELECT student_id \n",
      "FROM Student_Course_Registrations \n",
      "WHERE course_id = 301;\n",
      "```\n",
      "-------------------------89 sample_num\n",
      "```sql\n",
      "SELECT student_id\n",
      "FROM Student_Course_Registrations\n",
      "WHERE course_id = 301\n",
      "ORDER BY registration_date DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------90 sample_num\n",
      "```sql\n",
      "SELECT student_id\n",
      "FROM Student_Course_Registrations\n",
      "WHERE course_id = 301\n",
      "ORDER BY registration_date DESC\n",
      "LIMIT 1;\n",
      "```\n",
      "-------------------------91 sample_num\n",
      "```sql\n",
      "SELECT DISTINCT city\n",
      "FROM Addresses\n",
      "JOIN People_Addresses ON Addresses.address_id = People_Addresses.address_id\n",
      "JOIN People ON People_Addresses.person_id = People.person_id;\n",
      "```\n",
      "-------------------------92 sample_num\n",
      "```sql\n",
      "SELECT DISTINCT city \n",
      "FROM Addresses;\n",
      "```\n",
      "-------------------------93 sample_num\n",
      "```sql\n",
      "SELECT DISTINCT a.city\n",
      "FROM Addresses a\n",
      "JOIN People_Addresses pa ON a.address_id = pa.address_id\n",
      "JOIN People p ON pa.person_id = p.person_id\n",
      "JOIN Students s ON p.person_id = s.student_id;\n",
      "```\n",
      "-------------------------94 sample_num\n",
      "```sql\n",
      "SELECT DISTINCT a.city\n",
      "FROM Addresses a\n",
      "JOIN People_Addresses pa ON a.address_id = pa.address_id\n",
      "JOIN People p ON pa.person_id = p.person_id\n",
      "JOIN Students s ON p.person_id = s.student_id;\n",
      "```\n",
      "-------------------------95 sample_num\n",
      "```sql\n",
      "SELECT course_name\n",
      "FROM Courses\n",
      "ORDER BY course_name ASC;\n",
      "```\n",
      "-------------------------96 sample_num\n",
      "```sql\n",
      "SELECT course_name\n",
      "FROM Courses\n",
      "ORDER BY course_name ASC;\n",
      "```\n",
      "-------------------------97 sample_num\n",
      "```sql\n",
      "SELECT first_name\n",
      "FROM People\n",
      "ORDER BY first_name ASC;\n",
      "```\n",
      "-------------------------98 sample_num\n",
      "```sql\n",
      "SELECT first_name\n",
      "FROM People\n",
      "ORDER BY first_name ASC;\n",
      "```\n",
      "-------------------------99 sample_num\n",
      "```sql\n",
      "SELECT student_id\n",
      "FROM Student_Course_Registrations\n",
      "UNION\n",
      "SELECT student_id\n",
      "FROM Student_Course_Attendance;\n",
      "```\n",
      "-------------------------100 sample_num\n",
      "```sql\n",
      "SELECT student_id\n",
      "FROM Student_Course_Registrations\n",
      "UNION\n",
      "SELECT student_id\n",
      "FROM Student_Course_Attendance;\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "save_path = f'./results/sql_gen_ori/{model}_{dataset}_ori.txt'\n",
    "\n",
    "with open(save_path, 'a')as f:\n",
    "    for index, prompt in enumerate(prompt_list):\n",
    "        if index >= 0 : # 如果数据太多，可以设置 index， 分段生成，不必一次性生成\n",
    "            print(\"-------------------------\" + str(index) + ' sample_num')\n",
    "            output = LLM(prompt)\n",
    "            print(output)\n",
    "\n",
    "            f.write(f'-------------------------{str(index)} sample_num\\n')\n",
    "            f.write(output + '\\n')\n",
    "\n",
    "            # time.sleep(0.5)           # 如果报错：访问频率过快，则强制休眠，模型较小时使用"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. 输出清洗"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "101\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'SELECT student_id FROM Student_Course_Registrations UNION SELECT student_id FROM Student_Course_Attendance;'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from core.utils import *\n",
    "\n",
    "# save_path = f'./results/sql_gen_ori/{model}_{dataset}_ori.txt'\n",
    "sql_list = sql_match(save_path) \n",
    "\n",
    "# 存入文件\n",
    "with open(f'./results/sql_gen_clean/{model}_{dataset}_clean.txt', 'w') as f:\n",
    "    for sql in sql_list:\n",
    "        f.write(sql + '\\n')\n",
    "\n",
    "print(len(sql_list))\n",
    "sql_list[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. 结果评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pre_num: 101;\n",
      "ex: 0.7129\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[0,\n",
       " 0,\n",
       " -1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " -1,\n",
       " 1,\n",
       " -1,\n",
       " 1,\n",
       " 0,\n",
       " 1,\n",
       " 0,\n",
       " -1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 0,\n",
       " 1,\n",
       " 1,\n",
       " 0,\n",
       " 0,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 0,\n",
       " 0,\n",
       " 0,\n",
       " 0,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 0,\n",
       " 1,\n",
       " 0,\n",
       " 0,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 0,\n",
       " 0,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 0,\n",
       " 0,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 0,\n",
       " -1,\n",
       " 1,\n",
       " -1,\n",
       " -1,\n",
       " 1,\n",
       " 1,\n",
       " 0,\n",
       " 0,\n",
       " 1,\n",
       " 0,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1,\n",
       " 1]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from core.evaluate import *\n",
    "\n",
    "pre_sql_path = './results/sql_gen_clean/llama3_70b_spider_clean.txt'\n",
    "gold_sql_path = './dataset/data_instance/spider_all_gold.txt'\n",
    "db_path = './dataset/spider/database'\n",
    "\n",
    "res_list = ex_evaluation(pre_sql_path, gold_sql_path, db_path)\n",
    "\n",
    "print_context = f\"\"\"pre_num: {len(res_list)};\n",
    "ex: {res_list.count(1)/len(res_list):.4f}\"\"\"\n",
    "\n",
    "print(print_context)\n",
    "res_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. 存入结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "sql_incor_path = f\"./results/sql_incor/{model}_{dataset}_incor.json\"\n",
    "\n",
    "with open(pre_sql_path, 'r') as f:\n",
    "    pre_sql_list = [s.strip() for s in f.readlines()]\n",
    "\n",
    "sql_incor_list = []\n",
    "for index, _ in enumerate(res_list):\n",
    "    if res_list[index] == 0:  # 只考虑执行结果不正确的，排除无法执行以及超时错误\n",
    "        sql_incor_list.append(pre_sql_list[index])\n",
    "\n",
    "data = {\n",
    "    \"res_list\": res_list,\n",
    "    \"ex\": res_list.count(1)/len(res_list),\n",
    "    \"sql_incor\": sql_incor_list,\n",
    "    \"sql_gen_clean\": pre_sql_list\n",
    "}\n",
    "\n",
    "with open(sql_incor_path, 'w')as f:\n",
    "    json.dump(data, f, indent=4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cot",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}