code: """Generate SQL queries using LlamaIndex."""
  import argparse
  import json
  import logging
  import os
  import re
  from typing import Any, cast
  
  from llama_index import LLMPredictor, SQLDatabase
  from llama_index.indices import SQLStructStoreIndex
  from llama_index.llms.openai import OpenAI
  from sqlalchemy import create_engine, text
  from tqdm import tqdm
  
  logging.getLogger("root").setLevel(logging.WARNING)
  
  
  _spaces = re.compile(r"\s+")
  _newlines = re.compile(r"\n+")
  
  
  def _generate_sql(
      llama_index: SQLStructStoreIndex,
      nl_query_text: str,
  ) -> str:
      """Generate SQL query for the given NL query text."""
      query_engine = llama_index.as_query_engine()
      response = query_engine.query(nl_query_text)
      if (
          response.metadata is None
          or "sql_query" not in response.metadata
          or response.metadata["sql_query"] is None
      ):
          raise RuntimeError("No SQL query generated.")
      query = response.metadata["sql_query"]
      # Remove newlines and extra spaces.
      query = _newlines.sub(" ", query)
      query = _spaces.sub(" ", query)
      return query.strip()
  
  
  def generate_sql(llama_indexes: dict, examples: list, output_file: str) -> None:
      """Generate SQL queries for the given examples and write them to the output file."""
      with open(output_file, "w") as f:
          for example in tqdm(examples, desc=f"Generating {output_file}"):
              db_name = example["db_id"]
              nl_query_text = example["question"]
              try:
                  sql_query = _generate_sql(llama_indexes[db_name], nl_query_text)
              except Exception as e:
                  print(
                      f"Failed to generate SQL query for question: "
                      f"{example['question']} on database: {example['db_id']}."
                  )
                  print(e)
                  sql_query = "ERROR"
              f.write(sql_query + "\n")
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser(
          description="Generate SQL queries using LlamaIndex."
      )
      parser.add_argument(
          "--input", type=str, required=True, help="Path to the spider dataset directory."
      )
      parser.add_argument(
          "--output",
          type=str,
          required=True,
          help="Path to the output directory of generated SQL files,"
          " one query on each line, "
          "to be compared with the *_gold.sql files in the input directory.",
      )
      parser.add_argument(
          "--model",
          type=str,
          choices=["gpt-4", "gpt-3.5-turbo", "text-davinci-003", "code-davinci-002"],
          required=True,
          help="The model to use for generating SQL queries.",
      )
      args = parser.parse_args()
  
      # Create the output directory if it does not exist.
      if not os.path.exists(args.output):
          os.makedirs(args.output)
  
      # Load the Spider dataset from the input directory.
      with open(os.path.join(args.input, "train_spider.json")) as f:
          train_spider = json.load(f)
      with open(os.path.join(args.input, "train_others.json")) as f:
          train_others = json.load(f)
      with open(os.path.join(args.input, "dev.json")) as f:
          dev = json.load(f)
  
      # Create all necessary SQL database objects.
      databases = {}
      for db in train_spider + train_others + dev:
          db_name = db["db_id"]
          if db_name in databases:
              continue
          db_path = os.path.join(args.input, "database", db_name, db_name + ".sqlite")
          engine = create_engine("sqlite:///" + db_path)
          databases[db_name] = (SQLDatabase(engine=engine), engine)
  
      # Create the LlamaIndexes for all databases.
      llm = OpenAI(model=args.model, temperature=0)
      llm_predictor = LLMPredictor(llm=llm)
      llm_indexes = {}
      for db_name, (db, engine) in databases.items():
          # Get the name of the first table in the database.
          # This is a hack to get a table name for the index, which can use any
          # table in the database.
          with engine.connect() as connection:
              table_name = cast(
                  Any,
                  connection.execute(
                      text("select name from sqlite_master where type = 'table'")
                  ).fetchone(),
              )[0]
              llm_indexes[db_name] = SQLStructStoreIndex.from_documents(
                  documents=[],
                  llm_predictor=llm_predictor,
                  sql_database=db,
                  table_name=table_name,
              )
  
      # Generate SQL queries.
      generate_sql(
          llama_indexes=llm_indexes,
          examples=train_spider + train_others,
          output_file=os.path.join(args.output, "train_pred.sql"),
      )
      generate_sql(
          llama_indexes=llm_indexes,
          examples=dev,
          output_file=os.path.join(args.output, "dev_pred.sql"),
      )
  
apis: llama_index.LLMPredictor llama_index.SQLDatabase llama_index.indices.SQLStructStoreIndex.from_documents llama_index.llms.openai.OpenAI
extract_api: [((413, 431), 're.compile', 're.compile', (['"""\\\\s+"""'], {}), "('\\\\s+')\n", (423, 431), False, 'import re\n'), ((444, 462), 're.compile', 're.compile', (['"""\\\\n+"""'], {}), "('\\\\n+')\n", (454, 462), False, 'import re\n'), ((1926, 2003), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Generate SQL queries using LlamaIndex."""'}), "(description='Generate SQL queries using LlamaIndex.')\n", (1949, 2003), False, 'import argparse\n'), ((3623, 3662), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': 'args.model', 'temperature': '(0)'}), '(model=args.model, temperature=0)\n', (3629, 3662), False, 'from llama_index.llms.openai import OpenAI\n'), ((3683, 3704), 'llama_index.LLMPredictor', 'LLMPredictor', ([], {'llm': 'llm'}), '(llm=llm)\n', (3695, 3704), False, 'from llama_index import LLMPredictor, SQLDatabase\n'), ((349, 374), 'logging.getLogger', 'logging.getLogger', (['"""root"""'], {}), "('root')\n", (366, 374), False, 'import logging\n'), ((1329, 1377), 'tqdm.tqdm', 'tqdm', (['examples'], {'desc': 'f"""Generating {output_file}"""'}), "(examples, desc=f'Generating {output_file}')\n", (1333, 1377), False, 'from tqdm import tqdm\n'), ((2745, 2772), 'os.path.exists', 'os.path.exists', (['args.output'], {}), '(args.output)\n', (2759, 2772), False, 'import os\n'), ((2782, 2806), 'os.makedirs', 'os.makedirs', (['args.output'], {}), '(args.output)\n', (2793, 2806), False, 'import os\n'), ((2954, 2966), 'json.load', 'json.load', (['f'], {}), '(f)\n', (2963, 2966), False, 'import json\n'), ((3057, 3069), 'json.load', 'json.load', (['f'], {}), '(f)\n', (3066, 3069), False, 'import json\n'), ((3142, 3154), 'json.load', 'json.load', (['f'], {}), '(f)\n', (3151, 3154), False, 'import json\n'), ((3375, 3441), 'os.path.join', 'os.path.join', (['args.input', '"""database"""', 'db_name', "(db_name + '.sqlite')"], {}), "(args.input, 'database', db_name, db_name + '.sqlite')\n", (3387, 3441), False, 'import os\n'), ((3459, 3496), 'sqlalchemy.create_engine', 'create_engine', (["('sqlite:///' + db_path)"], {}), "('sqlite:///' + db_path)\n", (3472, 3496), False, 'from sqlalchemy import create_engine, text\n'), ((2878, 2923), 'os.path.join', 'os.path.join', (['args.input', '"""train_spider.json"""'], {}), "(args.input, 'train_spider.json')\n", (2890, 2923), False, 'import os\n'), ((2981, 3026), 'os.path.join', 'os.path.join', (['args.input', '"""train_others.json"""'], {}), "(args.input, 'train_others.json')\n", (2993, 3026), False, 'import os\n'), ((3084, 3120), 'os.path.join', 'os.path.join', (['args.input', '"""dev.json"""'], {}), "(args.input, 'dev.json')\n", (3096, 3120), False, 'import os\n'), ((3527, 3553), 'llama_index.SQLDatabase', 'SQLDatabase', ([], {'engine': 'engine'}), '(engine=engine)\n', (3538, 3553), False, 'from llama_index import LLMPredictor, SQLDatabase\n'), ((4243, 4365), 'llama_index.indices.SQLStructStoreIndex.from_documents', 'SQLStructStoreIndex.from_documents', ([], {'documents': '[]', 'llm_predictor': 'llm_predictor', 'sql_database': 'db', 'table_name': 'table_name'}), '(documents=[], llm_predictor=\n    llm_predictor, sql_database=db, table_name=table_name)\n', (4277, 4365), False, 'from llama_index.indices import SQLStructStoreIndex\n'), ((4588, 4631), 'os.path.join', 'os.path.join', (['args.output', '"""train_pred.sql"""'], {}), "(args.output, 'train_pred.sql')\n", (4600, 4631), False, 'import os\n'), ((4734, 4775), 'os.path.join', 'os.path.join', (['args.output', '"""dev_pred.sql"""'], {}), "(args.output, 'dev_pred.sql')\n", (4746, 4775), False, 'import os\n'), ((4101, 4160), 'sqlalchemy.text', 'text', (['"""select name from sqlite_master where type = \'table\'"""'], {}), '("select name from sqlite_master where type = \'table\'")\n', (4105, 4160), False, 'from sqlalchemy import create_engine, text\n')]