added duckdb-nsql benchmark setup

.gitignore

@@ -0,0 +1,4 @@
+.DS_Store
+__pycache__/
+*.egg-info
+.ipynb_checkpoints/

duckdb-nsql/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

duckdb-nsql/README.md

@@ -0,0 +1,63 @@
+# DuckDB-NSQL
+Numbers Station Text to SQL model for DuckDB.
+
+NSQL is a family of autoregressive open-source foundational models (FMs) that are particularly designed for SQL generation tasks. We are thrilled to introduce DuckDB-NSQL in this repository, an FM tailored for local DuckDB SQL analytics tasks. All model weights can be found on HuggingFace.
+
+| Model Name                            | Size | Link                                                           |
+| --------------------------------------| ---- | -------------------------------------------------------------- |
+| motherduckdb/DuckDB-NSQL-7B-v0.1      | 7B   | [link](https://huggingface.co/motherduckdb/DuckDB-NSQL-7B-v0.1) |
+| motherduckdb/DuckDB-NSQL-7B-v0.1-GGUF | 7B   | [link](https://huggingface.co/motherduckdb/DuckDB-NSQL-7B-v0.1-GGUF)|
+
+## Setup
+To install all the necessary dependencies, please run
+```
+pip install -r requirements.txt
+```
+
+## Usage
+Please refer to the examples in the `examples/` folder to learn how to connect to a local DuckDB and directly query your data. A simple notebook is provided in the `examples/` directory for reference.
+
+To host the model with llama.cpp, please execute the following:
+
+```python
+# Import necessary modules
+from llama_cpp import Llama
+from wurlitzer import pipes
+
+# Set up client with model path and context size
+with pipes() as (out, err):
+    client = Llama(
+        model_path="DuckDB-NSQL-7B-v0.1-q8_0.gguf",
+        n_ctx=2048,
+    )
+```
+
+To load the DuckDB database and query against it, please execute the following:
+
+```python
+# Import necessary modules
+import duckdb
+from utils import generate_sql
+
+# Connect to DuckDB database
+con = duckdb.connect("nyc.duckdb")
+
+# Sample question for SQL generation
+question = "alter taxi table and add struct column with name test and keys a:int, b:double"
+
+# Generate SQL, check validity, and print
+sql = generate_sql(question, con, client)
+print(sql)
+```
+
+## Training Data
+
+The training data for this model consists of two parts: 1) 200k synthetically generated DuckDB SQL queries, based on the DuckDB v.0.9.2 documentation, and 2) labeled text-to-SQL pairs from [NSText2SQL](https://huggingface.co/datasets/NumbersStation/NSText2SQL) transpiled to DuckDB SQL using [sqlglot](https://github.com/tobymao/sqlglot).
+
+## Evaluate the benchmark
+
+Please refer to the `eval/` folder to check the details for evaluating the model against our proposed DuckDB benchmark.
+
+## Acknowledgement
+
+We would like to express our appreciation to all authors of the evaluation scripts. Their work made this project possible.

duckdb-nsql/eval/README.md

@@ -0,0 +1,90 @@
+This folder contains the suite for evaluating the DuckDB-Text2SQL model.
+
+Please install the dependencies listed in the requirements.txt file located in the parent folder.
+
+## Setup
+To evaluate against the benchmark dataset, you need to prepare the evaluation script using this benchmark.
+
+```
+mkdir metrics
+cd metrics
+git clone [email protected]:ElementAI/test-suite-sql-eval.git test_suite_sql_eval
+cd ..
+```
+
+You need to add a new remote to evaluate against duckdb in the test-suite-sql-eval folder. And check the latest duckdb-only branch (640a12975abf75a94e917caca149d56dbc6bcdd7).
+
+```
+git remote add till https://github.com/tdoehmen/test-suite-sql-eval.git
+git fetch till
+git checkout till/duckdb-only
+```
+
+Next, prepare the docs for retrieval.
+```
+mkdir docs
+cd docs
+git clone https://github.com/duckdb/duckdb-web.git
+cd ..
+```
+
+#### Dataset
+The benchmark dataset is located in the `data/` folder and includes all databases (`data/databases`), table schemas (`data/tables.json`), and examples (`data/dev.json`).
+
+#### Eval
+Start a manifest session with the model you want to evaluate.
+
+```bash
+python -m manifest.api.app \
+    --model_type huggingface \
+    --model_generation_type text-generation \
+    --model_name_or_path motherduckdb/DuckDB-NSQL-7B-v0.1 \
+    --fp16 \
+    --device 0
+```
+
+Then, from the `DuckDB-NSQL` main folder, run:
+
+```bash
+python eval/predict.py \
+    predict \
+    eval/data/dev.json \
+    eval/data/tables.json \
+    --output-dir output/ \
+    --stop-tokens ';' \
+    --stop-tokens '--' \
+    --stop-tokens '```' \
+    --stop-tokens '###' \
+    --overwrite-manifest \
+    --manifest-client huggingface \
+    --manifest-connection http://localhost:5000 \
+    --prompt-format duckdbinst
+```
+This will format the prompt using the duckdbinst style.
+
+To evaluate the prediction, first run the following in a Python shell:
+
+```python
+try:
+    import duckdb
+
+    con = duckdb.connect()
+    con.install_extension("httpfs")
+    con.load_extension("httpfs")
+except Exception as e:
+    print(f"Error loading duckdb extensions: {e}")
+```
+
+Then, run the evaluation script:
+
+```bash
+python eval/evaluate.py \
+    evaluate \
+    --gold eval/data/dev.json \
+    --db eval/data/databases/ \
+    --tables eval/data/tables.json \
+    --output-dir output/ \
+    --pred [PREDICITON_FILE]
+```
+
+To view the output, all the information is located in the prediction file in the [output-dir]. Here, `query` is gold and `pred` is predicted.

duckdb-nsql/eval/constants.py

@@ -0,0 +1,37 @@
+"""Constants."""
+
+from prompt_formatters import (
+    DuckDBFormatter,
+    MotherDuckFormatter,
+    DuckDBInstFormatter,
+    DuckDBInstNoShorthandFormatter,
+    RajkumarFormatter,
+    DuckDBChat,
+    DuckDBInstFormatterLlamaShort,
+    DuckDBInstFormatterGraniteShort,
+    DuckDBInstFormatterLlama,
+    DuckDBInstFormatterLlamaBasic,
+    DuckDBInstFormatterGranite,
+    DuckDBInstFormatterPhi,
+    DuckDBInstFormatterGPTmini,
+    DuckDBInstFormatterPhiAzure,
+    DuckDBInstFormatterLlamaSyntax,
+)
+
+PROMPT_FORMATTERS = {
+    "rajkumar": RajkumarFormatter,
+    "duckdb": DuckDBFormatter,
+    "motherduck": MotherDuckFormatter,
+    "duckdbinst": DuckDBInstFormatter,
+    "duckdbinstllamashort": DuckDBInstFormatterLlamaShort,
+    "duckdbinstgraniteshort": DuckDBInstFormatterGraniteShort,
+    "duckdbinstllama": DuckDBInstFormatterLlama,
+    "duckdbinstgranite": DuckDBInstFormatterGranite,
+    "duckdbinstnoshort": DuckDBInstNoShorthandFormatter,
+    "duckdbchat": DuckDBChat,
+    "duckdbinstphi": DuckDBInstFormatterPhi,
+    "duckdbinstgptmini": DuckDBInstFormatterPhi,
+    "duckdbinstphiazure": DuckDBInstFormatterPhiAzure,
+    "duckdbinstllamabasic": DuckDBInstFormatterLlamaBasic,
+    "duckdbinstllamasyntax": DuckDBInstFormatterLlamaSyntax
+}

duckdb-nsql/eval/data/dev.json

@@ -0,0 +1,602 @@
+[
+    {
+        "db_id": "hn",
+        "query": "SELECT COUNT(*) as domain_count, \nSUBSTRING(SPLIT_PART(url, '//', 2), 1, POSITION('/' IN SPLIT_PART(url, '//', 2)) - 1) as domain \nFROM hacker_news\nWHERE url IS NOT NULL GROUP BY domain ORDER BY domain_count DESC LIMIT 10;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "what are the top domains being shared on hacker_news?",
+        "category": "hard"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT c.firstname, c.lastname, COUNT(*) AS num_pcs_bought\nFROM customers c\nJOIN sales s ON c.customer_id = s.customer_id\nJOIN pcs p ON s.model = p.model\nGROUP BY c.customer_id, c.firstname, c.lastname\nORDER BY num_pcs_bought DESC\nLIMIT 1;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Who bought the most PCs, print also the users name?",
+        "category": "medium"
+    },
+    {
+        "db_id": "transactions",
+        "query": "select users.id, users.name, sum(transactions.amount) as balance    from users    join transactions on users.id = transactions.user_id    group by users.id, users.name having balance = 0",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "list the names off account holders who have negative balances",
+        "category": "easy"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT model FROM products WHERE maker = 'B';",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "List only the model number of all products made by maker B.",
+        "category": "easy"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT model FROM products WHERE maker <> 'B';",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "List the model numbers of all products not made by maker B.",
+        "category": "easy"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT AVG(speed) FROM pcs WHERE speed >= 3.00",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Return the average speed all PCs with speed >= 3.00",
+        "category": "easy"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT MAX(price) FROM printers WHERE color = 'TRUE' AND type='laser'",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Return the price of the most expensive color laser printer",
+        "category": "medium"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT MIN(paid) FROM sales WHERE type_of_payment LIKE '%visa%'",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Return the minimum amount paid by customers who used a visa card (debit or credit) to purchase a product",
+        "category": "medium"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT customer_id FROM customers  WHERE firstname LIKE '%e%' OR lastname LIKE '%e%'",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Find the customer_id of customers who have the letter 'e' either in their first name or in their last name",
+        "category": "medium"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT model, price/0.85 AS 'price (USD)'  FROM laptops  WHERE ram >= 1024",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Assume all prices in the table Laptops are in Euro. List the prices of laptops with at least 1024 ram. You should return the price in USD in a column called 'price (USD)'. Assume that 1 USD = 0.85 EURO. Name the price column 'price (USD)'.",
+        "category": "hard"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT maker FROM products GROUP BY maker HAVING COUNT(maker) > 4;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Return a list of makers that make more than four different products.",
+        "category": "medium"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT model FROM laptops WHERE speed > 1.7 ORDER BY speed DESC;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "List all the laptop model numbers that have a speed greater than 1.7 in descending order of speed.",
+        "category": "medium"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT firstname \n        FROM sales \n        JOIN customers ON sales.customer_id = customers.customer_id \n        GROUP BY firstname \n        ORDER BY COUNT(firstname);",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "List firstnames of customers in an ascending order based on the number of purchases made by customers with this firstname.",
+        "category": "medium"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT DISTINCT maker FROM products JOIN pcs ON products.model = pcs.model WHERE ram > 1500;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "List all the makers (with only one entry per maker) who make PCs with RAM greater than 1500.",
+        "category": "medium"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT city, AVG(paid) as 'avg_spend' FROM sales JOIN customers ON sales.customer_id = customers.customer_id GROUP BY city",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Find the city and the average amount of money spent by customers in each city. Name the column for the amount 'avg_spend'",
+        "category": "medium"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT color, MAX(price) as 'max_price' FROM printers GROUP BY color;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Find the maximum price for each color of printer. Name the column for the maximum price 'max_price'",
+        "category": "medium"
+    },
+    {
+        "db_id": "who",
+        "query": "select country_name, max(pm25_concentration) as worst_pm25_for_country\nfrom ambient_air_quality\ngroup by country_name\norder by worst_pm25_for_country desc\nlimit 1",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Find the country with the worst single reading of air quality (highest PM 2.5 value). Show the PM 2.5 value as well.",
+        "category": "medium"
+    },
+    {
+        "db_id": "who",
+        "query": "select country_name, avg(pm25_concentration) as worst_avg_pm25_for_country\nfrom ambient_air_quality\ngroup by country_name\norder by worst_avg_pm25_for_country desc\nlimit 1",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Find the country with the worst average air quality (highest PM 2.5 value). Show the PM 2.5 value as well.",
+        "category": "medium"
+    },
+    {
+        "db_id": "who",
+        "query": "select distinct country_name from ambient_air_quality order by country_name",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Find all countries for which WHO air quality data is available. Sort alphabetically.",
+        "category": "medium"
+    },
+    {
+        "db_id": "who",
+        "query": "select year, avg(pm25_concentration) from ambient_air_quality \nwhere country_name = 'Singapore'\ngroup by year\norder by year",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Find Singapore air quality defined as PM2.5 concentration over time",
+        "category": "medium"
+    },
+    {
+        "db_id": "nyc",
+        "query": "SELECT COLUMNS('^trip_') FROM rideshare LIMIT 10;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "select only the column names from the rideshare table that start with trip_ and return the first 10 values",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "nyc",
+        "query": "SELECT * FROM rideshare USING SAMPLE 1%;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "select a 1% sample from the nyc.rideshare table",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT * EXCLUDE (customer_id) FROM customers;\n",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "select all columns from the customer table, except customer_id",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "nyc",
+        "query": "SUMMARIZE rideshare;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "show summary statistics of the rideshare table",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "none",
+        "query": "SELECT * FROM read_csv_auto(\n'https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv')",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "read a CSV from https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "none",
+        "query": "COPY (SELECT * FROM read_csv_auto(\n'https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv'))\nTO 'titanic.parquet' (FORMAT 'parquet');",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM 'titanic.parquet'",
+        "question": "read a CSV from https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv and convert it to a parquet file called \"titanic\"",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "none",
+        "query": "CREATE TABLE titanic AS (SELECT * FROM read_csv_auto(\n'https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv'))",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM titanic;",
+        "question": "create a table called \"titanic\" from CSV file https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "none",
+        "query": "PRAGMA default_null_order='NULLS LAST';",
+        "setup_sql": ";",
+        "validation_sql": "SELECT current_setting('default_null_order');",
+        "question": "configure duckdb to put null values last when sorting",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "none",
+        "query": "CREATE TABLE IF NOT EXISTS products (\n    maker varchar(10),\n    model varchar(10),\n    type varchar(10));",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM products;",
+        "question": "create a table about products, that contains a maker, model and type column",
+        "category": "ddl"
+    },
+    {
+        "db_id": "product",
+        "query": "INSERT INTO products (maker, model, type)\nVALUES\n    ('A', '1001', 'pc');",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM products;",
+        "question": "add a row with values for model \"1001\" of type \"pc\", from maker \"A\" to products table",
+        "category": "ddl"
+    },
+    {
+        "db_id": "none",
+        "query": "CALL pragma_version();\n",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "get current version of duckdb",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "nyc",
+        "query": "PRAGMA table_info('rideshare');",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "list all columns in table nyc.rideshare",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "nyc",
+        "query": "PRAGMA show_tables;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "show all tables in the curent database",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT customer_id, model, sum(paid) FROM sales GROUP BY ALL",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "how much did each customer spend per model type?",
+        "category": "easy"
+    },
+    {
+        "db_id": "nyc",
+        "query": "SELECT Max(datediff('minute', tpep_pickup_datetime, tpep_dropoff_datetime)) from nyc.taxi",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "What was the longest taxi ride in minutes?",
+        "category": "hard"
+    },
+    {
+        "db_id": "who",
+        "query": "with per_region as (\n   select avg(pm10_concentration) as avg_pm10, who_region from ambient_air_quality group by who_region\n), max_region as (\n select who_region from per_region where  avg_pm10 = (select max(avg_pm10) from per_region)\n), min_city_value_in_max_region as (\n  select min(pm10_concentration) from ambient_air_quality where who_region in (from max_region)\n), min_city_in_max_region as (\n  select city from ambient_air_quality where pm10_concentration in (from min_city_value_in_max_region) and who_region in (from max_region)\n)\nfrom min_city_in_max_region",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "What is the city with the lowest pm10 concentration in the region with the highest average pm10 concentration?",
+        "category": "hard"
+    },
+    {
+        "db_id": "hn",
+        "query": "SELECT *, regexp_extract(text, '([a-z0-9_\\.-]+)@([\\da-z\\.-]+)\\.([a-z\\.]{2,63})',0) email from hacker_news where email[:4]='test'",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Get all posts on hn that contain an email address starting with test. Return all original columns, plus a new column containing the email address.",
+        "category": "hard"
+    },
+    {
+        "db_id": "json",
+        "query": "SELECT employee.id, employee.first_name FROM employee_json",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Extract id and first_name properties as individual columns from the employee struct",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "who",
+        "query": "SELECT who_region[1]::INT as region, * EXCLUDE (who_region)  FROM who.ambient_air_quality",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "count quality measurements per region. Make sure to return the region code (first char of who_region) as integer and sort by region.",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "flightinfo",
+        "query": "SELECT seat.seat_number FROM seat \nJOIN direct_flight ON direct_flight.flight_number = seat.flight_number \nJOIN airport AS departure_airport ON departure_airport.iata_code = direct_flight.departure_airport_iata_code \nJOIN airport AS arriving_airport ON arriving_airport.iata_code = direct_flight.arriving_airport_iata_code \nJOIN city AS departure_city ON departure_city.city_zipcode = departure_airport.city_zip_code \nJOIN city AS arriving_city ON arriving_city.city_zipcode = arriving_airport.city_zip_code \nWHERE departure_city.city_name = 'Bruxelles' AND arriving_city.city_name = 'Newark';",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "Which seats were available on the flight from Bruxelles to Newark?",
+        "category": "hard"
+    },
+    {
+        "db_id": "laptop",
+        "query": "COPY customers FROM 'customers_12_12_2023.csv';",
+        "setup_sql": "COPY customers TO 'customers_12_12_2023.csv';",
+        "validation_sql": "SELECT * FROM customers;",
+        "question": "copy content of csv file customers_12_12_2023.csv into customers table",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "COPY customers FROM 'customers_12_12_2023.csv' (DELIMITER '\\t');",
+        "setup_sql": "COPY customers TO 'customers_12_12_2023.csv' (FORMAT CSV, DELIMITER '\\t');",
+        "validation_sql": "SELECT * FROM customers;",
+        "question": "copy content of csv file costomers_12_12_2023.csv into customers table with tab separator",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "COPY customers FROM 'customers_partitioned/city=Amsterdam/*.parquet';",
+        "setup_sql": "COPY customers TO 'customers_partitioned' (FORMAT PARQUET, PARTITION_BY (city), OVERWRITE_OR_IGNORE True);",
+        "validation_sql": "SELECT * FROM customers;;",
+        "question": "copy any parquet files from 'customers_partitioned/city=Amsterdam/' into customers table",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "COPY customers(customer_id) FROM 'customers_customer_ids.csv';",
+        "setup_sql": "COPY customers(customer_id) TO 'customers_customer_ids.csv';",
+        "validation_sql": "SELECT * FROM customers;",
+        "question": "copy only the customer_id column from the customers_customer_ids.csv into the customers tables",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "CREATE TABLE test_tbl AS SELECT * FROM read_json_auto('test.json');",
+        "setup_sql": "COPY customers TO 'test.json'\n",
+        "validation_sql": "SELECT * FROM test_tbl;",
+        "question": "read json file from test.json and create new table from it called 'test_tbl'",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT * FROM read_csv_auto('test.csv');",
+        "setup_sql": "COPY customers TO 'test.csv';",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "read csv from test.csv",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT * FROM read_csv_auto('test.csv', columns={'customer_id': 'VARCHAR', 'firstname': 'VARCHAR', 'lastname': 'VARCHAR'});",
+        "setup_sql": "COPY customers(customer_id, firstname, lastname) TO 'test.csv';",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "read csv from test.csv with predefined column and types - customer_id: string, firstname: string, lastname: string",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT * EXCLUDE (ram, hd) FROM pcs;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "select all columns from pcs table except for ram and hd",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT COLUMNS('name$') FROM customers;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "select all columns ending with 'name' from customers table",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT LENGTH(COLUMNS('name$')) FROM customers",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "for each column ending with 'name' in the customers table, compute the string length",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT * REPLACE (upper(city) AS city) FROM customers;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "get all columns from customer table, and make all city names uppercase",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "EXPLAIN SELECT * FROM customers",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "show query plan for query: SELECT * from customers",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT ascii(lastname) FROM customers;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "get the first character of the firstname column and cast it to an INT",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT model, speed::INTEGER FROM laptops;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "get laptop name and speed, return the speed as integer",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop_array",
+        "query": "SELECT phone_numbers[1] FROM customers;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "get the first phone number of each customer",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop_array",
+        "query": "INSERT INTO customers(customer_id, phone_numbers) VALUES (5, ['12312323', '23123344']);",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM customers;",
+        "question": "insert two phone numbers to customer with id 5 [\\\"12312323\\\", and '23123344']",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "ALTER TABLE customers ADD COLUMN phone_numbers VARCHAR[];",
+        "setup_sql": ";",
+        "validation_sql": "DESCRIBE customers;",
+        "question": "how to add a new column phone_numbers to the customers table, with array type varchar",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT firstname[1] FROM customers;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "get the first letter of the customers firstname",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop_array",
+        "query": "SELECT phone_numbers[:2] FROM customers;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "get the first two phone numbers from the phone numbers array of each customer",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT {'a':1, 'b':2, 'c':3};",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "create a struct with keys a, b, c and values 1,2,3",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT [1,2,3];\n",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "create array with values 1,2,3",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "CREATE TABLE test (embeddings FLOAT[100]);",
+        "setup_sql": ";",
+        "validation_sql": "DESCRIBE test;",
+        "question": "create table test with a fix-sized array column with 100 dimenions, called embeddings",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "CREATE TABLE test (person STRUCT(name VARCHAR, id INTEGER));",
+        "setup_sql": ";",
+        "validation_sql": "DESCRIBE test;",
+        "question": "create table test with a struct column called person with properties name and id",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop_struct",
+        "query": "SELECT person.name, person.id FROM test;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "get persons name and persons id from the test table.",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "UPDATE customers SET email = NULL;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT email FROM customers;",
+        "question": "remove all values from email column in customers table",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop_json",
+        "query": "ALTER TABLE customers ALTER COLUMN email SET DATA TYPE VARCHAR;",
+        "setup_sql": ";",
+        "validation_sql": "DESCRIBE customers;",
+        "question": "make customer email of type VARCHAR",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop_json",
+        "query": "INSERT INTO customers (customer_id, email) VALUES (5,'{\"from\": \"[email protected]\", \"to\": \"[email protected]\"}');",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM customers;",
+        "question": "insert json into customer email for customer id 5: {'from': '[email protected]', 'to': '[email protected]'}",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop_json",
+        "query": "SELECT customers.email->>'from' FROM customers;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "get 'from' field from customer email json",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SUMMARIZE customers;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "summarize the customer table",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT * FROM customers USING SAMPLE 10% (reservoir);",
+        "setup_sql": ";",
+        "validation_sql": "SELECT count(*) FROM ddb_benchmark_result;",
+        "question": "sample 10% from the customers table using reservoir sampling",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SET threads = 10;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT current_setting('threads');",
+        "question": "set number of threads to 10",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SET memory_limit = '20G';\n",
+        "setup_sql": ";",
+        "validation_sql": "SELECT current_setting('memory_limit');",
+        "question": "set memory limit to 20 gigabyte",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT * EXCLUDE (price), avg(price) FROM laptops GROUP BY ALL;",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "show the average price of laptop and group by the remaining columns",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "SELECT * FROM laptops WHERE price > 1000 ORDER BY ALL;\n",
+        "setup_sql": ";",
+        "validation_sql": "SELECT * FROM ddb_benchmark_result;",
+        "question": "show all laptops with price above 1000 and order by all columns",
+        "category": "duckdb"
+    },
+    {
+        "db_id": "laptop",
+        "query": "ATTACH 'who.ddb';",
+        "setup_sql": ";",
+        "validation_sql": "SHOW DATABASES;",
+        "question": "attach database file who.ddb",
+        "category": "duckdb"
+    }
+]

duckdb-nsql/eval/data/tables.json

@@ -0,0 +1,3465 @@
+[
+    {
+        "db_id": "hn",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "title"
+            ],
+            [
+                0,
+                "url"
+            ],
+            [
+                0,
+                "text"
+            ],
+            [
+                0,
+                "dead"
+            ],
+            [
+                0,
+                "by"
+            ],
+            [
+                0,
+                "score"
+            ],
+            [
+                0,
+                "time"
+            ],
+            [
+                0,
+                "timestamp"
+            ],
+            [
+                0,
+                "type"
+            ],
+            [
+                0,
+                "id"
+            ],
+            [
+                0,
+                "parent"
+            ],
+            [
+                0,
+                "descendants"
+            ],
+            [
+                0,
+                "ranking"
+            ],
+            [
+                0,
+                "deleted"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "title"
+            ],
+            [
+                0,
+                "url"
+            ],
+            [
+                0,
+                "text"
+            ],
+            [
+                0,
+                "dead"
+            ],
+            [
+                0,
+                "by"
+            ],
+            [
+                0,
+                "score"
+            ],
+            [
+                0,
+                "time"
+            ],
+            [
+                0,
+                "timestamp"
+            ],
+            [
+                0,
+                "type"
+            ],
+            [
+                0,
+                "id"
+            ],
+            [
+                0,
+                "parent"
+            ],
+            [
+                0,
+                "descendants"
+            ],
+            [
+                0,
+                "ranking"
+            ],
+            [
+                0,
+                "deleted"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "varchar",
+            "varchar",
+            "varchar",
+            "boolean",
+            "varchar",
+            "bigint",
+            "bigint",
+            "timestamp",
+            "varchar",
+            "bigint",
+            "bigint",
+            "bigint",
+            "bigint",
+            "boolean"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "hacker_news"
+        ],
+        "table_names_original": [
+            "hacker_news"
+        ]
+    },
+    {
+        "db_id": "laptop",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "customer_id"
+            ],
+            [
+                0,
+                "firstname"
+            ],
+            [
+                0,
+                "lastname"
+            ],
+            [
+                0,
+                "city"
+            ],
+            [
+                0,
+                "address"
+            ],
+            [
+                0,
+                "email"
+            ],
+            [
+                1,
+                "model"
+            ],
+            [
+                1,
+                "speed"
+            ],
+            [
+                1,
+                "ram"
+            ],
+            [
+                1,
+                "hd"
+            ],
+            [
+                1,
+                "screen"
+            ],
+            [
+                1,
+                "price"
+            ],
+            [
+                2,
+                "model"
+            ],
+            [
+                2,
+                "speed"
+            ],
+            [
+                2,
+                "ram"
+            ],
+            [
+                2,
+                "hd"
+            ],
+            [
+                2,
+                "price"
+            ],
+            [
+                3,
+                "model"
+            ],
+            [
+                3,
+                "color"
+            ],
+            [
+                3,
+                "type"
+            ],
+            [
+                3,
+                "price"
+            ],
+            [
+                4,
+                "maker"
+            ],
+            [
+                4,
+                "model"
+            ],
+            [
+                4,
+                "type"
+            ],
+            [
+                5,
+                "customer_id"
+            ],
+            [
+                5,
+                "model"
+            ],
+            [
+                5,
+                "quantity"
+            ],
+            [
+                5,
+                "day"
+            ],
+            [
+                5,
+                "paid"
+            ],
+            [
+                5,
+                "type_of_payment"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "customer_id"
+            ],
+            [
+                0,
+                "firstname"
+            ],
+            [
+                0,
+                "lastname"
+            ],
+            [
+                0,
+                "city"
+            ],
+            [
+                0,
+                "address"
+            ],
+            [
+                0,
+                "email"
+            ],
+            [
+                1,
+                "model"
+            ],
+            [
+                1,
+                "speed"
+            ],
+            [
+                1,
+                "ram"
+            ],
+            [
+                1,
+                "hd"
+            ],
+            [
+                1,
+                "screen"
+            ],
+            [
+                1,
+                "price"
+            ],
+            [
+                2,
+                "model"
+            ],
+            [
+                2,
+                "speed"
+            ],
+            [
+                2,
+                "ram"
+            ],
+            [
+                2,
+                "hd"
+            ],
+            [
+                2,
+                "price"
+            ],
+            [
+                3,
+                "model"
+            ],
+            [
+                3,
+                "color"
+            ],
+            [
+                3,
+                "type"
+            ],
+            [
+                3,
+                "price"
+            ],
+            [
+                4,
+                "maker"
+            ],
+            [
+                4,
+                "model"
+            ],
+            [
+                4,
+                "type"
+            ],
+            [
+                5,
+                "customer_id"
+            ],
+            [
+                5,
+                "model"
+            ],
+            [
+                5,
+                "quantity"
+            ],
+            [
+                5,
+                "day"
+            ],
+            [
+                5,
+                "paid"
+            ],
+            [
+                5,
+                "type_of_payment"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "char",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "char",
+            "double",
+            "int",
+            "int",
+            "double",
+            "double",
+            "char",
+            "double",
+            "int",
+            "int",
+            "double",
+            "char",
+            "varchar",
+            "varchar",
+            "double",
+            "char",
+            "char",
+            "varchar",
+            "char",
+            "char",
+            "int",
+            "date",
+            "double",
+            "varchar"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "customers",
+            "laptops",
+            "pcs",
+            "printers",
+            "products",
+            "sales"
+        ],
+        "table_names_original": [
+            "customers",
+            "laptops",
+            "pcs",
+            "printers",
+            "products",
+            "sales"
+        ]
+    },
+    {
+        "db_id": "transactions",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "id"
+            ],
+            [
+                0,
+                "name"
+            ],
+            [
+                1,
+                "user_id"
+            ],
+            [
+                1,
+                "amount"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "id"
+            ],
+            [
+                0,
+                "name"
+            ],
+            [
+                1,
+                "user_id"
+            ],
+            [
+                1,
+                "amount"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "int",
+            "varchar",
+            "int",
+            "int"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "users",
+            "transactions"
+        ],
+        "table_names_original": [
+            "users",
+            "transactions"
+        ]
+    },
+    {
+        "db_id": "who",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "who_region"
+            ],
+            [
+                0,
+                "iso3"
+            ],
+            [
+                0,
+                "country_name"
+            ],
+            [
+                0,
+                "city"
+            ],
+            [
+                0,
+                "year"
+            ],
+            [
+                0,
+                "version"
+            ],
+            [
+                0,
+                "pm10_concentration"
+            ],
+            [
+                0,
+                "pm25_concentration"
+            ],
+            [
+                0,
+                "no2_concentration"
+            ],
+            [
+                0,
+                "pm10_tempcov"
+            ],
+            [
+                0,
+                "pm25_tempcov"
+            ],
+            [
+                0,
+                "no2_tempcov"
+            ],
+            [
+                0,
+                "type_of_stations"
+            ],
+            [
+                0,
+                "reference"
+            ],
+            [
+                0,
+                "web_link"
+            ],
+            [
+                0,
+                "population"
+            ],
+            [
+                0,
+                "population_source"
+            ],
+            [
+                0,
+                "latitude"
+            ],
+            [
+                0,
+                "longitude"
+            ],
+            [
+                0,
+                "who_ms"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "who_region"
+            ],
+            [
+                0,
+                "iso3"
+            ],
+            [
+                0,
+                "country_name"
+            ],
+            [
+                0,
+                "city"
+            ],
+            [
+                0,
+                "year"
+            ],
+            [
+                0,
+                "version"
+            ],
+            [
+                0,
+                "pm10_concentration"
+            ],
+            [
+                0,
+                "pm25_concentration"
+            ],
+            [
+                0,
+                "no2_concentration"
+            ],
+            [
+                0,
+                "pm10_tempcov"
+            ],
+            [
+                0,
+                "pm25_tempcov"
+            ],
+            [
+                0,
+                "no2_tempcov"
+            ],
+            [
+                0,
+                "type_of_stations"
+            ],
+            [
+                0,
+                "reference"
+            ],
+            [
+                0,
+                "web_link"
+            ],
+            [
+                0,
+                "population"
+            ],
+            [
+                0,
+                "population_source"
+            ],
+            [
+                0,
+                "latitude"
+            ],
+            [
+                0,
+                "longitude"
+            ],
+            [
+                0,
+                "who_ms"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "bigint",
+            "varchar",
+            "bigint",
+            "bigint",
+            "bigint",
+            "bigint",
+            "bigint",
+            "bigint",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "float",
+            "float",
+            "bigint"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "ambient_air_quality"
+        ],
+        "table_names_original": [
+            "ambient_air_quality"
+        ]
+    },
+    {
+        "db_id": "nyc",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "unique_key"
+            ],
+            [
+                0,
+                "created_date"
+            ],
+            [
+                0,
+                "closed_date"
+            ],
+            [
+                0,
+                "agency"
+            ],
+            [
+                0,
+                "agency_name"
+            ],
+            [
+                0,
+                "complaint_type"
+            ],
+            [
+                0,
+                "descriptor"
+            ],
+            [
+                0,
+                "location_type"
+            ],
+            [
+                0,
+                "incident_zip"
+            ],
+            [
+                0,
+                "incident_address"
+            ],
+            [
+                0,
+                "street_name"
+            ],
+            [
+                0,
+                "cross_street_1"
+            ],
+            [
+                0,
+                "cross_street_2"
+            ],
+            [
+                0,
+                "intersection_street_1"
+            ],
+            [
+                0,
+                "intersection_street_2"
+            ],
+            [
+                0,
+                "address_type"
+            ],
+            [
+                0,
+                "city"
+            ],
+            [
+                0,
+                "landmark"
+            ],
+            [
+                0,
+                "facility_type"
+            ],
+            [
+                0,
+                "status"
+            ],
+            [
+                0,
+                "due_date"
+            ],
+            [
+                0,
+                "resolution_description"
+            ],
+            [
+                0,
+                "resolution_action_updated_date"
+            ],
+            [
+                0,
+                "community_board"
+            ],
+            [
+                0,
+                "bbl"
+            ],
+            [
+                0,
+                "borough"
+            ],
+            [
+                0,
+                "x_coordinate_state_plane"
+            ],
+            [
+                0,
+                "y_coordinate_state_plane"
+            ],
+            [
+                0,
+                "open_data_channel_type"
+            ],
+            [
+                0,
+                "park_facility_name"
+            ],
+            [
+                0,
+                "park_borough"
+            ],
+            [
+                0,
+                "vehicle_type"
+            ],
+            [
+                0,
+                "taxi_company_borough"
+            ],
+            [
+                0,
+                "taxi_pick_up_location"
+            ],
+            [
+                0,
+                "bridge_highway_name"
+            ],
+            [
+                0,
+                "bridge_highway_direction"
+            ],
+            [
+                0,
+                "road_ramp"
+            ],
+            [
+                0,
+                "bridge_highway_segment"
+            ],
+            [
+                0,
+                "latitude"
+            ],
+            [
+                0,
+                "longitude"
+            ],
+            [
+                1,
+                "hvfhs_license_num"
+            ],
+            [
+                1,
+                "dispatching_base_num"
+            ],
+            [
+                1,
+                "originating_base_num"
+            ],
+            [
+                1,
+                "request_datetime"
+            ],
+            [
+                1,
+                "on_scene_datetime"
+            ],
+            [
+                1,
+                "pickup_datetime"
+            ],
+            [
+                1,
+                "dropoff_datetime"
+            ],
+            [
+                1,
+                "PULocationID"
+            ],
+            [
+                1,
+                "DOLocationID"
+            ],
+            [
+                1,
+                "trip_miles"
+            ],
+            [
+                1,
+                "trip_time"
+            ],
+            [
+                1,
+                "base_passenger_fare"
+            ],
+            [
+                1,
+                "tolls"
+            ],
+            [
+                1,
+                "bcf"
+            ],
+            [
+                1,
+                "sales_tax"
+            ],
+            [
+                1,
+                "congestion_surcharge"
+            ],
+            [
+                1,
+                "airport_fee"
+            ],
+            [
+                1,
+                "tips"
+            ],
+            [
+                1,
+                "driver_pay"
+            ],
+            [
+                1,
+                "shared_request_flag"
+            ],
+            [
+                1,
+                "shared_match_flag"
+            ],
+            [
+                1,
+                "access_a_ride_flag"
+            ],
+            [
+                1,
+                "wav_request_flag"
+            ],
+            [
+                1,
+                "wav_match_flag"
+            ],
+            [
+                2,
+                "VendorID"
+            ],
+            [
+                2,
+                "tpep_pickup_datetime"
+            ],
+            [
+                2,
+                "tpep_dropoff_datetime"
+            ],
+            [
+                2,
+                "passenger_count"
+            ],
+            [
+                2,
+                "trip_distance"
+            ],
+            [
+                2,
+                "RatecodeID"
+            ],
+            [
+                2,
+                "store_and_fwd_flag"
+            ],
+            [
+                2,
+                "PULocationID"
+            ],
+            [
+                2,
+                "DOLocationID"
+            ],
+            [
+                2,
+                "payment_type"
+            ],
+            [
+                2,
+                "fare_amount"
+            ],
+            [
+                2,
+                "extra"
+            ],
+            [
+                2,
+                "mta_tax"
+            ],
+            [
+                2,
+                "tip_amount"
+            ],
+            [
+                2,
+                "tolls_amount"
+            ],
+            [
+                2,
+                "improvement_surcharge"
+            ],
+            [
+                2,
+                "total_amount"
+            ],
+            [
+                2,
+                "congestion_surcharge"
+            ],
+            [
+                2,
+                "airport_fee"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "unique_key"
+            ],
+            [
+                0,
+                "created_date"
+            ],
+            [
+                0,
+                "closed_date"
+            ],
+            [
+                0,
+                "agency"
+            ],
+            [
+                0,
+                "agency_name"
+            ],
+            [
+                0,
+                "complaint_type"
+            ],
+            [
+                0,
+                "descriptor"
+            ],
+            [
+                0,
+                "location_type"
+            ],
+            [
+                0,
+                "incident_zip"
+            ],
+            [
+                0,
+                "incident_address"
+            ],
+            [
+                0,
+                "street_name"
+            ],
+            [
+                0,
+                "cross_street_1"
+            ],
+            [
+                0,
+                "cross_street_2"
+            ],
+            [
+                0,
+                "intersection_street_1"
+            ],
+            [
+                0,
+                "intersection_street_2"
+            ],
+            [
+                0,
+                "address_type"
+            ],
+            [
+                0,
+                "city"
+            ],
+            [
+                0,
+                "landmark"
+            ],
+            [
+                0,
+                "facility_type"
+            ],
+            [
+                0,
+                "status"
+            ],
+            [
+                0,
+                "due_date"
+            ],
+            [
+                0,
+                "resolution_description"
+            ],
+            [
+                0,
+                "resolution_action_updated_date"
+            ],
+            [
+                0,
+                "community_board"
+            ],
+            [
+                0,
+                "bbl"
+            ],
+            [
+                0,
+                "borough"
+            ],
+            [
+                0,
+                "x_coordinate_state_plane"
+            ],
+            [
+                0,
+                "y_coordinate_state_plane"
+            ],
+            [
+                0,
+                "open_data_channel_type"
+            ],
+            [
+                0,
+                "park_facility_name"
+            ],
+            [
+                0,
+                "park_borough"
+            ],
+            [
+                0,
+                "vehicle_type"
+            ],
+            [
+                0,
+                "taxi_company_borough"
+            ],
+            [
+                0,
+                "taxi_pick_up_location"
+            ],
+            [
+                0,
+                "bridge_highway_name"
+            ],
+            [
+                0,
+                "bridge_highway_direction"
+            ],
+            [
+                0,
+                "road_ramp"
+            ],
+            [
+                0,
+                "bridge_highway_segment"
+            ],
+            [
+                0,
+                "latitude"
+            ],
+            [
+                0,
+                "longitude"
+            ],
+            [
+                1,
+                "hvfhs_license_num"
+            ],
+            [
+                1,
+                "dispatching_base_num"
+            ],
+            [
+                1,
+                "originating_base_num"
+            ],
+            [
+                1,
+                "request_datetime"
+            ],
+            [
+                1,
+                "on_scene_datetime"
+            ],
+            [
+                1,
+                "pickup_datetime"
+            ],
+            [
+                1,
+                "dropoff_datetime"
+            ],
+            [
+                1,
+                "PULocationID"
+            ],
+            [
+                1,
+                "DOLocationID"
+            ],
+            [
+                1,
+                "trip_miles"
+            ],
+            [
+                1,
+                "trip_time"
+            ],
+            [
+                1,
+                "base_passenger_fare"
+            ],
+            [
+                1,
+                "tolls"
+            ],
+            [
+                1,
+                "bcf"
+            ],
+            [
+                1,
+                "sales_tax"
+            ],
+            [
+                1,
+                "congestion_surcharge"
+            ],
+            [
+                1,
+                "airport_fee"
+            ],
+            [
+                1,
+                "tips"
+            ],
+            [
+                1,
+                "driver_pay"
+            ],
+            [
+                1,
+                "shared_request_flag"
+            ],
+            [
+                1,
+                "shared_match_flag"
+            ],
+            [
+                1,
+                "access_a_ride_flag"
+            ],
+            [
+                1,
+                "wav_request_flag"
+            ],
+            [
+                1,
+                "wav_match_flag"
+            ],
+            [
+                2,
+                "VendorID"
+            ],
+            [
+                2,
+                "tpep_pickup_datetime"
+            ],
+            [
+                2,
+                "tpep_dropoff_datetime"
+            ],
+            [
+                2,
+                "passenger_count"
+            ],
+            [
+                2,
+                "trip_distance"
+            ],
+            [
+                2,
+                "RatecodeID"
+            ],
+            [
+                2,
+                "store_and_fwd_flag"
+            ],
+            [
+                2,
+                "PULocationID"
+            ],
+            [
+                2,
+                "DOLocationID"
+            ],
+            [
+                2,
+                "payment_type"
+            ],
+            [
+                2,
+                "fare_amount"
+            ],
+            [
+                2,
+                "extra"
+            ],
+            [
+                2,
+                "mta_tax"
+            ],
+            [
+                2,
+                "tip_amount"
+            ],
+            [
+                2,
+                "tolls_amount"
+            ],
+            [
+                2,
+                "improvement_surcharge"
+            ],
+            [
+                2,
+                "total_amount"
+            ],
+            [
+                2,
+                "congestion_surcharge"
+            ],
+            [
+                2,
+                "airport_fee"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "bigint",
+            "timestamp",
+            "timestamp",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "timestamp",
+            "varchar",
+            "timestamp",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "double",
+            "double",
+            "varchar",
+            "varchar",
+            "varchar",
+            "timestamp",
+            "timestamp",
+            "timestamp",
+            "timestamp",
+            "bigint",
+            "bigint",
+            "double",
+            "bigint",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "bigint",
+            "timestamp",
+            "timestamp",
+            "double",
+            "double",
+            "double",
+            "varchar",
+            "bigint",
+            "bigint",
+            "bigint",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "service_requests",
+            "rideshare",
+            "taxi"
+        ],
+        "table_names_original": [
+            "service_requests",
+            "rideshare",
+            "taxi"
+        ]
+    },
+    {
+        "db_id": "product",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "maker"
+            ],
+            [
+                0,
+                "model"
+            ],
+            [
+                0,
+                "type"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "maker"
+            ],
+            [
+                0,
+                "model"
+            ],
+            [
+                0,
+                "type"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "varchar",
+            "varchar",
+            "varchar"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "products"
+        ],
+        "table_names_original": [
+            "products"
+        ]
+    },
+    {
+        "db_id": "json",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "employee"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "employee"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "struct(id int, first_name text, last_name text, email text)"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "employee_json"
+        ],
+        "table_names_original": [
+            "employee_json"
+        ]
+    },
+    {
+        "db_id": "flightinfo",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "icao_code"
+            ],
+            [
+                0,
+                "email"
+            ],
+            [
+                0,
+                "name"
+            ],
+            [
+                0,
+                "phone_number"
+            ],
+            [
+                0,
+                "iata_code"
+            ],
+            [
+                1,
+                "title"
+            ],
+            [
+                1,
+                "description"
+            ],
+            [
+                1,
+                "price"
+            ],
+            [
+                1,
+                "service_type"
+            ],
+            [
+                1,
+                "airline_icao_code"
+            ],
+            [
+                2,
+                "iata_code"
+            ],
+            [
+                2,
+                "address"
+            ],
+            [
+                2,
+                "name"
+            ],
+            [
+                2,
+                "phone_number"
+            ],
+            [
+                2,
+                "email"
+            ],
+            [
+                2,
+                "city_zip_code"
+            ],
+            [
+                2,
+                "city_dbpedia"
+            ],
+            [
+                3,
+                "title"
+            ],
+            [
+                3,
+                "cabin_bag_dimension_cm"
+            ],
+            [
+                3,
+                "cabin_bags_no"
+            ],
+            [
+                3,
+                "cabin_bg_weight_kg"
+            ],
+            [
+                3,
+                "checked_bag_dimension_cm"
+            ],
+            [
+                3,
+                "checked_bags_no"
+            ],
+            [
+                3,
+                "checked_bag_weight_kg"
+            ],
+            [
+                3,
+                "excessive_price_perkg"
+            ],
+            [
+                3,
+                "flight_type"
+            ],
+            [
+                3,
+                "airline_icao_code"
+            ],
+            [
+                4,
+                "title"
+            ],
+            [
+                4,
+                "description"
+            ],
+            [
+                4,
+                "airline_icao_code"
+            ],
+            [
+                5,
+                "title"
+            ],
+            [
+                5,
+                "description"
+            ],
+            [
+                5,
+                "due_date"
+            ],
+            [
+                5,
+                "refund_postdue_percentage"
+            ],
+            [
+                5,
+                "refund_predue_percentage"
+            ],
+            [
+                5,
+                "airline_icao_code"
+            ],
+            [
+                6,
+                "city_zipcode"
+            ],
+            [
+                6,
+                "city_name"
+            ],
+            [
+                6,
+                "country_iso_code"
+            ],
+            [
+                7,
+                "country_iso_code"
+            ],
+            [
+                7,
+                "country_name"
+            ],
+            [
+                8,
+                "flight_number"
+            ],
+            [
+                8,
+                "departure_airport_iata_code"
+            ],
+            [
+                8,
+                "arriving_airport_iata_code"
+            ],
+            [
+                9,
+                "number"
+            ],
+            [
+                9,
+                "departure_date"
+            ],
+            [
+                9,
+                "arrival_date"
+            ],
+            [
+                9,
+                "distance_km"
+            ],
+            [
+                9,
+                "is_available"
+            ],
+            [
+                9,
+                "duration_min"
+            ],
+            [
+                9,
+                "airline_icao_code"
+            ],
+            [
+                9,
+                "type"
+            ],
+            [
+                10,
+                "title"
+            ],
+            [
+                10,
+                "description"
+            ],
+            [
+                10,
+                "cabin_class_title"
+            ],
+            [
+                10,
+                "baggage_policy_title"
+            ],
+            [
+                10,
+                "cancelation_policy_title"
+            ],
+            [
+                11,
+                "subflight_number"
+            ],
+            [
+                11,
+                "flight_number"
+            ],
+            [
+                12,
+                "flight_package_title"
+            ],
+            [
+                12,
+                "airline_service_title"
+            ],
+            [
+                13,
+                "seat_number"
+            ],
+            [
+                13,
+                "is_available"
+            ],
+            [
+                13,
+                "flight_number"
+            ],
+            [
+                14,
+                "meal_type"
+            ],
+            [
+                14,
+                "airline_service_title"
+            ],
+            [
+                15,
+                "duration_min"
+            ],
+            [
+                15,
+                "duration_from"
+            ],
+            [
+                15,
+                "duration_to"
+            ],
+            [
+                15,
+                "airport_iatacode"
+            ],
+            [
+                15,
+                "flight_number"
+            ],
+            [
+                16,
+                "flight_number"
+            ],
+            [
+                16,
+                "package_title"
+            ],
+            [
+                16,
+                "trip_id"
+            ],
+            [
+                16,
+                "requested_excessive_baggage_kg"
+            ],
+            [
+                16,
+                "seat_number"
+            ],
+            [
+                16,
+                "chosen_meal_service_price"
+            ],
+            [
+                16,
+                "chosen_wifi_service_price"
+            ],
+            [
+                16,
+                "price"
+            ],
+            [
+                17,
+                "id"
+            ],
+            [
+                17,
+                "tax"
+            ],
+            [
+                17,
+                "booking_date"
+            ],
+            [
+                17,
+                "user_email"
+            ],
+            [
+                17,
+                "type"
+            ],
+            [
+                18,
+                "email"
+            ],
+            [
+                18,
+                "first_name"
+            ],
+            [
+                18,
+                "last_name"
+            ],
+            [
+                18,
+                "birthdate"
+            ],
+            [
+                18,
+                "passport_number"
+            ],
+            [
+                18,
+                "address"
+            ],
+            [
+                18,
+                "password"
+            ],
+            [
+                18,
+                "phone_number"
+            ],
+            [
+                19,
+                "wifi_onboard_service_bandwidth_MB"
+            ],
+            [
+                19,
+                "airline_service_title"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "icao_code"
+            ],
+            [
+                0,
+                "email"
+            ],
+            [
+                0,
+                "name"
+            ],
+            [
+                0,
+                "phone_number"
+            ],
+            [
+                0,
+                "iata_code"
+            ],
+            [
+                1,
+                "title"
+            ],
+            [
+                1,
+                "description"
+            ],
+            [
+                1,
+                "price"
+            ],
+            [
+                1,
+                "service_type"
+            ],
+            [
+                1,
+                "airline_icao_code"
+            ],
+            [
+                2,
+                "iata_code"
+            ],
+            [
+                2,
+                "address"
+            ],
+            [
+                2,
+                "name"
+            ],
+            [
+                2,
+                "phone_number"
+            ],
+            [
+                2,
+                "email"
+            ],
+            [
+                2,
+                "city_zip_code"
+            ],
+            [
+                2,
+                "city_dbpedia"
+            ],
+            [
+                3,
+                "title"
+            ],
+            [
+                3,
+                "cabin_bag_dimension_cm"
+            ],
+            [
+                3,
+                "cabin_bags_no"
+            ],
+            [
+                3,
+                "cabin_bg_weight_kg"
+            ],
+            [
+                3,
+                "checked_bag_dimension_cm"
+            ],
+            [
+                3,
+                "checked_bags_no"
+            ],
+            [
+                3,
+                "checked_bag_weight_kg"
+            ],
+            [
+                3,
+                "excessive_price_perkg"
+            ],
+            [
+                3,
+                "flight_type"
+            ],
+            [
+                3,
+                "airline_icao_code"
+            ],
+            [
+                4,
+                "title"
+            ],
+            [
+                4,
+                "description"
+            ],
+            [
+                4,
+                "airline_icao_code"
+            ],
+            [
+                5,
+                "title"
+            ],
+            [
+                5,
+                "description"
+            ],
+            [
+                5,
+                "due_date"
+            ],
+            [
+                5,
+                "refund_postdue_percentage"
+            ],
+            [
+                5,
+                "refund_predue_percentage"
+            ],
+            [
+                5,
+                "airline_icao_code"
+            ],
+            [
+                6,
+                "city_zipcode"
+            ],
+            [
+                6,
+                "city_name"
+            ],
+            [
+                6,
+                "country_iso_code"
+            ],
+            [
+                7,
+                "country_iso_code"
+            ],
+            [
+                7,
+                "country_name"
+            ],
+            [
+                8,
+                "flight_number"
+            ],
+            [
+                8,
+                "departure_airport_iata_code"
+            ],
+            [
+                8,
+                "arriving_airport_iata_code"
+            ],
+            [
+                9,
+                "number"
+            ],
+            [
+                9,
+                "departure_date"
+            ],
+            [
+                9,
+                "arrival_date"
+            ],
+            [
+                9,
+                "distance_km"
+            ],
+            [
+                9,
+                "is_available"
+            ],
+            [
+                9,
+                "duration_min"
+            ],
+            [
+                9,
+                "airline_icao_code"
+            ],
+            [
+                9,
+                "type"
+            ],
+            [
+                10,
+                "title"
+            ],
+            [
+                10,
+                "description"
+            ],
+            [
+                10,
+                "cabin_class_title"
+            ],
+            [
+                10,
+                "baggage_policy_title"
+            ],
+            [
+                10,
+                "cancelation_policy_title"
+            ],
+            [
+                11,
+                "subflight_number"
+            ],
+            [
+                11,
+                "flight_number"
+            ],
+            [
+                12,
+                "flight_package_title"
+            ],
+            [
+                12,
+                "airline_service_title"
+            ],
+            [
+                13,
+                "seat_number"
+            ],
+            [
+                13,
+                "is_available"
+            ],
+            [
+                13,
+                "flight_number"
+            ],
+            [
+                14,
+                "meal_type"
+            ],
+            [
+                14,
+                "airline_service_title"
+            ],
+            [
+                15,
+                "duration_min"
+            ],
+            [
+                15,
+                "duration_from"
+            ],
+            [
+                15,
+                "duration_to"
+            ],
+            [
+                15,
+                "airport_iatacode"
+            ],
+            [
+                15,
+                "flight_number"
+            ],
+            [
+                16,
+                "flight_number"
+            ],
+            [
+                16,
+                "package_title"
+            ],
+            [
+                16,
+                "trip_id"
+            ],
+            [
+                16,
+                "requested_excessive_baggage_kg"
+            ],
+            [
+                16,
+                "seat_number"
+            ],
+            [
+                16,
+                "chosen_meal_service_price"
+            ],
+            [
+                16,
+                "chosen_wifi_service_price"
+            ],
+            [
+                16,
+                "price"
+            ],
+            [
+                17,
+                "id"
+            ],
+            [
+                17,
+                "tax"
+            ],
+            [
+                17,
+                "booking_date"
+            ],
+            [
+                17,
+                "user_email"
+            ],
+            [
+                17,
+                "type"
+            ],
+            [
+                18,
+                "email"
+            ],
+            [
+                18,
+                "first_name"
+            ],
+            [
+                18,
+                "last_name"
+            ],
+            [
+                18,
+                "birthdate"
+            ],
+            [
+                18,
+                "passport_number"
+            ],
+            [
+                18,
+                "address"
+            ],
+            [
+                18,
+                "password"
+            ],
+            [
+                18,
+                "phone_number"
+            ],
+            [
+                19,
+                "wifi_onboard_service_bandwidth_MB"
+            ],
+            [
+                19,
+                "airline_service_title"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "double",
+            "varchar",
+            "varchar",
+            "varchar",
+            "text",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "double",
+            "varchar",
+            "varchar",
+            "varchar",
+            "text",
+            "varchar",
+            "varchar",
+            "text",
+            "text",
+            "int",
+            "int",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "text",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "datetime",
+            "datetime",
+            "double",
+            "tinyint",
+            "double",
+            "varchar",
+            "varchar",
+            "varchar",
+            "text",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "tinyint",
+            "varchar",
+            "varchar",
+            "varchar",
+            "double",
+            "datetime",
+            "datetime",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "int",
+            "int",
+            "varchar",
+            "int",
+            "int",
+            "double",
+            "int",
+            "double",
+            "datetime",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "date",
+            "varchar",
+            "varchar",
+            "varchar",
+            "double",
+            "double",
+            "varchar"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "airline",
+            "airline_service",
+            "airport",
+            "baggage_policy",
+            "cabin_class",
+            "cancellation_policy",
+            "city",
+            "country",
+            "direct_flight",
+            "flight",
+            "flight_package",
+            "non_direct_flight",
+            "package_service",
+            "seat",
+            "special_meal_type",
+            "stopping",
+            "ticke",
+            "trip",
+            "user",
+            "wifi_onboard_service"
+        ],
+        "table_names_original": [
+            "airline",
+            "airline_service",
+            "airport",
+            "baggage_policy",
+            "cabin_class",
+            "cancellation_policy",
+            "city",
+            "country",
+            "direct_flight",
+            "flight",
+            "flight_package",
+            "non_direct_flight",
+            "package_service",
+            "seat",
+            "special_meal_type",
+            "stopping",
+            "ticke",
+            "trip",
+            "user",
+            "wifi_onboard_service"
+        ]
+    },
+    {
+        "db_id": "none",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ]
+        ],
+        "column_types": [
+            "text"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [],
+        "table_names_original": []
+    },
+    {
+        "db_id": "laptop_array",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "customer_id"
+            ],
+            [
+                0,
+                "firstname"
+            ],
+            [
+                0,
+                "lastname"
+            ],
+            [
+                0,
+                "city"
+            ],
+            [
+                0,
+                "address"
+            ],
+            [
+                0,
+                "email"
+            ],
+            [
+                0,
+                "phone_numbers"
+            ],
+            [
+                1,
+                "model"
+            ],
+            [
+                1,
+                "speed"
+            ],
+            [
+                1,
+                "ram"
+            ],
+            [
+                1,
+                "hd"
+            ],
+            [
+                1,
+                "screen"
+            ],
+            [
+                1,
+                "price"
+            ],
+            [
+                2,
+                "model"
+            ],
+            [
+                2,
+                "speed"
+            ],
+            [
+                2,
+                "ram"
+            ],
+            [
+                2,
+                "hd"
+            ],
+            [
+                2,
+                "price"
+            ],
+            [
+                3,
+                "model"
+            ],
+            [
+                3,
+                "color"
+            ],
+            [
+                3,
+                "type"
+            ],
+            [
+                3,
+                "price"
+            ],
+            [
+                4,
+                "maker"
+            ],
+            [
+                4,
+                "model"
+            ],
+            [
+                4,
+                "type"
+            ],
+            [
+                5,
+                "customer_id"
+            ],
+            [
+                5,
+                "model"
+            ],
+            [
+                5,
+                "quantity"
+            ],
+            [
+                5,
+                "day"
+            ],
+            [
+                5,
+                "paid"
+            ],
+            [
+                5,
+                "type_of_payment"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "customer_id"
+            ],
+            [
+                0,
+                "firstname"
+            ],
+            [
+                0,
+                "lastname"
+            ],
+            [
+                0,
+                "city"
+            ],
+            [
+                0,
+                "address"
+            ],
+            [
+                0,
+                "email"
+            ],
+            [
+                0,
+                "phone_number"
+            ],
+            [
+                1,
+                "model"
+            ],
+            [
+                1,
+                "speed"
+            ],
+            [
+                1,
+                "ram"
+            ],
+            [
+                1,
+                "hd"
+            ],
+            [
+                1,
+                "screen"
+            ],
+            [
+                1,
+                "price"
+            ],
+            [
+                2,
+                "model"
+            ],
+            [
+                2,
+                "speed"
+            ],
+            [
+                2,
+                "ram"
+            ],
+            [
+                2,
+                "hd"
+            ],
+            [
+                2,
+                "price"
+            ],
+            [
+                3,
+                "model"
+            ],
+            [
+                3,
+                "color"
+            ],
+            [
+                3,
+                "type"
+            ],
+            [
+                3,
+                "price"
+            ],
+            [
+                4,
+                "maker"
+            ],
+            [
+                4,
+                "model"
+            ],
+            [
+                4,
+                "type"
+            ],
+            [
+                5,
+                "customer_id"
+            ],
+            [
+                5,
+                "model"
+            ],
+            [
+                5,
+                "quantity"
+            ],
+            [
+                5,
+                "day"
+            ],
+            [
+                5,
+                "paid"
+            ],
+            [
+                5,
+                "type_of_payment"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "char",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "array",
+            "char",
+            "double",
+            "int",
+            "int",
+            "double",
+            "double",
+            "char",
+            "double",
+            "int",
+            "int",
+            "double",
+            "char",
+            "varchar",
+            "varchar",
+            "double",
+            "char",
+            "char",
+            "varchar",
+            "char",
+            "char",
+            "int",
+            "date",
+            "double",
+            "varchar"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "customers",
+            "laptops",
+            "pcs",
+            "printers",
+            "products",
+            "sales"
+        ],
+        "table_names_original": [
+            "customers",
+            "laptops",
+            "pcs",
+            "printers",
+            "products",
+            "sales"
+        ]
+    },
+    {
+        "db_id": "laptop_struct",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "person"
+            ],
+            [
+                1,
+                "customer_id"
+            ],
+            [
+                1,
+                "firstname"
+            ],
+            [
+                1,
+                "lastname"
+            ],
+            [
+                1,
+                "city"
+            ],
+            [
+                1,
+                "address"
+            ],
+            [
+                1,
+                "email"
+            ],
+            [
+                2,
+                "model"
+            ],
+            [
+                2,
+                "speed"
+            ],
+            [
+                2,
+                "ram"
+            ],
+            [
+                2,
+                "hd"
+            ],
+            [
+                2,
+                "screen"
+            ],
+            [
+                2,
+                "price"
+            ],
+            [
+                3,
+                "model"
+            ],
+            [
+                3,
+                "speed"
+            ],
+            [
+                3,
+                "ram"
+            ],
+            [
+                3,
+                "hd"
+            ],
+            [
+                3,
+                "price"
+            ],
+            [
+                4,
+                "model"
+            ],
+            [
+                4,
+                "color"
+            ],
+            [
+                4,
+                "type"
+            ],
+            [
+                4,
+                "price"
+            ],
+            [
+                5,
+                "maker"
+            ],
+            [
+                5,
+                "model"
+            ],
+            [
+                5,
+                "type"
+            ],
+            [
+                6,
+                "customer_id"
+            ],
+            [
+                6,
+                "model"
+            ],
+            [
+                6,
+                "quantity"
+            ],
+            [
+                6,
+                "day"
+            ],
+            [
+                6,
+                "paid"
+            ],
+            [
+                6,
+                "type_of_payment"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "person"
+            ],
+            [
+                1,
+                "customer_id"
+            ],
+            [
+                1,
+                "firstname"
+            ],
+            [
+                1,
+                "lastname"
+            ],
+            [
+                1,
+                "city"
+            ],
+            [
+                1,
+                "address"
+            ],
+            [
+                1,
+                "email"
+            ],
+            [
+                2,
+                "model"
+            ],
+            [
+                2,
+                "speed"
+            ],
+            [
+                2,
+                "ram"
+            ],
+            [
+                2,
+                "hd"
+            ],
+            [
+                2,
+                "screen"
+            ],
+            [
+                2,
+                "price"
+            ],
+            [
+                3,
+                "model"
+            ],
+            [
+                3,
+                "speed"
+            ],
+            [
+                3,
+                "ram"
+            ],
+            [
+                3,
+                "hd"
+            ],
+            [
+                3,
+                "price"
+            ],
+            [
+                4,
+                "model"
+            ],
+            [
+                4,
+                "color"
+            ],
+            [
+                4,
+                "type"
+            ],
+            [
+                4,
+                "price"
+            ],
+            [
+                5,
+                "maker"
+            ],
+            [
+                5,
+                "model"
+            ],
+            [
+                5,
+                "type"
+            ],
+            [
+                6,
+                "customer_id"
+            ],
+            [
+                6,
+                "model"
+            ],
+            [
+                6,
+                "quantity"
+            ],
+            [
+                6,
+                "day"
+            ],
+            [
+                6,
+                "paid"
+            ],
+            [
+                6,
+                "type_of_payment"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "struct(id int, name: text)",
+            "char",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "char",
+            "double",
+            "int",
+            "int",
+            "double",
+            "double",
+            "char",
+            "double",
+            "int",
+            "int",
+            "double",
+            "char",
+            "varchar",
+            "varchar",
+            "double",
+            "char",
+            "char",
+            "varchar",
+            "char",
+            "char",
+            "int",
+            "date",
+            "double",
+            "varchar"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "test",
+            "customers",
+            "laptops",
+            "pcs",
+            "printers",
+            "products",
+            "sales"
+        ],
+        "table_names_original": [
+            "test",
+            "customers",
+            "laptops",
+            "pcs",
+            "printers",
+            "products",
+            "sales"
+        ]
+    },
+    {
+        "db_id": "laptop_json",
+        "column_names": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "customer_id"
+            ],
+            [
+                0,
+                "firstname"
+            ],
+            [
+                0,
+                "lastname"
+            ],
+            [
+                0,
+                "city"
+            ],
+            [
+                0,
+                "address"
+            ],
+            [
+                0,
+                "email"
+            ],
+            [
+                1,
+                "model"
+            ],
+            [
+                1,
+                "speed"
+            ],
+            [
+                1,
+                "ram"
+            ],
+            [
+                1,
+                "hd"
+            ],
+            [
+                1,
+                "screen"
+            ],
+            [
+                1,
+                "price"
+            ],
+            [
+                2,
+                "model"
+            ],
+            [
+                2,
+                "speed"
+            ],
+            [
+                2,
+                "ram"
+            ],
+            [
+                2,
+                "hd"
+            ],
+            [
+                2,
+                "price"
+            ],
+            [
+                3,
+                "model"
+            ],
+            [
+                3,
+                "color"
+            ],
+            [
+                3,
+                "type"
+            ],
+            [
+                3,
+                "price"
+            ],
+            [
+                4,
+                "maker"
+            ],
+            [
+                4,
+                "model"
+            ],
+            [
+                4,
+                "type"
+            ],
+            [
+                5,
+                "customer_id"
+            ],
+            [
+                5,
+                "model"
+            ],
+            [
+                5,
+                "quantity"
+            ],
+            [
+                5,
+                "day"
+            ],
+            [
+                5,
+                "paid"
+            ],
+            [
+                5,
+                "type_of_payment"
+            ]
+        ],
+        "column_names_original": [
+            [
+                -1,
+                "*"
+            ],
+            [
+                0,
+                "customer_id"
+            ],
+            [
+                0,
+                "firstname"
+            ],
+            [
+                0,
+                "lastname"
+            ],
+            [
+                0,
+                "city"
+            ],
+            [
+                0,
+                "address"
+            ],
+            [
+                0,
+                "email"
+            ],
+            [
+                1,
+                "model"
+            ],
+            [
+                1,
+                "speed"
+            ],
+            [
+                1,
+                "ram"
+            ],
+            [
+                1,
+                "hd"
+            ],
+            [
+                1,
+                "screen"
+            ],
+            [
+                1,
+                "price"
+            ],
+            [
+                2,
+                "model"
+            ],
+            [
+                2,
+                "speed"
+            ],
+            [
+                2,
+                "ram"
+            ],
+            [
+                2,
+                "hd"
+            ],
+            [
+                2,
+                "price"
+            ],
+            [
+                3,
+                "model"
+            ],
+            [
+                3,
+                "color"
+            ],
+            [
+                3,
+                "type"
+            ],
+            [
+                3,
+                "price"
+            ],
+            [
+                4,
+                "maker"
+            ],
+            [
+                4,
+                "model"
+            ],
+            [
+                4,
+                "type"
+            ],
+            [
+                5,
+                "customer_id"
+            ],
+            [
+                5,
+                "model"
+            ],
+            [
+                5,
+                "quantity"
+            ],
+            [
+                5,
+                "day"
+            ],
+            [
+                5,
+                "paid"
+            ],
+            [
+                5,
+                "type_of_payment"
+            ]
+        ],
+        "column_types": [
+            "text",
+            "char",
+            "varchar",
+            "varchar",
+            "varchar",
+            "varchar",
+            "json",
+            "char",
+            "double",
+            "int",
+            "int",
+            "double",
+            "double",
+            "char",
+            "double",
+            "int",
+            "int",
+            "double",
+            "char",
+            "varchar",
+            "varchar",
+            "double",
+            "char",
+            "char",
+            "varchar",
+            "char",
+            "char",
+            "int",
+            "date",
+            "double",
+            "varchar"
+        ],
+        "foreign_keys": {},
+        "primary_keys": {},
+        "table_names": [
+            "customers",
+            "laptops",
+            "pcs",
+            "printers",
+            "products",
+            "sales"
+        ],
+        "table_names_original": [
+            "customers",
+            "laptops",
+            "pcs",
+            "printers",
+            "products",
+            "sales"
+        ]
+    }
+]

duckdb-nsql/eval/data_utils.py

@@ -0,0 +1,79 @@
+"""Training data prep utils."""
+import json
+import re
+from collections import defaultdict
+from schema import ForeignKey, Table, TableColumn
+
+
+def read_tables_json(
+    schema_file: str,
+    lowercase: bool = False,
+) -> dict[str, dict[str, Table]]:
+    """Read tables json."""
+    data = json.load(open(schema_file))
+    db_to_tables = {}
+    for db in data:
+        db_name = db["db_id"]
+        table_names = db["table_names_original"]
+        db["column_names_original"] = [
+            [x[0], x[1]] for x in db["column_names_original"]
+        ]
+        db["column_types"] = db["column_types"]
+        if lowercase:
+            table_names = [tn.lower() for tn in table_names]
+        pks = db["primary_keys"]
+        fks = db["foreign_keys"]
+        tables = defaultdict(list)
+        tables_pks = defaultdict(list)
+        tables_fks = defaultdict(list)
+        for idx, ((ti, col_name), col_type) in enumerate(
+            zip(db["column_names_original"], db["column_types"])
+        ):
+            if ti == -1:
+                continue
+            if lowercase:
+                col_name = col_name.lower()
+                col_type = col_type.lower()
+            if idx in pks:
+                tables_pks[table_names[ti]].append(
+                    TableColumn(name=col_name, dtype=col_type)
+                )
+            for fk in fks:
+                if idx == fk[0]:
+                    other_column = db["column_names_original"][fk[1]]
+                    other_column_type = db["column_types"][fk[1]]
+                    other_table = table_names[other_column[0]]
+                    tables_fks[table_names[ti]].append(
+                        ForeignKey(
+                            column=TableColumn(name=col_name, dtype=col_type),
+                            references_name=other_table,
+                            references_column=TableColumn(
+                                name=other_column[1], dtype=other_column_type
+                            ),
+                        )
+                    )
+            tables[table_names[ti]].append(TableColumn(name=col_name, dtype=col_type))
+        db_to_tables[db_name] = {
+            table_name: Table(
+                name=table_name,
+                columns=tables[table_name],
+                pks=tables_pks[table_name],
+                fks=tables_fks[table_name],
+                examples=None,
+            )
+            for table_name in tables
+        }
+    return db_to_tables
+
+
+def clean_str(target: str) -> str:
+    """Clean string for question."""
+    if not target:
+        return target
+
+    target = re.sub(r"[^\x00-\x7f]", r" ", target)
+    line = re.sub(r"''", r" ", target)
+    line = re.sub(r"``", r" ", line)
+    line = re.sub(r"\"", r"'", line)
+    line = re.sub(r"[\t ]+", " ", line)
+    return line.strip()

duckdb-nsql/eval/doc_retriever.py

@@ -0,0 +1,92 @@
+"""Retrieve documentation for a given query."""
+
+from pathlib import Path
+from typing import Any
+from rich.console import Console
+from tqdm import tqdm
+import numpy as np
+from manifest import Manifest
+from langchain.text_splitter import MarkdownHeaderTextSplitter
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+console = Console(soft_wrap=True)
+
+try:
+    EMBEDDING_MODEL = Manifest(
+        client_name="openaiembedding",
+    )
+except Exception as e:
+    console.print(e)
+    console.print(
+        "Failed to load embedding model. Likely OPENAI API key is not set. Please set to run document retrieval.",
+        style="bold red",
+    )
+
+
+def load_documentation(path: Path) -> dict[str, str]:
+    """Load documentation from path."""
+    content = {}
+    for file in path.glob("**/*.md"):
+        with open(file, "r") as f:
+            data = f.read()
+            key = str(file).replace(str(path), "")
+            content[key] = data
+    return content
+
+
+def split_documents(content: dict[str, str]) -> dict[str, Any]:
+    """Split documents into chunks."""
+    md_splitted_docs = []
+    markdown_splitter = MarkdownHeaderTextSplitter(
+        headers_to_split_on=[("#", "Header 1"), ("##", "Header 2"), ("###", "Header 3")]
+    )
+    text_splitter = RecursiveCharacterTextSplitter(
+        separators=["\n"], chunk_size=500, chunk_overlap=50, length_function=len
+    )
+
+    for file, raw_doc in content.items():
+        splitted_text = markdown_splitter.split_text(raw_doc)
+        for t in splitted_text:
+            t.metadata["source"] = file
+        md_splitted_docs.extend(splitted_text)
+
+    docs = text_splitter.split_documents(md_splitted_docs)
+    docs_as_dict = [doc.dict() for doc in docs]
+    return docs_as_dict
+
+
+def get_embeddings(text: str) -> np.ndarray:
+    """Get embeddings."""
+    return np.array(EMBEDDING_MODEL.run(text))
+
+
+def embed_documents(
+    chunked_docs: dict[str, Any], key: str = "page_content"
+) -> tuple[dict[str, Any], np.ndarray]:
+    """Embed documents."""
+    all_embeddings = []
+    for doc in tqdm(chunked_docs):
+        emb = get_embeddings(doc[key])
+        doc["embedding"] = emb
+        all_embeddings.append(doc["embedding"])
+    full_embedding_mat = np.vstack(all_embeddings)
+    return chunked_docs, full_embedding_mat
+
+
+def query_docs(
+    query: str,
+    docs: dict[str, Any],
+    embedding_mat: np.ndarray,
+    top_n: int = 10,
+    key: str = "page_content",
+) -> tuple[list[int], list[str]]:
+    """Query documents."""
+    query_embedding = get_embeddings(query)
+    scores = embedding_mat.dot(query_embedding)
+    sorted_indices = np.argsort(scores)[::-1]
+    top_n_indices = sorted_indices[:top_n]
+    top_n_indices_rev = top_n_indices[::-1]
+    returned_docs = []
+    for i in top_n_indices_rev:
+        returned_docs.append(docs[i][key])
+    return top_n_indices_rev.tolist(), returned_docs

duckdb-nsql/eval/evaluate.py

@@ -0,0 +1,392 @@
+"""Evaluate text2sql spider model predictions."""
+import json
+import os
+import re
+import signal
+import sys
+import traceback
+from pathlib import Path
+from typing import Any
+
+import click
+import pandas as pd
+from rich.console import Console
+from tqdm.auto import tqdm
+
+sys.path.append(os.path.join(os.path.dirname(__file__), "."))
+# from metrics.spider import evaluation as spider_evaluation  # type: ignore # noqa: E402
+from metrics.test_suite_sql_eval import (  # type: ignore # noqa: E402
+    evaluation as test_suite_evaluation,
+)
+from data_utils import read_tables_json  # type: ignore  # noqa: E402
+from metric_utils import (  # type: ignore  # noqa: E402
+    correct_casing,
+    edit_distance,
+)
+
+console = Console(soft_wrap=True)
+
+LEVELS = ["easy", "medium", "hard", "duckdb", "ddl", "all"]
+PARTIAL_TYPES = [
+    "select",
+    "select(no AGG)",
+    "where",
+    "where(no OP)",
+    "group(no Having)",
+    "group",
+    "order",
+    "and/or",
+    "IUEN",
+    "keywords",
+]
+TIMEOUT_SECONDS = 30
+
+
+def timeout_handler(signum: int, frame: Any) -> None:
+    raise TimeoutError("Function execution timed out.")
+
+
+def print_scores(scores: dict, model_name: str, metric_type: str = "exec") -> None:
+    """Print scores."""
+
+    def print_formated_s(
+        row_name: str, l: list[str], element_format: str = "{}", sep: str = "\t"
+    ) -> None:
+        template = "{}" + sep + sep.join([element_format] * len(l))
+        console.print(template.format(row_name, *l))
+
+    # Add empty scores for each level if not present
+    for level in LEVELS:
+        if level not in scores:
+            scores[level] = {}
+            scores[level]["count"] = 0
+            scores[level]["exec"] = 0
+            scores[level]["exact"] = 0
+
+    print_formated_s("", LEVELS)
+    counts = [scores[level]["count"] for level in LEVELS]
+    print_formated_s("count", counts)
+    console.print(f">======================   {model_name}     =====================")
+    if metric_type == "exec":
+        console.print(
+            ">=====================   EXECUTION ACCURACY     ====================="
+        )
+        exec_scores = [scores[level]["exec"] for level in LEVELS]
+        print_formated_s("execution", exec_scores, element_format="{:.3f}")
+
+    elif metric_type == "exact":
+        console.print(
+            "\n>====================== EXACT MATCHING ACCURACY ====================="
+        )
+        exact_scores = [scores[level]["exact"] for level in LEVELS]
+        print_formated_s("exact match", exact_scores, element_format="{:.3f}")
+
+
+def compute_exact_match_metric(
+    predictions: list,
+    references: list,
+    gold_dbs: list,
+    kmaps: dict,
+    db_dir: str,
+    categories,
+) -> dict:
+    """Compute exact match metric."""
+    exact_match = {}
+    exact_match["all"] = {}
+    exact_match["all"]["count"] = 0
+    exact_match["all"]["exact"] = 0
+    for prediction, reference, gold_db, category in tqdm(
+        zip(predictions, references, gold_dbs, categories), total=len(predictions)
+    ):
+        if category not in exact_match:
+            exact_match[category] = {}
+            exact_match[category]["count"] = 0
+            exact_match[category]["exact"] = 0
+        exact_match["all"]["count"] += 1
+        exact_match[category]["count"] += 1
+        try:
+            match = int(prediction.trim() == reference.trim())
+            exact_match[category]["exact"] += match
+            exact_match["all"]["exact"] += match
+        except Exception:
+            pass
+    return exact_match
+
+
+def compute_test_suite_metric(
+    predictions: list,
+    references: list,
+    gold_dbs: list,
+    setup_sqls: list,
+    validate_sqls: list,
+    kmaps: dict,
+    db_dir: str,
+    categories: list[str] = None,
+) -> tuple[Any, list[int | None]]:
+    """Compute test suite execution metric."""
+    evaluator = test_suite_evaluation.Evaluator(
+        db_dir=db_dir,
+        kmaps=kmaps,
+        etype="exec",
+        plug_value=False,
+        keep_distinct=False,
+        progress_bar_for_each_datapoint=False,
+    )
+    # Only used for Sparc/CoSQL
+    turn_scores: dict[str, list] = {"exec": [], "exact": []}
+    by_row_metrics: list[int | None] = []
+    for prediction, reference, gold_db, setup_sql, validate_sql, category in tqdm(
+        zip(predictions, references, gold_dbs, setup_sqls, validate_sqls, categories),
+        total=len(predictions),
+    ):
+        turn_idx = 0
+        # skip final utterance-query pairs
+        if turn_idx < 0:
+            continue
+
+        # Register the timeout handler function
+        signal.signal(signal.SIGALRM, timeout_handler)
+        signal.alarm(TIMEOUT_SECONDS)
+
+        try:
+            ex_metrics = evaluator.evaluate_one(
+                gold_db,
+                reference,
+                prediction,
+                setup_sql,
+                validate_sql,
+                turn_scores,
+                idx=turn_idx,
+                category=category,
+            )
+            signal.alarm(0)
+
+            by_row_metrics.append(int(ex_metrics["exec"]))
+        except Exception as e:
+            raise e
+            by_row_metrics.append(None)
+            pass
+    evaluator.finalize()
+    return evaluator.scores, by_row_metrics
+
+
+def compute_metrics(
+    gold_sqls: list[str],
+    pred_sqls: list[str],
+    gold_dbs: list[str],
+    setup_sqls: list[str],
+    validate_sqls: list[str],
+    kmaps: dict,
+    db_schemas: dict,
+    database_dir: str,
+    lowercase_schema_match: bool,
+    model_name: str,
+    categories: list[str] = None,
+) -> dict[str, str]:
+    """Compute all metrics for data slice."""
+    if len(gold_sqls) != len(pred_sqls):
+        raise ValueError(
+            f"Gold {len(gold_sqls)} and pred {len(pred_sqls)} have different number of lines!"
+        )
+    all_metrics: dict[str, Any] = {}
+
+    # Execution Accuracy
+    metrics, by_row_metrics = compute_test_suite_metric(
+        pred_sqls,
+        gold_sqls,
+        gold_dbs,
+        setup_sqls,
+        validate_sqls,
+        kmaps,
+        database_dir,
+        categories,
+    )
+    all_metrics["exec"] = metrics
+    all_metrics["by_row_exec"] = by_row_metrics
+    print_scores(metrics, model_name, "exec")
+
+    # Exact Match Accuracy
+    metrics = compute_exact_match_metric(
+        pred_sqls, gold_sqls, gold_dbs, kmaps, database_dir, categories
+    )
+    all_metrics["exact"] = metrics
+    print_scores(metrics, model_name, "exact")
+
+    # Equality Accuracy
+    per_row_match = [
+        int(gold.lower() == pred.lower()) for gold, pred in zip(gold_sqls, pred_sqls)
+    ]
+    all_metrics["equality"] = {"equality": sum(per_row_match) / len(gold_sqls)}
+    all_metrics["by_row_equality"] = per_row_match
+
+    # Edit Distance
+    per_row_edit_dist = [
+        edit_distance(gold, pred) for gold, pred in zip(gold_sqls, pred_sqls)
+    ]
+    edit_dist = sum(per_row_edit_dist) / len(gold_sqls)
+    all_metrics["edit_distance"] = {"edit_distance": edit_dist}
+    all_metrics["by_row_edit_distance"] = per_row_edit_dist
+
+    return all_metrics
+
+
+def get_to_print(metrics: dict, key: str, model_name: str, num_rows: int) -> dict:
+    """Get pretty print dictionary of metrics."""
+    return {
+        "slice": key,
+        "model": model_name,
+        "support": num_rows,
+        "exec": f"{metrics[key]['exec']['all']['exec']:.3f}",
+        "exact": f"{metrics[key]['exact']['all']['exact']:.3f}",
+        "equality": f"{metrics[key]['equality']['equality']:.3f}",
+        "edit_distance": f"{metrics[key]['edit_distance']['edit_distance']:.3f}",
+    }
+
+
+@click.group()
+def cli() -> None:
+    """Entrypoint."""
+    pass
+
+
+@cli.command()
+@click.option("--gold", type=str, required=True)
+@click.option("--pred", type=str, required=True)
+@click.option("--tables", type=str, required=True)
+@click.option("--db", type=str, default="")
+@click.option("--slice-attribute", type=str, default=None)
+@click.option("--output-dir", type=str, default="")
+@click.option("--output-filename", type=str, default="")
+@click.option(
+    "--correct-sql-casing", type=bool, is_flag=True, default=False, required=False
+)
+@click.option(
+    "--lowercase-schema-match", type=bool, is_flag=True, default=False, required=False
+)
+def evaluate(
+    gold: str,
+    pred: str,
+    tables: str,
+    db: str,
+    slice_attribute: str,
+    output_dir: str,
+    output_filename: str,
+    correct_sql_casing: bool,
+    lowercase_schema_match: bool,
+) -> None:
+    """Evaluate SQL.
+
+    Args:
+        gold: path to gold sql file.
+        pred: path to predicted json lines file.
+        tables: the json path of the table metadata.
+        db: path to database dir.
+        slice_attribute: json attribute in gold data to slice on.
+        output_dir: the prediction output directory
+        output_filename: the prediction output filename
+        correct_sql_casing: whether to correct casing of SQL keywords
+        lowercase_schema_match: whether to lowercase schema match
+    """
+    gold_path = Path(gold)
+    pred_path = Path(pred)
+    model_name = pred_path.stem
+    if not output_filename:
+        output_filename = pred_path.stem + "_eval.json"
+    console.print(f"Saving to {Path(output_dir) / output_filename}")
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+
+    kmaps = test_suite_evaluation.build_foreign_key_map_from_json(tables)
+    db_schemas = read_tables_json(tables)
+
+    gold_sqls_dict = json.load(gold_path.open("r", encoding="utf-8"))
+    pred_sqls_dict = [json.loads(l) for l in pred_path.open("r").readlines()]
+
+    # Data validation
+    assert len(gold_sqls_dict) == len(
+        pred_sqls_dict
+    ), "Sample size doesn't match between pred and gold file"
+
+    # Keep track of everything
+    full_results = []
+    for gold_sql, pred_sql in zip(gold_sqls_dict, pred_sqls_dict):
+        merged_res = {**pred_sql, **gold_sql}
+        full_results.append(merged_res)
+
+    gold_sqls = [
+        re.sub(r"[\s\t\n]+", " ", p.get("gold", p.get("query", p.get("sql", ""))))
+        for p in gold_sqls_dict
+    ]
+    setup_sqls = [re.sub(r"[\s\t\n]+", " ", p["setup_sql"]) for p in gold_sqls_dict]
+    validate_sqls = [
+        re.sub(r"[\s\t\n]+", " ", p["validation_sql"]) for p in gold_sqls_dict
+    ]
+    gold_dbs = [p.get("db_id", p.get("db", "")) for p in gold_sqls_dict]
+    pred_sqls = [re.sub(r"[\s\t\n]+", " ", p["pred"]) for p in pred_sqls_dict]
+    categories = [p.get("category", "") for p in gold_sqls_dict]
+    if correct_sql_casing:
+        # One line to correct casing of SQL keywords using correct_casing(sql)
+        gold_sqls = [correct_casing(sql) for sql in gold_sqls]
+        pred_sqls = [correct_casing(sql) for sql in pred_sqls]
+
+    final_metrics: dict[str, dict[str, Any]] = {}
+    to_print = []
+    final_metrics["all"] = compute_metrics(
+        gold_sqls=gold_sqls,
+        pred_sqls=pred_sqls,
+        gold_dbs=gold_dbs,
+        setup_sqls=setup_sqls,
+        validate_sqls=validate_sqls,
+        kmaps=kmaps,
+        db_schemas=db_schemas,
+        database_dir=db,
+        lowercase_schema_match=lowercase_schema_match,
+        model_name=model_name + "(all)",
+        categories=categories,
+    )
+
+    for k, v in final_metrics["all"].items():
+        if k.startswith("by_row"):
+            assert len(v) == len(gold_sqls)
+            for dct, val in zip(full_results, v):
+                dct[k[len("by_row_") :]] = val
+    to_print.append(get_to_print(final_metrics, "all", model_name, len(gold_sqls)))
+    # TODO: could be way more efficient if we subsliced the results but...whatever
+    if slice_attribute:
+        for unq_value in sorted(set([g[slice_attribute] for g in gold_sqls_dict])):
+            idx_set = [
+                i
+                for i, g in enumerate(gold_sqls_dict)
+                if g[slice_attribute] == unq_value
+            ]
+            print(f"Processing {unq_value} with {len(idx_set)} samples")
+            final_metrics[unq_value] = compute_metrics(
+                gold_sqls=[gold_sqls[i] for i in idx_set],
+                pred_sqls=[pred_sqls[i] for i in idx_set],
+                gold_dbs=[gold_dbs[i] for i in idx_set],
+                setup_sqls=[setup_sqls[i] for i in idx_set],
+                validate_sqls=[validate_sqls[i] for i in idx_set],
+                kmaps=kmaps,
+                db_schemas=db_schemas,
+                database_dir=db,
+                lowercase_schema_match=lowercase_schema_match,
+                model_name=model_name + f"({unq_value})",
+                categories=[categories[i] for i in idx_set],
+            )
+            to_print.append(
+                get_to_print(final_metrics, unq_value, model_name, len(idx_set))
+            )
+
+    df = pd.DataFrame(to_print)
+    console.print(df.to_csv(sep=",", index=False))
+    console.print("******")
+    console.print(f"Saved metrics to {Path(output_dir) / output_filename}")
+    json.dump(final_metrics, open(Path(output_dir) / output_filename, "w"), indent=4)
+    output_filename = str(output_filename).replace("_eval.json", "_fd.jsonl")
+    console.print(f"Saved dump to {Path(output_dir) / output_filename}")
+    with open(Path(output_dir) / output_filename, "w") as f:
+        for dct in full_results:
+            f.write(json.dumps(dct) + "\n")
+
+
+if __name__ == "__main__":
+    cli()

duckdb-nsql/eval/get_manifest.py

@@ -0,0 +1,29 @@
+"""Manifest utils."""
+from manifest import Manifest
+from manifest.connections.client_pool import ClientConnection
+
+
+def get_manifest(
+    manifest_client: str,
+    manifest_connection: str,
+    manifest_engine: str,
+) -> Manifest:
+    """Get manifest engine."""
+    if manifest_client in {"openai", "openaichat", "openai_mock", "openrouter", "azureendpoint"}:
+        manifest = Manifest(
+            client_name=manifest_client,
+            engine=manifest_engine,
+        )
+    elif manifest_client in {"huggingface"}:
+        manifest = Manifest(
+            client_pool=[
+                ClientConnection(
+                    client_name=manifest_client,
+                    client_connection=manifest_conn,
+                )
+                for manifest_conn in manifest_connection.split(";")
+            ],
+        )
+    else:
+        raise ValueError(f"Unknown manifest client {manifest_client}")
+    return manifest

duckdb-nsql/eval/loaders.py

@@ -0,0 +1,78 @@
+"""Data loaders."""
+import json
+import re
+import string
+from abc import ABC, abstractmethod
+
+from rich.console import Console
+from data_utils import read_tables_json
+from schema import Table
+
+RE_COLUMN = re.compile(r"^select (.+?) from")
+RE_CONDS = re.compile(r"where (.+?)$")
+RE_COND = re.compile(r"^(.+?)\s*([=><])\s*(.+?)$")
+
+translator = str.maketrans(
+    string.punctuation, " " * len(string.punctuation)
+)  # map punctuation to space
+
+console = Console(soft_wrap=True)
+
+
+def standardize_column(col: str) -> str:
+    """Standardize the column name to SQL compatible."""
+    col_name = col.replace("#", "num").replace("%", "perc")
+    col_name = col_name.strip().lower().translate(translator)
+    col_name = re.sub("[^0-9a-z ]", " ", col_name).strip()
+    col_name = re.sub(" +", "_", col_name)
+    if not col_name:
+        console.print(f"original {col}, new {col_name}")
+    return col_name
+
+
+def clean_col(col: str) -> str:
+    """Remove table name and standardize column name."""
+    if "." in col and not col.endswith("."):
+        col = col.split(".")[-1]
+    return standardize_column(col)
+
+
+class Loader(ABC):
+    """Loader abstract class."""
+
+    @classmethod
+    @abstractmethod
+    def load_data(cls, path: str) -> list[dict]:
+        """Load data from path."""
+
+    @classmethod
+    @abstractmethod
+    def load_table_metadata(cls, path: str) -> dict[str, dict[str, Table]]:
+        """Extract table metadata from table-metadata-path."""
+
+    @classmethod
+    def format_output(cls, prediction: dict) -> dict:
+        """Parse for spider format."""
+        return prediction
+
+
+class DefaultLoader(Loader):
+    """Spider loader and writer."""
+
+    @classmethod
+    def load_data(cls, path: str) -> list[dict]:
+        """Load data from path."""
+        try:
+            with open(path) as f:
+                data = json.loads(f.read())
+        except json.decoder.JSONDecodeError:
+            # Try with jsonl
+            data = [json.loads(line) for line in open(path)]
+        return data
+
+    @classmethod
+    def load_table_metadata(cls, path: str) -> dict[str, dict[str, Table]]:
+        """Extract table metadata from table-metadata-path."""
+        # load the tables
+        db_to_tables = read_tables_json(path, lowercase=True)
+        return db_to_tables

duckdb-nsql/eval/metric_utils.py

@@ -0,0 +1,40 @@
+"""Utility metrics."""
+import sqlglot
+from rich.console import Console
+from sqlglot import parse_one
+
+console = Console(soft_wrap=True)
+
+
+def correct_casing(sql: str) -> str:
+    """Correct casing of SQL."""
+    parse: sqlglot.expressions.Expression = parse_one(sql, read="sqlite")
+    return parse.sql()
+
+
+def prec_recall_f1(gold: set, pred: set) -> dict[str, float]:
+    """Compute precision, recall and F1 score."""
+    prec = len(gold.intersection(pred)) / len(pred) if pred else 0.0
+    recall = len(gold.intersection(pred)) / len(gold) if gold else 0.0
+    f1 = 2 * prec * recall / (prec + recall) if prec + recall else 0.0
+    return {"prec": prec, "recall": recall, "f1": f1}
+
+
+def edit_distance(s1: str, s2: str) -> int:
+    """Compute edit distance between two strings."""
+    # Make sure s1 is the shorter string
+    if len(s1) > len(s2):
+        s1, s2 = s2, s1
+
+    distances: list[int] = list(range(len(s1) + 1))
+    for i2, c2 in enumerate(s2):
+        distances_ = [i2 + 1]
+        for i1, c1 in enumerate(s1):
+            if c1 == c2:
+                distances_.append(distances[i1])
+            else:
+                distances_.append(
+                    1 + min((distances[i1], distances[i1 + 1], distances_[-1]))
+                )
+        distances = distances_
+    return distances[-1]

duckdb-nsql/eval/metrics/test_suite_sql_eval

@@ -0,0 +1 @@
+Subproject commit 640a12975abf75a94e917caca149d56dbc6bcdd7

duckdb-nsql/eval/predict.py

@@ -0,0 +1,345 @@
+"""Run dataset on text2sql zazu experiment.
+
+See README.md for more details.
+"""
+import datetime
+import json
+import multiprocessing
+import random
+import re
+from pathlib import Path
+
+import click
+import numpy as np
+from constants import PROMPT_FORMATTERS
+from loaders import DefaultLoader
+from get_manifest import get_manifest
+from manifest import Manifest
+from prompt_formatters import RajkumarFormatter
+from rich.console import Console
+from schema import Table, TextToSQLModelResponse, TextToSQLParams
+from text_to_sql import instruction_to_sql, instruction_to_sql_list
+from doc_retriever import (
+    load_documentation,
+    split_documents,
+    embed_documents,
+    query_docs,
+)
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+console = Console(soft_wrap=True)
+
+
+def generate_sql(
+    manifest: Manifest,
+    text_to_sql_in: list[TextToSQLParams],
+    retrieved_docs: list[list[str]],
+    prompt_formatter: RajkumarFormatter,
+    stop_tokens: list[str] | None = None,
+    overwrite_manifest: bool = False,
+    max_tokens: int = 300,
+    temperature: float = 0.1,
+    num_beams: int = 2,
+    parallel: bool = False,
+) -> list[tuple[str, TextToSQLModelResponse]]:
+    """Call our text2sql function with manifest of our choice."""
+    if parallel:
+        instruction_to_sql_resps: list[
+            TextToSQLModelResponse
+        ] = instruction_to_sql_list(
+            params=text_to_sql_in,
+            extra_context=retrieved_docs,
+            manifest=manifest,
+            prompt_formatter=prompt_formatter,
+            overwrite_manifest=overwrite_manifest,
+            max_tokens=max_tokens,
+            temperature=0.1,
+            stop_sequences=stop_tokens,
+            num_beams=num_beams,
+        )
+    else:
+        instruction_to_sql_resps = [
+            instruction_to_sql(
+                params=_text_to_sql_in,
+                extra_context=_retrieved_docs,
+                manifest=manifest,
+                prompt_formatter=prompt_formatter,
+                overwrite_manifest=overwrite_manifest,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                stop_sequences=stop_tokens,
+                num_beams=num_beams,
+            )
+            for _retrieved_docs, _text_to_sql_in in tqdm(
+                zip(retrieved_docs, text_to_sql_in),
+                desc="Generating SQL",
+                total=len(text_to_sql_in),
+                disable=(len(text_to_sql_in) <= 1),
+            )
+        ]
+    assert len(instruction_to_sql_resps) == len(text_to_sql_in)
+
+    sql_statements = []
+    for i in range(len(instruction_to_sql_resps)):
+        sql_statement = instruction_to_sql_resps[i].output.strip()
+        if "<>" in sql_statement:
+            sql_statement.replace("<>", "!=")
+        # Models sometime train to predict <databasename/schema> | <sql>
+        sql_statement = sql_statement.split("|")[-1].strip()
+        sql_statements.append(sql_statement)
+    return list(zip(sql_statements, instruction_to_sql_resps))
+
+
+def get_text_to_sql_in(
+    input_question: dict, db_to_tables: dict[str, dict[str, Table]]
+) -> TextToSQLParams:
+    """Format input question for text2sql function."""
+    question = input_question["question"]
+    db_id = input_question.get("db_id", None)
+    if db_id != "none":
+        table_params = list(db_to_tables.get(db_id, {}).values())
+    else:
+        table_params = []
+    if len(table_params) == 0:
+        console.print(f"[red] WARNING: No tables found for {db_id} [/red]")
+    text_to_sql_in = TextToSQLParams(
+        instruction=question,
+        database=db_id,
+        tables=table_params,
+    )
+    return text_to_sql_in
+
+
+@click.group()
+def cli() -> None:
+    """Entrypoint."""
+    pass
+
+
+@cli.command()
+@click.argument("dataset-path")
+@click.argument("table-meta-path")
+@click.option("--output-dir", type=str, default="")
+@click.option("--run-name", type=str, default="")
+@click.option("--num-run", type=int, default=-1)
+@click.option("--num-print", type=int, default=20)
+# Format options
+@click.option("--prompt-format", type=str, default="spider")
+# Prompt options
+@click.option("--stop-tokens", type=str, default=[], multiple=True)
+@click.option("--max-tokens", type=int, default=200)
+@click.option("--temperature", type=float, default=0)
+@click.option("--num-beams", type=int, default=-1)  # use whatever is in manifest
+@click.option("--max-context-length", type=int, default=-1)
+# Docs options
+@click.option(
+    "--markdown-docs-path",
+    #type=click.Path(
+    #    exists=True, file_okay=True, dir_okay=True, readable=True, path_type=Path
+    #),
+    default="eval/docs/duckdb-web/docs/archive/0.9.2/sql",
+)
+@click.option("--num-retrieved-docs", type=int, default=0)
+# Manifest options
+@click.option("--manifest-client", type=str, default="openai")
+@click.option("--manifest-engine", type=str, default="gpt-4o")
+@click.option("--manifest-connection", type=str, default="http://localhost:5005")
+@click.option("--overwrite-manifest", is_flag=True, default=False)
+@click.option("--parallel", is_flag=True, default=False)
+def predict(
+    dataset_path: str,
+    table_meta_path: str,
+    output_dir: str,
+    run_name: str,
+    num_run: int,
+    num_print: int,
+    prompt_format: str,
+    stop_tokens: list[str],
+    max_tokens: int,
+    temperature: float,
+    num_beams: int,
+    max_context_length: int,
+    markdown_docs_path: Path,
+    num_retrieved_docs: int,
+    manifest_client: str,
+    manifest_engine: str,
+    manifest_connection: str,
+    overwrite_manifest: bool,
+    parallel: bool,
+) -> None:
+    """Predict SQL.
+
+    Args:
+        dataset_path: the dataset path.
+        table_meta_path: the json path of the table metadata.
+        database_path: the database path for sqlite.
+        output_dir: the prediction output directory
+        run_name: special prefix to add to filename
+        num_run: the number of examples to run
+        num_print: the number of examples to print
+        prompt_format: the format of the prompt. E.g., "rajkumar"
+        stop_tokens: the stop tokens to try
+        max_tokens: the max tokens
+        temperature: the temperature
+        num_beams: the number of beams
+        max_context_length: max context length for demonstration truncation (-1 means None)
+        markdown_docs_path: path to duckdb sql docs
+        num_retrieved_docs: number of docs to retrieve
+        manifest_client: the manifest client
+        manifest_engine: the manifest engine
+        manifest_connection: the manifest connection
+    """
+    multiprocessing.set_start_method("spawn", force=True)
+    random.seed(0)
+    np.random.seed(0)
+    locals_dict = locals()
+    locals_dict["markdown_docs_path"] = str(markdown_docs_path)
+    console.print(json.dumps(locals_dict, indent=2))
+
+    data_formatter = DefaultLoader()
+
+    if prompt_format not in PROMPT_FORMATTERS:
+        raise ValueError(f"Unknown prompt format {prompt_format}")
+    prompt_formatter = PROMPT_FORMATTERS[prompt_format]()
+
+    # load manifest
+    manifest = get_manifest(
+        manifest_client=manifest_client,
+        manifest_connection=manifest_connection,
+        manifest_engine=manifest_engine,
+    )
+    manifest_params = manifest.client_pool.get_current_client().get_model_params()
+    console.print(f"Running with {manifest_params} manifest.")
+    model_name = manifest_params.get("engine", manifest_params["model_name"])
+
+    if manifest_client in {"openai", "openaichat", "openrouter", "azureendpoint"}:
+        tokenizer = AutoTokenizer.from_pretrained("gpt2", trust_remote_code=True)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+    if stop_tokens:
+        stop_tokens = [st.strip("'") for st in stop_tokens]
+    console.print(f"Stop tokens: {stop_tokens}")
+
+    # Get output filename
+    full_dataset_path = Path(dataset_path)
+    # Get todays date
+    date_today = datetime.datetime.now().strftime("%y-%m-%d")
+    if run_name:
+        run_name = f"{run_name}_"
+    suffix = f"{run_name}{full_dataset_path.stem}_{date_today}.json"  # noqa: E501
+    prefix = f"{prompt_format}_{num_retrieved_docs}docs"
+    if manifest_client in {"openai", "openaiazure"}:
+        middleix = manifest_engine
+    elif manifest_client in {"huggingface", "ray"}:
+        middleix = Path(manifest_params.get("model_path", "")).name.replace("/", "-")
+    elif manifest_client in {"toma", "openrouter", "openaichat", "azureendpoint"}:
+        middleix = manifest_engine.split("/")[-1]
+    else:
+        raise ValueError(f"Unknown manifest client {manifest_client}")
+    output_filename = f"{prefix}_{middleix}_{suffix}"
+    console.print(f"Saving to {Path(output_dir) / output_filename}")
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+
+    console.print("Loading metadata...")
+    db_to_tables = data_formatter.load_table_metadata(table_meta_path)
+
+    console.print("Loading data...")
+    data = data_formatter.load_data(dataset_path)
+    if num_run > 0:
+        console.print(f"Running on {min(len(data), num_run)} examples")
+        data = data[:num_run]
+    original_data = data
+
+    # load the examples
+    console.print("Formatting data...")
+    num_print = min(num_print, len(data))
+    token_lengths = []
+    text_to_sql_in = [
+        get_text_to_sql_in(input_question, db_to_tables) for input_question in data
+    ]
+
+    if num_retrieved_docs > 0:
+        console.print("Loading documenration and indexing...")
+        retrieved_docs = []
+        doc_contents = load_documentation(markdown_docs_path)
+        chunked_docs = split_documents(doc_contents)
+        embedded_docs, full_embedding_mat = embed_documents(chunked_docs)
+        for i in tqdm(range(len(text_to_sql_in)), desc="Retrieving docs"):
+            _, retrieved_docs_strings = query_docs(
+                text_to_sql_in[i].instruction,
+                embedded_docs,
+                full_embedding_mat,
+                top_n=num_retrieved_docs,
+            )
+            retrieved_docs.append(retrieved_docs_strings)
+    else:
+        retrieved_docs = [[] for _ in range(len(text_to_sql_in))]
+
+    for i in range(num_print):
+        # Run a few to get some examples to print
+        generated_responses = generate_sql(
+            manifest=manifest,
+            text_to_sql_in=[text_to_sql_in[i]],
+            retrieved_docs=[retrieved_docs[i]],
+            stop_tokens=stop_tokens,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            num_beams=num_beams,
+            prompt_formatter=prompt_formatter,
+            overwrite_manifest=overwrite_manifest,
+            parallel=parallel,
+        )
+        for prediction, model_response in generated_responses:
+            prediction = re.sub(r"[\s\t\n]+", " ", prediction)
+            token_lengths.append(len(tokenizer(prediction).input_ids))
+            console.print(f"[blue]Prompt:[/blue] {model_response.final_prompt}")
+            console.print(f"[red]Prediction:[/red] {prediction}")
+            if data[i].get("query") or data[i].get("sql"):
+                console.print(
+                    "[purple]Gold:[/purple] "
+                    f"{data[i].get('query') or data[i].get('sql')}"
+                )
+            console.print("\n****\n")
+
+    # Run the entire thing now - the to_print results will be in cache and fast
+    generated_sqls = generate_sql(
+        manifest=manifest,
+        text_to_sql_in=text_to_sql_in,
+        retrieved_docs=retrieved_docs,
+        stop_tokens=stop_tokens,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        num_beams=num_beams,
+        prompt_formatter=prompt_formatter,
+        overwrite_manifest=overwrite_manifest,
+        parallel=parallel,
+    )
+
+    with open(Path(output_dir) / output_filename, "w") as fout:
+        for i, (prediction, model_response) in enumerate(generated_sqls):
+            if isinstance(model_response.final_prompt, str):
+                token_lengths.append(
+                    len(tokenizer(model_response.final_prompt).input_ids)
+                )
+            else:
+                for prompt in model_response.final_prompt:
+                    token_lengths.append(len(tokenizer(prompt["content"]).input_ids))
+            entry = {
+                **original_data[i],
+                "pred": prediction,
+                "raw_pred": model_response.output,
+                "raw_output": model_response.raw_output,
+                "prompt": model_response.final_prompt,
+                "tables": [tbl.dict() for tbl in text_to_sql_in[i].tables or []],
+            }
+            formatted_entry = data_formatter.format_output(entry)
+            print(json.dumps(formatted_entry), file=fout)
+    overflow = len([tl for tl in token_lengths if tl > 2048]) / len(token_lengths)
+    console.print(f"Overflow 2048 prompt {100*overflow:.2f}%")
+    console.print(f"Saved to {Path(output_dir) / output_filename}")
+
+
+if __name__ == "__main__":
+    cli()

duckdb-nsql/eval/prompt_formatters.py

@@ -0,0 +1,1056 @@
+"""Rajkumar prompt formatter."""
+
+from random import shuffle
+from manifest import Manifest
+from schema import Table
+import re
+
+
+class RajkumarFormatter:
+    """RajkumarFormatter class.
+
+    From https://arxiv.org/pdf/2204.00498.pdf.
+    """
+
+    table_sep: str = "\n\n"
+    shuffle_table_order: bool = True
+    _cache: dict[tuple[str, str, str], list[str]] = {}
+    clean_whitespace = False
+
+    @classmethod
+    def format_table(cls, table: Table) -> str:
+        """Get table format."""
+        table_fmt = []
+        for col in table.columns or []:
+            # This is technically an incorrect type, but it should be a catchall word
+            table_fmt.append(f"    {col.name} {col.dtype or 'any'}")
+        if table_fmt:
+            all_cols = ",\n".join(table_fmt)
+            create_tbl = f"CREATE TABLE {table.name} (\n{all_cols}\n)"
+        else:
+            create_tbl = f"CREATE TABLE {table.name}"
+        return create_tbl
+
+    @classmethod
+    def format_all_tables(cls, tables: list[Table], instruction: str) -> list[str]:
+        """Get all tables format."""
+        table_texts = [cls.format_table(table) for table in tables]
+        key = ("tables", instruction, str(tables))
+        if key not in cls._cache:
+            shuffle(table_texts)
+            cls._cache[key] = table_texts
+        else:
+            table_texts = cls._cache[key]
+        return table_texts
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n\n/*\nHere is additional documentation about DuckDB that could be useful.\n--------\n{context_str}\n--------\n*/"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        return f"""{table_text}\n\n\n-- Using valid DuckDB SQL, answer the following question for the tables provided above.{context_text}\n\n-- {instruction}\n"""  # noqa: E501
+
+    @classmethod
+    def format_model_output(cls, output_sql: str, prompt: str) -> str:
+        """Format model output."""
+        clean_sql = (output_sql
+            .replace('```sql\n', '')
+            .replace('```duckdb\n', '')
+            .replace('```\n', '')
+            .replace('```', '')).strip()
+
+        if clean_sql.find(';') != -1:
+            clean_sql[:clean_sql.find(';')].strip()
+
+        if not clean_sql.endswith(";"):
+            clean_sql += ";"
+
+        return clean_sql
+
+    @classmethod
+    def format_gold_output(cls, output_sql: str) -> str:
+        """Format gold output for demonstration."""
+        return output_sql
+
+class MotherDuckFormatter(RajkumarFormatter):
+    """MotherDuck class."""
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        return f"""{table_text}\n\n\n-- Using valid DuckDB SQL, answer the following question for the tables provided above.{context_text}\n\n-- {instruction}\n```sql\n"""  # noqa: E501
+
+
+class DuckDBFormatter(RajkumarFormatter):
+    """DuckDB class."""
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        return f"""{table_text}\n\n\n-- Using valid DuckDB SQL, answer the following question for the tables provided above.{context_text}\n\n-- {instruction}\n```sql\n"""  # noqa: E501
+
+
+class DuckDBInstFormatter(RajkumarFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """### Instruction:\n{instruction}\n\n### Input:\n{input}{context}\n### Question:\n{question}\n\n### Response (use duckdb shorthand if possible):\n"""
+    INSTRUCTION_TEMPLATE = """Your task is to generate valid duckdb SQL to answer the following question{has_schema}"""  # noqa: E501
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n### Documentation:\n{context_str}\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        input = ""
+        if table_text:
+            input = """Here is the database schema that the SQL query will run on:\n{schema}\n""".format(  # noqa: E501
+                schema=table_text
+            )
+        instruction = cls.PROMPT_TEMPLATE.format(
+            instruction=cls.INSTRUCTION_TEMPLATE.format(
+                has_schema="."
+                if table_text == ""
+                else ", given a duckdb database schema."
+            ),
+            context=context_text,
+            input=input,
+            question=instruction,
+        )
+        return instruction
+
+class DuckDBInstFormatterLlamaShort(RajkumarFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+Your task is to generate valid DuckDB SQL to answer the question that the user asks. You should only respond with a valid DuckDB SQL query.
+
+Here are some DuckDB SQL syntax specifics you should be aware of:
+
+- DuckDB uses double quotes (") for identifiers that contain spaces or special characters, or to force case-sensitivity and single quotes (') to define string literals
+- DuckDB can query CSV, Parquet, and JSON directly without loading them first, e.g. `SELECT * FROM 'data.csv';`
+- DuckDB supports CREATE TABLE AS (CTAS): `CREATE TABLE new_table AS SELECT * FROM old_table;`
+- DuckDB queries can start with FROM, and optionally omit SELECT *, e.g. `FROM my_table WHERE condition;` is equivalent to `SELECT * FROM my_table WHERE condition;`
+- DuckDB allows you to use SELECT without a FROM clause to generate a single row of results or to work with expressions directly, e.g. `SELECT 1 + 1 AS result;`
+- DuckDB supports attaching multiple databases, unsing the ATTACH statement: `ATTACH 'my_database.duckdb' AS mydb;`. Tables within attached databases can be accessed using the dot notation (.), e.g. `SELECT * FROM mydb.table_name syntax`. The default databases doesn't require the do notation to access tables. The default database can be changed with the USE statement, e.g. `USE my_db;`.
+- DuckDB is generally more lenient with implicit type conversions (e.g. `SELECT '42' + 1;` - Implicit cast, result is 43), but you can always be explicit using `::`, e.g. `SELECT '42'::INTEGER + 1;`
+- DuckDB can extract parts of strings and lists using [start:end] or [start:end:step] syntax. Indexes start at 1. String slicing: `SELECT 'DuckDB'[1:4];`. Array/List slicing: `SELECT [1, 2, 3, 4][1:3];`
+- DuckDB has a powerful way to select or transform multiple columns using patterns or functions. You can select columns matching a pattern: `SELECT COLUMNS('sales_.*') FROM sales_data;` or transform multiple columns with a function: `SELECT AVG(COLUMNS('sales_.*')) FROM sales_data;`
+- DuckDB has an easy way to include/exclude or modify columns when selecting all: e.g. Exclude: `SELECT * EXCLUDE (sensitive_data) FROM users;` Replace: `SELECT * REPLACE (UPPER(name) AS name) FROM users;`
+- DuckDB has a shorthand for grouping/ordering by all non-aggregated/all columns. e.g `SELECT category, SUM(sales) FROM sales_data GROUP BY ALL;` and `SELECT * FROM my_table ORDER BY ALL;`
+- DuckDB can combine tables by matching column names, not just their positions using UNION BY NAME. E.g. `SELECT * FROM table1 UNION BY NAME SELECT * FROM table2;`
+- DuckDB has an inutitive syntax to create List/Struct/Map and Array types. Create complex types using intuitive syntax. List: `SELECT [1, 2, 3] AS my_list;`, Struct: `{{'a': 1, 'b': 'text'}} AS my_struct;`, Map: `MAP([1,2],['one','two']) as my_map;`. All types can also be nested into each other. Array types are fixed size, while list types have variable size. Compared to Structs, MAPs do not need to have the same keys present for each row, but keys can only be of type Integer or Varchar. Example: `CREATE TABLE example (my_list INTEGER[], my_struct STRUCT(a INTEGER, b TEXT), my_map MAP(INTEGER, VARCHAR),  my_array INTEGER[3], my_nested_struct STRUCT(a INTEGER, b Integer[3]));`
+- DuckDB has an inutive syntax to access struct fields using dot notation (.) or brackets ([]) with the field name. Maps fields can be accessed by brackets ([]).
+- DuckDB's way of converting between text and timestamps, and extract date parts. Current date as 'YYYY-MM-DD': `SELECT strftime(NOW(), '%Y-%m-%d');` String to timestamp: `SELECT strptime('2023-07-23', '%Y-%m-%d')::TIMESTAMP;`, Extract Year from date: `SELECT EXTRACT(YEAR FROM DATE '2023-07-23');`
+- Column Aliases in WHERE/GROUP BY/HAVING: You can use column aliases defined in the SELECT clause within the WHERE, GROUP BY, and HAVING clauses. E.g.: `SELECT a + b AS total FROM my_table WHERE total > 10 GROUP BY total HAVING total < 20;`
+- DuckDB allows generating lists using expressions similar to Python list comprehensions. E.g. `SELECT [x*2 FOR x IN [1, 2, 3]];` Returns [2, 4, 6].
+- DuckDB allows chaining multiple function calls together using the dot (.) operator. E.g.: `SELECT 'DuckDB'.replace('Duck', 'Goose').upper(); -- Returns 'GOOSEDB';`
+- DuckDB has a JSON data type. It supports selecting fields from the JSON with a JSON-Path expression using the arrow operator, -> (returns JSON) or ->> (returns text) with JSONPath expressions. For example: `SELECT data->'$.user.id' AS user_id, data->>'$.event_type' AS event_type FROM events;`
+- DuckDB has built-in functions for regex regexp_matches(column, regex), regexp_replace(column, regex), and regexp_extract(column, regex).
+- DuckDB has a way to quickly get a subset of your data with `SELECT * FROM large_table USING SAMPLE 10%;`<|eot_id|>
+<|start_header_id|>user<|end_header_id|>
+
+Database Schema:
+Here is the schema of the DuckDB database that the SQL query will run on:
+```sql
+{schema}
+```
+
+Question:
+Here is the question or an instruction the user provided:
+{question}
+
+Task:
+Write a DuckDB SQL query for the given question!<|eot_id|>
+<|start_header_id|>assistant<|end_header_id|>
+
+```sql
+"""
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n### Documentation:\n{context_str}\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        instruction = cls.PROMPT_TEMPLATE.format(
+            schema=table_text,
+            question=instruction
+        )
+        return instruction
+
+class DuckDBInstFormatterLlamaBasic(RajkumarFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+Your task is to generate valid DuckDB SQL to answer the question that the user asks. You should only respond with a valid DuckDB SQL query.<|eot_id|>
+<|start_header_id|>user<|end_header_id|>
+
+Database Schema:
+Here is the schema of the DuckDB database that the SQL query will run on:
+```sql
+{schema}
+```
+
+Question:
+Here is the question or an instruction the user provided:
+{question}
+
+Task:
+Write a DuckDB SQL query for the given question!<|eot_id|>
+<|start_header_id|>assistant<|end_header_id|>
+
+```sql
+"""
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n### Documentation:\n{context_str}\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        instruction = cls.PROMPT_TEMPLATE.format(
+            schema=table_text,
+            question=instruction
+        )
+        return instruction
+
+class DuckDBInstFormatterLlamaSyntax(RajkumarFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+Your task is to generate valid DuckDB SQL to answer the question that the user asks. You should only respond with a valid DuckDB SQL query.
+
+Here are some DuckDB SQL syntax specifics you should be aware of:
+
+- DuckDB uses double quotes (") for identifiers that contain spaces or special characters, or to force case-sensitivity and single quotes (') to define string literals
+- DuckDB can query CSV, Parquet, and JSON directly without loading them first, e.g. `SELECT * FROM 'data.csv';`
+- DuckDB supports CREATE TABLE AS (CTAS): `CREATE TABLE new_table AS SELECT * FROM old_table;`
+- DuckDB queries can start with FROM, and optionally omit SELECT *, e.g. `FROM my_table WHERE condition;` is equivalent to `SELECT * FROM my_table WHERE condition;`
+- DuckDB allows you to use SELECT without a FROM clause to generate a single row of results or to work with expressions directly, e.g. `SELECT 1 + 1 AS result;`
+- DuckDB supports attaching multiple databases, unsing the ATTACH statement: `ATTACH 'my_database.duckdb' AS mydb;`. Tables within attached databases can be accessed using the dot notation (.), e.g. `SELECT * FROM mydb.table_name syntax`. The default databases doesn't require the do notation to access tables. The default database can be changed with the USE statement, e.g. `USE my_db;`.
+- DuckDB is generally more lenient with implicit type conversions (e.g. `SELECT '42' + 1;` - Implicit cast, result is 43), but you can always be explicit using `::`, e.g. `SELECT '42'::INTEGER + 1;`
+- DuckDB can extract parts of strings and lists using [start:end] or [start:end:step] syntax. Indexes start at 1. String slicing: `SELECT 'DuckDB'[1:4];`. Array/List slicing: `SELECT [1, 2, 3, 4][1:3];`
+- DuckDB has a powerful way to select or transform multiple columns using patterns or functions. You can select columns matching a pattern: `SELECT COLUMNS('sales_.*') FROM sales_data;` or transform multiple columns with a function: `SELECT AVG(COLUMNS('sales_.*')) FROM sales_data;`
+- DuckDB has an easy way to include/exclude or modify columns when selecting all: e.g. Exclude: `SELECT * EXCLUDE (sensitive_data) FROM users;` Replace: `SELECT * REPLACE (UPPER(name) AS name) FROM users;`
+- DuckDB has a shorthand for grouping/ordering by all non-aggregated/all columns. e.g `SELECT category, SUM(sales) FROM sales_data GROUP BY ALL;` and `SELECT * FROM my_table ORDER BY ALL;`
+- DuckDB can combine tables by matching column names, not just their positions using UNION BY NAME. E.g. `SELECT * FROM table1 UNION BY NAME SELECT * FROM table2;`
+- DuckDB has an inutitive syntax to create List/Struct/Map and Array types. Create complex types using intuitive syntax. List: `SELECT [1, 2, 3] AS my_list;`, Struct: `{{'a': 1, 'b': 'text'}} AS my_struct;`, Map: `MAP([1,2],['one','two']) as my_map;`. All types can also be nested into each other. Array types are fixed size, while list types have variable size. Compared to Structs, MAPs do not need to have the same keys present for each row, but keys can only be of type Integer or Varchar. Example: `CREATE TABLE example (my_list INTEGER[], my_struct STRUCT(a INTEGER, b TEXT), my_map MAP(INTEGER, VARCHAR),  my_array INTEGER[3], my_nested_struct STRUCT(a INTEGER, b Integer[3]));`
+- DuckDB has an inutive syntax to access struct fields using dot notation (.) or brackets ([]) with the field name. Maps fields can be accessed by brackets ([]).
+- DuckDB's way of converting between text and timestamps, and extract date parts. Current date as 'YYYY-MM-DD': `SELECT strftime(NOW(), '%Y-%m-%d');` String to timestamp: `SELECT strptime('2023-07-23', '%Y-%m-%d')::TIMESTAMP;`, Extract Year from date: `SELECT EXTRACT(YEAR FROM DATE '2023-07-23');`
+- Column Aliases in WHERE/GROUP BY/HAVING: You can use column aliases defined in the SELECT clause within the WHERE, GROUP BY, and HAVING clauses. E.g.: `SELECT a + b AS total FROM my_table WHERE total > 10 GROUP BY total HAVING total < 20;`
+- DuckDB allows generating lists using expressions similar to Python list comprehensions. E.g. `SELECT [x*2 FOR x IN [1, 2, 3]];` Returns [2, 4, 6].
+- DuckDB allows chaining multiple function calls together using the dot (.) operator. E.g.: `SELECT 'DuckDB'.replace('Duck', 'Goose').upper(); -- Returns 'GOOSEDB';`
+- DuckDB has a JSON data type. It supports selecting fields from the JSON with a JSON-Path expression using the arrow operator, -> (returns JSON) or ->> (returns text) with JSONPath expressions. For example: `SELECT data->'$.user.id' AS user_id, data->>'$.event_type' AS event_type FROM events;`
+- DuckDB has built-in functions for regex regexp_matches(column, regex), regexp_replace(column, regex), and regexp_extract(column, regex).
+- DuckDB has a way to quickly get a subset of your data with `SELECT * FROM large_table USING SAMPLE 10%;`<|eot_id|>
+<|start_header_id|>user<|end_header_id|>
+
+Database Schema:
+Here is the schema of the DuckDB database that the SQL query will run on:
+```sql
+{schema}
+```
+
+Question:
+Here is the question or an instruction the user provided:
+{question}
+
+Task:
+Write a DuckDB SQL query for the given question!<|eot_id|>
+<|start_header_id|>assistant<|end_header_id|>
+
+```sql
+"""
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n### Documentation:\n{context_str}\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        instruction = cls.PROMPT_TEMPLATE.format(
+            schema=table_text,
+            question=instruction
+        )
+        return instruction
+
+class DuckDBInstFormatterGraniteShort(RajkumarFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """System:
+Your task is to generate valid DuckDB SQL to answer the question that the user asks. You should only respond with a valid DuckDB SQL query.
+
+Here are some DuckDB SQL syntax specifics you should be aware of:
+
+
+- DuckDB use double quotes (") for identifiers that contain spaces or special characters, or to force case-sensitivity and single quotes (') to define string literals
+- DuckDB can query CSV, Parquet, and JSON directly without loading them first, e.g. `SELECT * FROM 'data.csv';`
+- DuckDB supports CREATE TABLE AS (CTAS): `CREATE TABLE new_table AS SELECT * FROM old_table;`
+- DuckDB queries can start with FROM, and optionally omit SELECT *, e.g. `FROM my_table WHERE condition;` is equivalent to `SELECT * FROM my_table WHERE condition;`
+- DuckDB allows you to use SELECT without a FROM clause to generate a single row of results or to work with expressions directly, e.g. `SELECT 1 + 1 AS result;`
+- DuckDB supports attaching multiple databases, unsing the ATTACH statement: `ATTACH 'my_database.duckdb' AS mydb;`. Tables within attached databases can be accessed using the dot notation (.), e.g. `SELECT * FROM mydb.table_name syntax`. The default databases doesn't require the do notation to access tables. The default database can be changed with the USE statement, e.g. `USE my_db;`.
+- DuckDB is generally more lenient with implicit type conversions (e.g. `SELECT '42' + 1;` - Implicit cast, result is 43), but you can always be explicit using `::`, e.g. `SELECT '42'::INTEGER + 1;`
+- DuckDB can extract parts of strings and lists using [start:end] or [start:end:step] syntax. Indexes start at 1. String slicing: `SELECT 'DuckDB'[1:4];`. Array/List slicing: `SELECT [1, 2, 3, 4][1:3];`
+- DuckDB has a powerful way to select or transform multiple columns using patterns or functions. You can select columns matching a pattern: `SELECT COLUMNS('sales_.*') FROM sales_data;` or transform multiple columns with a function: `SELECT AVG(COLUMNS('sales_.*')) FROM sales_data;`
+- DuckDB an easy way to include/exclude or modify columns when selecting all: e.g. Exclude: `SELECT * EXCLUDE (sensitive_data) FROM users;` Replace: `SELECT * REPLACE (UPPER(name) AS name) FROM users;`
+- DuckDB has a shorthand for grouping/ordering by all non-aggregated/all columns. e.g `SELECT category, SUM(sales) FROM sales_data GROUP BY ALL;` and `SELECT * FROM my_table ORDER BY ALL;`
+- DuckDB can combine tables by matching column names, not just their positions using UNION BY NAME. E.g. `SELECT * FROM table1 UNION BY NAME SELECT * FROM table2;`
+- DuckDB has an inutitive syntax to create List/Struct/Map and Array types. Create complex types using intuitive syntax. List: `SELECT [1, 2, 3] AS my_list;`, Struct: `{{'a': 1, 'b': 'text'}} AS my_struct;`, Map: `MAP([1,2],['one','two']) as my_map;`. All types can also be nested into each other. Array types are fixed size, while list types have variable size. Compared to Structs, MAPs do not need to have the same keys present for each row, but keys can only be of type Integer or Varchar. Example: `CREATE TABLE example (my_list INTEGER[], my_struct STRUCT(a INTEGER, b TEXT), my_map MAP(INTEGER, VARCHAR),  my_array INTEGER[3], my_nested_struct STRUCT(a INTEGER, b Integer[3]));`
+- DuckDB has an inutive syntax to access struct fields using dot notation (.) or brackets ([]) with the field name. Maps fields can be accessed by brackets ([]).
+- DuckDB's way of converting between text and timestamps, and extract date parts. Current date as 'YYYY-MM-DD': `SELECT strftime(NOW(), '%Y-%m-%d');` String to timestamp: `SELECT strptime('2023-07-23', '%Y-%m-%d')::TIMESTAMP;`, Extract Year from date: `SELECT EXTRACT(YEAR FROM DATE '2023-07-23');`
+- Column Aliases in WHERE/GROUP BY/HAVING: You can use column aliases defined in the SELECT clause within the WHERE, GROUP BY, and HAVING clauses. E.g.: `SELECT a + b AS total FROM my_table WHERE total > 10 GROUP BY total HAVING total < 20;`
+- DuckDB allows generating lists using expressions similar to Python list comprehensions. E.g. `SELECT [x*2 FOR x IN [1, 2, 3]];` Returns [2, 4, 6].
+- DuckDB allows chaining multiple function calls together using the dot (.) operator. E.g.: `SELECT 'DuckDB'.replace('Duck', 'Goose').upper(); -- Returns 'GOOSEDB';`
+- DuckDB has a JSON data type. It supports selecting fields from the JSON with a JSON-Path expression using the arrow operator, -> (returns JSON) or ->> (returns text) with JSONPath expressions. For example: `SELECT data->'$.user.id' AS user_id, data->>'$.event_type' AS event_type FROM events;`
+- DuckDB has built-in functions for regex regexp_matches(column, regex), regexp_replace(column, regex), and regexp_extract(column, regex).
+- DuckDB has a way to quickly get a subset of your data with `SELECT * FROM large_table USING SAMPLE 10%;`
+
+Here is the schema of the DuckDB database that the SQL query will run on:
+{schema}
+
+Question:
+Here is the question or an instruction the user provided:
+{question}
+
+Write a DuckDB SQL query for the given question!
+
+Answer:
+```
+"""
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n### Documentation:\n{context_str}\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        instruction = cls.PROMPT_TEMPLATE.format(
+            schema=table_text,
+            question=instruction
+        )
+        return instruction
+
+class DuckDBInstFormatterLlama(RajkumarFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """<|begin_of_text|>
+
+Your task is to generate valid DuckDB SQL to answer the following question, given a DuckDB database schema.
+
+## DuckDB SQL syntax specifics you should be aware of:
+
+### Case Insensitivity and Quoting:
+
+Identifiers (tables, columns): Case-insensitive, but DuckDB remembers the case you use. Use double quotes (") for identifiers that contain spaces or special characters, or to force case-sensitivity.
+```
+CREATE TABLE "My Table" ("column_name" VARCHAR); -- Spaces and mixed case
+SELECT "column_name" FROM "My Table";
+```
+
+### String Literals: Always use single quotes (') to define string literals.
+```
+SELECT 'This is a string' AS text;
+```
+
+### Direct File Querying: Query CSV, Parquet, and JSON files directly without loading them first.
+
+```
+SELECT * FROM 'data.csv';
+SELECT * FROM 'data.parquet';
+SELECT * FROM 'data.json';
+```
+
+### CREATE TABLE AS (CTAS): Create tables from query results.
+
+```
+CREATE TABLE squares AS SELECT i, i * i AS square FROM generate_series(1, 10) t(i);
+```
+
+### FROM-First Syntax (Optional SELECT): Start queries with FROM, and optionally omit SELECT *.
+
+```
+FROM my_table WHERE condition;  -- Equivalent to SELECT * FROM my_table WHERE condition
+```
+
+### SELECT without FROM: DuckDB allows you to use SELECT without a FROM clause to generate a single row of results or to work with expressions directly.
+
+```
+SELECT 1 + 1 AS result;
+```
+
+### GROUP BY ALL/ORDER BY ALL:  Shorthand for grouping/ordering by all non-aggregated/all columns.
+
+```
+SELECT category, SUM(sales) FROM sales_data GROUP BY ALL;
+SELECT * FROM my_table ORDER BY ALL;
+```
+
+### SELECT COLUMNS(): Powerful way to select or transform multiple columns using patterns or functions.
+
+```
+-- Select columns matching a pattern
+SELECT COLUMNS('sales_.*') FROM sales_data;
+
+-- Transform multiple columns with a function
+SELECT AVG(COLUMNS(*)) FROM sales_data;
+```
+
+### UNION BY NAME: Combine tables by matching column names, not just their positions.
+
+```
+SELECT * FROM table1 UNION BY NAME SELECT * FROM table2;
+```
+
+### Implicit/Explicit Casting: DuckDB is generally more lenient with implicit type conversions, but you can always be explicit using ::
+
+```
+SELECT '42' + 1;  -- Implicit cast, result is 43
+SELECT '42'::INTEGER + 1; -- Explicit cast, result is 43
+```
+
+### String/List Slicing: Extract parts of strings and lists using [start:end] or [start:end:step] syntax.
+
+```
+SELECT 'DuckDB'[1:4];  -- Returns 'Duck'
+SELECT [1, 2, 3, 4][1:3]; -- Returns [1, 2, 3]
+```
+
+### Simple List/Struct/Map/Array Creation: Create complex types using intuitive syntax.
+
+In a SELECT statement:
+```
+SELECT [1, 2, 3] AS my_list, {{'a': 1, 'b': 'text'}} AS my_struct, MAP([1,2],['one','two']) as my_map;
+```
+
+When creating a table:
+```
+CREATE TABLE data (
+    my_list INTEGER[],
+    my_struct STRUCT(a INTEGER, b TEXT),
+    my_map MAP(INTEGER, VARCHAR),
+    my_array INTEGER[3]
+);
+```
+
+### Timestamp Conversions and Extraction: Convert between text and timestamps, and extract date parts.
+
+```
+SELECT strftime(NOW(), '%Y-%m-%d');  -- Current date as 'YYYY-MM-DD'
+SELECT strptime('2023-07-23', '%Y-%m-%d')::TIMESTAMP; -- String to timestamp
+SELECT EXTRACT(YEAR FROM DATE '2023-07-23'); -- Extract year
+```
+
+### Column Aliases in WHERE/GROUP BY/HAVING: You can use column aliases defined in the SELECT clause within the WHERE, GROUP BY, and HAVING clauses.
+
+```
+SELECT a + b AS total
+FROM my_table
+WHERE total > 10
+GROUP BY total
+HAVING total < 20;
+```
+
+### List Comprehensions:  Generate lists using expressions similar to Python list comprehensions.
+
+```
+SELECT [x*2 FOR x IN [1, 2, 3]];  -- Returns [2, 4, 6]
+```
+
+### Function Chaining: Chain multiple function calls together using the dot (.) operator.
+
+```
+SELECT 'DuckDB'.replace('Duck', 'Goose').upper(); -- Returns 'GOOSEDB'
+```
+
+### Regular Expressions: DuckDB has built-in functions for regex matching, replacement, and extraction.
+
+```
+SELECT regexp_matches('DuckDB', 'Duck'); -- Returns true
+SELECT regexp_replace('DuckDB', 'Duck', 'Goose'); -- Returns 'GooseDB'
+SELECT regexp_extract('DuckDB', '(\w+)(DB)', 1); -- Returns 'Duck'
+```
+
+### Sampling: Quickly get a subset of your data with SAMPLE or TABLESAMPLE.
+
+```
+SELECT * FROM large_table USING SAMPLE 10%; -- Random 10% sample
+SELECT * FROM large_table TABLESAMPLE BERNOULLI(10); -- Bernoulli sampling
+```
+
+### ATTACH and Access: Attach external databases and reference their objects using databasename.table_name syntax.
+
+```
+ATTACH 'my_database.duckdb' AS mydb;
+SELECT * FROM mydb.my_table;
+```
+
+### SUMMARIZE: Get summary statistics (min, max, unique count, average, standard deviation, quartiles, and count) of a table.
+
+```
+SUMMARIZE table_name;
+```
+
+### DESCRIBE: Get schema of a table (column_name, column_type, null, key, default, extra).
+
+```
+DESCRIBE table_name;
+```
+
+Database Schema:
+Here is the schema of the DuckDB database that the SQL query will run on:
+{schema}
+
+Question:
+Here is the question or an instruction the user provided:
+{question}
+
+Task:
+Write a DuckDB SQL query for the given question!
+
+Here is the valid DuckDB SQL query:
+```
+"""
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n### Documentation:\n{context_str}\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        instruction = cls.PROMPT_TEMPLATE.format(
+            schema=table_text,
+            question=instruction
+        )
+        return instruction
+
+class DuckDBInstFormatterGranite(RajkumarFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """System:
+
+Your task is to generate valid DuckDB SQL to answer the following question, given a DuckDB database schema.
+
+## DuckDB SQL syntax specifics you should be aware of:
+
+### Case Insensitivity and Quoting:
+
+Identifiers (tables, columns): Case-insensitive, but DuckDB remembers the case you use. Use double quotes (") for identifiers that contain spaces or special characters, or to force case-sensitivity.
+```
+CREATE TABLE "My Table" ("column_name" VARCHAR); -- Spaces and mixed case
+SELECT "column_name" FROM "My Table";
+```
+
+### String Literals: Always use single quotes (') to define string literals.
+```
+SELECT 'This is a string' AS text;
+```
+
+### Direct File Querying: Query CSV, Parquet, and JSON files directly without loading them first.
+
+```
+SELECT * FROM 'data.csv';
+SELECT * FROM 'data.parquet';
+SELECT * FROM 'data.json';
+```
+
+### CREATE TABLE AS (CTAS): Create tables from query results.
+
+```
+CREATE TABLE squares AS SELECT i, i * i AS square FROM generate_series(1, 10) t(i);
+```
+
+### FROM-First Syntax (Optional SELECT): Start queries with FROM, and optionally omit SELECT *.
+
+```
+FROM my_table WHERE condition;  -- Equivalent to SELECT * FROM my_table WHERE condition
+```
+
+### SELECT without FROM: DuckDB allows you to use SELECT without a FROM clause to generate a single row of results or to work with expressions directly.
+
+```
+SELECT 1 + 1 AS result;
+```
+
+### GROUP BY ALL/ORDER BY ALL:  Shorthand for grouping/ordering by all non-aggregated/all columns.
+
+```
+SELECT category, SUM(sales) FROM sales_data GROUP BY ALL;
+SELECT * FROM my_table ORDER BY ALL;
+```
+
+### SELECT COLUMNS(): Powerful way to select or transform multiple columns using patterns or functions.
+
+```
+-- Select columns matching a pattern
+SELECT COLUMNS('sales_.*') FROM sales_data;
+
+-- Transform multiple columns with a function
+SELECT AVG(COLUMNS(*)) FROM sales_data;
+```
+
+### UNION BY NAME: Combine tables by matching column names, not just their positions.
+
+```
+SELECT * FROM table1 UNION BY NAME SELECT * FROM table2;
+```
+
+### Implicit/Explicit Casting: DuckDB is generally more lenient with implicit type conversions, but you can always be explicit using ::
+
+```
+SELECT '42' + 1;  -- Implicit cast, result is 43
+SELECT '42'::INTEGER + 1; -- Explicit cast, result is 43
+```
+
+### String/List Slicing: Extract parts of strings and lists using [start:end] or [start:end:step] syntax.
+
+```
+SELECT 'DuckDB'[1:4];  -- Returns 'Duck'
+SELECT [1, 2, 3, 4][1:3]; -- Returns [1, 2, 3]
+```
+
+### Simple List/Struct/Map/Array Creation: Create complex types using intuitive syntax.
+
+In a SELECT statement:
+```
+SELECT [1, 2, 3] AS my_list, {{'a': 1, 'b': 'text'}} AS my_struct, MAP([1,2],['one','two']) as my_map;
+```
+
+When creating a table:
+```
+CREATE TABLE data (
+    my_list INTEGER[],
+    my_struct STRUCT(a INTEGER, b TEXT),
+    my_map MAP(INTEGER, VARCHAR),
+    my_array INTEGER[3]
+);
+```
+
+### Timestamp Conversions and Extraction: Convert between text and timestamps, and extract date parts.
+
+```
+SELECT strftime(NOW(), '%Y-%m-%d');  -- Current date as 'YYYY-MM-DD'
+SELECT strptime('2023-07-23', '%Y-%m-%d')::TIMESTAMP; -- String to timestamp
+SELECT EXTRACT(YEAR FROM DATE '2023-07-23'); -- Extract year
+```
+
+### Column Aliases in WHERE/GROUP BY/HAVING: You can use column aliases defined in the SELECT clause within the WHERE, GROUP BY, and HAVING clauses.
+
+```
+SELECT a + b AS total
+FROM my_table
+WHERE total > 10
+GROUP BY total
+HAVING total < 20;
+```
+
+### List Comprehensions:  Generate lists using expressions similar to Python list comprehensions.
+
+```
+SELECT [x*2 FOR x IN [1, 2, 3]];  -- Returns [2, 4, 6]
+```
+
+### Function Chaining: Chain multiple function calls together using the dot (.) operator.
+
+```
+SELECT 'DuckDB'.replace('Duck', 'Goose').upper(); -- Returns 'GOOSEDB'
+```
+
+### Regular Expressions: DuckDB has built-in functions for regex matching, replacement, and extraction.
+
+```
+SELECT regexp_matches('DuckDB', 'Duck'); -- Returns true
+SELECT regexp_replace('DuckDB', 'Duck', 'Goose'); -- Returns 'GooseDB'
+SELECT regexp_extract('DuckDB', '(\w+)(DB)', 1); -- Returns 'Duck'
+```
+
+### Sampling: Quickly get a subset of your data with SAMPLE or TABLESAMPLE.
+
+```
+SELECT * FROM large_table USING SAMPLE 10%; -- Random 10% sample
+SELECT * FROM large_table TABLESAMPLE BERNOULLI(10); -- Bernoulli sampling
+```
+
+### ATTACH and Access: Attach external databases and reference their objects using databasename.table_name syntax.
+
+```
+ATTACH 'my_database.duckdb' AS mydb;
+SELECT * FROM mydb.my_table;
+```
+
+### SUMMARIZE: Get summary statistics (min, max, unique count, average, standard deviation, quartiles, and count) of a table.
+
+```
+SUMMARIZE table_name;
+```
+
+### DESCRIBE: Get schema of a table (column_name, column_type, null, key, default, extra).
+
+```
+DESCRIBE table_name;
+```
+
+Here is the schema of the DuckDB database that the SQL query will run on:
+{schema}
+
+Question:
+Here is the question or an instruction the user provided:
+{question}
+
+Please write a DuckDB SQL query that answers the user's question or instruction. Use DuckDB-specific syntax if possible.
+
+Answer:
+```
+"""
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n### Documentation:\n{context_str}\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        instruction = cls.PROMPT_TEMPLATE.format(
+            schema=table_text,
+            question=instruction
+        )
+        return instruction
+
+class DuckDBInstFormatterPhi(RajkumarFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """<|endoftext|><|user|>
+Your task is to generate valid DuckDB SQL to answer the question that the user asks. You should only respond with a valid DuckDB SQL query.
+
+Here are some DuckDB SQL syntax specifics you should be aware of:
+
+
+- DuckDB use double quotes (") for identifiers that contain spaces or special characters, or to force case-sensitivity and single quotes (') to define string literals
+- DuckDB can query CSV, Parquet, and JSON directly without loading them first, e.g. `SELECT * FROM 'data.csv';`
+- DuckDB supports CREATE TABLE AS (CTAS): `CREATE TABLE new_table AS SELECT * FROM old_table;`
+- DuckDB queries can start with FROM, and optionally omit SELECT *, e.g. `FROM my_table WHERE condition;` is equivalent to `SELECT * FROM my_table WHERE condition;`
+- DuckDB allows you to use SELECT without a FROM clause to generate a single row of results or to work with expressions directly, e.g. `SELECT 1 + 1 AS result;`
+- DuckDB supports attaching multiple databases, unsing the ATTACH statement: `ATTACH 'my_database.duckdb' AS mydb;`. Tables within attached databases can be accessed using the dot notation (.), e.g. `SELECT * FROM mydb.table_name syntax`. The default databases doesn't require the do notation to access tables. The default database can be changed with the USE statement, e.g. `USE my_db;`.
+- DuckDB is generally more lenient with implicit type conversions (e.g. `SELECT '42' + 1;` - Implicit cast, result is 43), but you can always be explicit using `::`, e.g. `SELECT '42'::INTEGER + 1;`
+- DuckDB can extract parts of strings and lists using [start:end] or [start:end:step] syntax. Indexes start at 1. String slicing: `SELECT 'DuckDB'[1:4];`. Array/List slicing: `SELECT [1, 2, 3, 4][1:3];`
+- DuckDB has a powerful way to select or transform multiple columns using patterns or functions. You can select columns matching a pattern: `SELECT COLUMNS('sales_.*') FROM sales_data;` or transform multiple columns with a function: `SELECT AVG(COLUMNS('sales_.*')) FROM sales_data;`
+- DuckDB an easy way to include/exclude or modify columns when selecting all: e.g. Exclude: `SELECT * EXCLUDE (sensitive_data) FROM users;` Replace: `SELECT * REPLACE (UPPER(name) AS name) FROM users;`
+- DuckDB has a shorthand for grouping/ordering by all non-aggregated/all columns. e.g `SELECT category, SUM(sales) FROM sales_data GROUP BY ALL;` and `SELECT * FROM my_table ORDER BY ALL;`
+- DuckDB can combine tables by matching column names, not just their positions using UNION BY NAME. E.g. `SELECT * FROM table1 UNION BY NAME SELECT * FROM table2;`
+- DuckDB has an inutitive syntax to create List/Struct/Map and Array types. Create complex types using intuitive syntax. List: `SELECT [1, 2, 3] AS my_list;`, Struct: `{{'a': 1, 'b': 'text'}} AS my_struct;`, Map: `MAP([1,2],['one','two']) as my_map;`. All types can also be nested into each other. Array types are fixed size, while list types have variable size. Compared to Structs, MAPs do not need to have the same keys present for each row, but keys can only be of type Integer or Varchar. Example: `CREATE TABLE example (my_list INTEGER[], my_struct STRUCT(a INTEGER, b TEXT), my_map MAP(INTEGER, VARCHAR),  my_array INTEGER[3], my_nested_struct STRUCT(a INTEGER, b Integer[3]));`
+- DuckDB has an inutive syntax to access struct fields using dot notation (.) or brackets ([]) with the field name. Maps fields can be accessed by brackets ([]).
+- DuckDB's way of converting between text and timestamps, and extract date parts. Current date as 'YYYY-MM-DD': `SELECT strftime(NOW(), '%Y-%m-%d');` String to timestamp: `SELECT strptime('2023-07-23', '%Y-%m-%d')::TIMESTAMP;`, Extract Year from date: `SELECT EXTRACT(YEAR FROM DATE '2023-07-23');`
+- Column Aliases in WHERE/GROUP BY/HAVING: You can use column aliases defined in the SELECT clause within the WHERE, GROUP BY, and HAVING clauses. E.g.: `SELECT a + b AS total FROM my_table WHERE total > 10 GROUP BY total HAVING total < 20;`
+- DuckDB allows generating lists using expressions similar to Python list comprehensions. E.g. `SELECT [x*2 FOR x IN [1, 2, 3]];` Returns [2, 4, 6].
+- DuckDB allows chaining multiple function calls together using the dot (.) operator. E.g.: `SELECT 'DuckDB'.replace('Duck', 'Goose').upper(); -- Returns 'GOOSEDB';`
+- DuckDB has a JSON data type. It supports selecting fields from the JSON with a JSON-Path expression using the arrow operator, -> (returns JSON) or ->> (returns text) with JSONPath expressions. For example: `SELECT data->'$.user.id' AS user_id, data->>'$.event_type' AS event_type FROM events;`
+- DuckDB has built-in functions for regex regexp_matches(column, regex), regexp_replace(column, regex), and regexp_extract(column, regex).
+- DuckDB has a way to quickly get a subset of your data with `SELECT * FROM large_table USING SAMPLE 10%;`
+
+Here is the schema of the DuckDB database that the SQL query will run on:
+{schema}
+
+Question:
+Here is the question or an instruction the user provided:
+{question}
+
+Write a DuckDB SQL query for the given question!<|end|>
+<|assistant|>
+```sql
+"""
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n### Documentation:\n{context_str}\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        instruction = cls.PROMPT_TEMPLATE.format(
+            schema=table_text,
+            question=instruction
+        )
+        return instruction
+
+class DuckDBInstFormatterGPTmini(RajkumarFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """Schema:
+```sql
+{schema}
+```
+
+Question:
+{question}
+
+Write a valid DuckDB SQL query to answer the question!
+"""
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n### Documentation:\n{context_str}\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        instruction = cls.PROMPT_TEMPLATE.format(
+            schema=table_text,
+            question=instruction
+        )
+        return instruction
+
+class DuckDBInstFormatterPhiAzure(RajkumarFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """Your task is to generate valid DuckDB SQL to answer the question that the user asks. You should only respond with a valid DuckDB SQL query.
+
+Here is the schema of the DuckDB database that the SQL query will run on:
+{schema}
+
+Question:
+Here is the question or an instruction the user provided:
+{question}
+
+Write a DuckDB SQL query for the given question!
+"""
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n### Documentation:\n{context_str}\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: str,
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        instruction = cls.PROMPT_TEMPLATE.format(
+            schema=table_text,
+            question=instruction
+        )
+        return instruction
+
+class DuckDBInstNoShorthandFormatter(DuckDBInstFormatter):
+    """DuckDB Inst class."""
+
+    PROMPT_TEMPLATE = """### Instruction:\n{instruction}\n\n### Input:\n{input}{context}\n### Question:\n{question}\n\n### Response:\n"""
+    INSTRUCTION_TEMPLATE = """Your task is to generate valid duckdb SQL to answer the following question{has_schema}"""  # noqa: E501
+
+
+class DuckDBChat:
+    """DuckDB Inst class."""
+
+    table_sep: str = "\n\n"
+    shuffle_table_order: bool = True
+    _cache: dict[tuple[str, str, str], list[str]] = {}
+    clean_whitespace = False
+    model = None
+
+    @classmethod
+    def format_table(cls, table: Table) -> str:
+        """Get table format."""
+        table_fmt = []
+        for col in table.columns or []:
+            # This is technically an incorrect type, but it should be a catchall word
+            table_fmt.append(f"    {col.name} {col.dtype or 'any'}")
+        if table_fmt:
+            all_cols = ",\n".join(table_fmt)
+            create_tbl = f"CREATE TABLE {table.name} (\n{all_cols}\n)"
+        else:
+            create_tbl = f"CREATE TABLE {table.name}"
+        return create_tbl
+
+    @classmethod
+    def format_all_tables(cls, tables: list[Table], instruction: str) -> list[dict]:
+        """Get all tables format."""
+        if not cls.model:
+            cls.model = Manifest(
+                engine="gpt-3.5-turbo",
+                client_name="openaichat",
+                cache_name="sqlite",
+                cache_connection=".manifest.sqlite",
+            )
+        table_texts = [cls.format_table(table) for table in tables]
+        full_schema = cls.table_sep.join(table_texts)
+        prompt = f"""SQL schema of my database:
+{full_schema}
+Explain in a few sentences what the data is about:
+        """
+        messages = [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant that can generate an human redable summary of database content based on the schema.",
+            },
+            {"role": "user", "content": prompt},
+        ]
+        explanation = cls.model.run(messages, temperature=0)
+        messages.append({"role": "assistant", "content": explanation})
+        return messages[1:]
+
+    @classmethod
+    def format_retrieved_context(
+        cls,
+        context: list[str],
+    ) -> str:
+        """Format retrieved context."""
+        context_str = "\n--------\n".join(context)
+        return f"\n\nHere is additional documentation about DuckDB that could be useful.\n--------\n{context_str}\n--------\n"
+
+    @classmethod
+    def format_prompt(
+        cls,
+        instruction: str,
+        table_text: list[dict],
+        context_text: str,
+    ) -> str | list[str]:
+        """Get prompt format."""
+        prompt = f"""Now output a single SQL query without any explanation and do not add anything
+to the query that was not part of the question, also do not use markdown. Make sure to only
+use information provided in the prompt, or tables and columns from the schema above and write a query to answer the question.{context_text}\n\nMy quesiton is \n`{instruction}`\n\nGenerate the DuckDB specific SQL query:"""  # noqa: E501
+        messages = [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant that can generate DuckDB sql queries, which is a superset of Postgresql, based on the user input. You do not respond with any human readable text, only SQL code.",
+            },
+            *table_text,
+            {"role": "user", "content": prompt},
+        ]
+        return messages
+
+    @classmethod
+    def format_model_output(cls, output_sql: str, prompt: str) -> str:
+        """Format model output."""
+        return output_sql
+
+    @classmethod
+    def format_gold_output(cls, output_sql: str) -> str:
+        """Format gold output for demonstration."""
+        return output_sql

duckdb-nsql/eval/schema.py

@@ -0,0 +1,115 @@
+"""Text2SQL schemas."""
+import enum
+
+from manifest.response import Usage
+from pydantic import BaseModel
+
+DEFAULT_TABLE_NAME: str = "db_table"
+
+
+class Dialect(str, enum.Enum):
+    """SQGFluff and SQLGlot dialects.
+
+    Lucky for us, the dialects match both parsers.
+
+    Ref: https://github.com/sqlfluff/sqlfluff/blob/main/src/sqlfluff/core/dialects/__init__.py  # noqa: E501
+    Ref: https://github.com/tobymao/sqlglot/blob/main/sqlglot/dialects/__init__.py  # noqa: E501
+    """
+
+    SNOWFLAKE = "snowflake"
+    BIGQUERY = "bigquery"
+    REDSHIFT = "redshift"
+    POSTGRES = "postgres"
+    UNKNOWN = "unknown"
+
+    @property
+    def dialect_str(self) -> str | None:
+        """Get the dialect string for validation.
+
+        We need to pass in dialect = None for UNKNOWN dialects.
+        """
+        if self != Dialect.UNKNOWN:
+            return self.value
+        else:
+            return None
+
+    @property
+    def quote_str(self) -> str:
+        """Get the quote string for the dialect."""
+        if self == Dialect.SNOWFLAKE:
+            return '"'
+        elif self == Dialect.BIGQUERY:
+            return "`"
+        elif self == Dialect.REDSHIFT:
+            return '"'
+        elif self == Dialect.POSTGRES:
+            return '"'
+        elif self == Dialect.UNKNOWN:
+            return '"'
+        raise NotImplementedError(f"Quote string not implemented for dialect {self}")
+
+    def quote(self, string: str) -> str:
+        """Quote a string."""
+        return f"{self.quote_str}{string}{self.quote_str}"
+
+
+class ColumnOrLiteral(BaseModel):
+    """Column that may or may not be a literal."""
+
+    name: str | None = None
+    literal: bool = False
+
+    def __hash__(self) -> int:
+        """Hash."""
+        return hash((self.name, self.literal))
+
+
+class TableColumn(BaseModel):
+    """Table column."""
+
+    name: str
+    dtype: str | None
+
+
+class ForeignKey(BaseModel):
+    """Foreign key."""
+
+    # Referenced column
+    column: TableColumn
+    # References table name
+    references_name: str
+    # References column
+    references_column: TableColumn
+
+
+class Table(BaseModel):
+    """Table."""
+
+    name: str | None
+    columns: list[TableColumn] | None
+    pks: list[TableColumn] | None
+    # FK from this table to another column in another table
+    fks: list[ForeignKey] | None
+    examples: list[dict] | None
+    # Is the table a source or intermediate reference table
+    is_reference_table: bool = False
+
+
+class TextToSQLParams(BaseModel):
+    """A text to sql request."""
+
+    instruction: str
+    database: str | None
+    # Default to unknown
+    dialect: Dialect = Dialect.UNKNOWN
+    tables: list[Table] | None
+
+
+class TextToSQLModelResponse(BaseModel):
+    """Model for Autocomplete Responses."""
+
+    output: str
+    final_prompt: str | list[dict]
+    raw_output: str
+    usage: Usage
+    metadata: str | None = None

duckdb-nsql/eval/text_to_sql.py

@@ -0,0 +1,260 @@
+"""Text-to-SQL running."""
+import asyncio
+import json
+import re
+import time
+from typing import cast
+import duckdb
+
+import structlog
+from manifest import Manifest
+from manifest.response import Response, Usage
+from prompt_formatters import RajkumarFormatter, MotherDuckFormatter
+from schema import DEFAULT_TABLE_NAME, TextToSQLModelResponse, TextToSQLParams
+from tqdm.auto import tqdm
+
+logger = structlog.get_logger()
+
+
+def clean_whitespace(sql: str) -> str:
+    """Clean whitespace."""
+    return re.sub(r"[\t\n\s]+", " ", sql)
+
+
+def instruction_to_sql(
+    params: TextToSQLParams,
+    extra_context: list[str],
+    manifest: Manifest,
+    prompt_formatter: RajkumarFormatter = None,
+    overwrite_manifest: bool = False,
+    max_tokens: int = 300,
+    temperature: float = 0.1,
+    stop_sequences: list[str] | None = None,
+    num_beams: int = 1,
+) -> TextToSQLModelResponse:
+    """Parse the instruction to a sql command."""
+    return instruction_to_sql_list(
+        params=[params],
+        extra_context=[extra_context],
+        manifest=manifest,
+        prompt_formatter=prompt_formatter,
+        overwrite_manifest=overwrite_manifest,
+        max_tokens=max_tokens,
+        temperature=0.1,
+        stop_sequences=stop_sequences,
+        num_beams=num_beams,
+    )[0]
+
+def run_motherduck_prompt_sql(params: list[TextToSQLParams]) -> list[TextToSQLModelResponse]:
+    results = []
+    for param in params:
+        con = duckdb.connect('md:')
+        try:
+            sql_query = con.execute("CALL prompt_sql(?);", [param.instruction]).fetchall()[0][0]
+        except Exception as e:
+            print(e)
+            sql_query = "SELECT * FROM hn.hacker_news LIMIT 1";
+        usage = Usage(
+                completion_tokens = 0,
+                prompt_tokens = 0,
+                total_tokens = 0
+        )
+        model_response = TextToSQLModelResponse(
+            output=sql_query,
+            raw_output=sql_query,
+            final_prompt=param.instruction,
+            usage=usage,
+        )
+        results.append(model_response)
+    return results
+
+
+
+def instruction_to_sql_list(
+    params: list[TextToSQLParams],
+    extra_context: list[list[str]],
+    manifest: Manifest,
+    prompt_formatter: RajkumarFormatter = None,
+    overwrite_manifest: bool = False,
+    max_tokens: int = 300,
+    temperature: float = 0.1,
+    stop_sequences: list[str] | None = None,
+    num_beams: int = 1,
+    verbose: bool = False,
+) -> list[TextToSQLModelResponse]:
+    """Parse the list of instructions to sql commands.
+
+    Connector is used for default retry handlers only.
+    """
+    if type(prompt_formatter) is MotherDuckFormatter:
+        return run_motherduck_prompt_sql(params)
+
+    if prompt_formatter is None:
+        raise ValueError("Prompt formatter is required.")
+
+    def construct_params(
+        params: TextToSQLParams,
+        context: list[str],
+    ) -> str | list[dict]:
+        """Turn params into prompt."""
+        if prompt_formatter.clean_whitespace:
+            instruction = clean_whitespace(params.instruction)
+        else:
+            instruction = params.instruction
+
+        table_texts = prompt_formatter.format_all_tables(
+            params.tables, instruction=instruction
+        )
+        # table_texts can be list of chat messages. Only join list of str.
+        if table_texts:
+            if isinstance(table_texts[0], str):
+                table_text = prompt_formatter.table_sep.join(table_texts)
+            else:
+                table_text = table_texts
+        else:
+            table_text = ""
+
+        if context:
+            context_text = prompt_formatter.format_retrieved_context(context)
+        else:
+            context_text = "" if isinstance(table_text, str) else []
+        prompt = prompt_formatter.format_prompt(
+            instruction,
+            table_text,
+            context_text,
+        )
+        return prompt
+
+    # If no inputs, return nothing
+    if not params:
+        return []
+
+    # Stitch together demonstrations and params
+    prompts: list[str | list[dict]] = []
+    for i, param in tqdm(
+        enumerate(params),
+        total=len(params),
+        desc="Constructing prompts",
+        disable=not verbose,
+    ):
+        predict_str = construct_params(param, extra_context[i] if extra_context else [])
+        if isinstance(predict_str, str):
+            prompt = predict_str.lstrip()
+        else:
+            prompt = predict_str
+        prompts.append(prompt)
+
+    manifest_params = dict(
+        max_tokens=max_tokens,
+        overwrite_cache=overwrite_manifest,
+        num_beams=num_beams,
+        logprobs=5,
+        temperature=0.1,
+        do_sample=False if 0.1 <= 0 else True,
+        stop_sequences=stop_sequences or prompt_formatter.stop_sequences,
+    )
+
+    ret: list[TextToSQLModelResponse] = []
+    if len(params) == 1:
+        prompt = prompts[0]
+        success = False
+        retries = 0
+        while not success and retries < 5:
+            try:
+                model_response = _run_manifest(
+                    prompt,
+                    manifest_params,
+                    prompt_formatter,
+                    manifest,
+                    stop_sequences=stop_sequences,
+                )
+                success = True
+            except:
+                retries +=1
+
+        usage = model_response.usage
+        model_response.usage = usage
+        ret.append(model_response)
+    else:
+        # We do not handle retry logic on parallel requests right now
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        response = cast(
+            Response,
+            loop.run_until_complete(
+                manifest.arun_batch(
+                    prompts,
+                    **manifest_params,  # type: ignore
+                ),
+            ),
+        )
+        loop.close()
+
+        response_usage = response.get_usage()
+        response_text = response.get_parsed_response()
+        for prompt, resp in zip(prompts, response_text):
+            # This will restitch the query in the case we force it to start with SELECT
+            sql_query = prompt_formatter.format_model_output(cast(str, resp), prompt)
+            for token in stop_sequences:
+                sql_query = sql_query.split(token)[0]
+            logger.info(f"FINAL OUTPUT: {sql_query}")
+            ret.append(
+                TextToSQLModelResponse(
+                    output=sql_query,
+                    raw_output=cast(str, resp),
+                    final_prompt=prompt,
+                    usage=response_usage,
+                )
+            )
+
+    return ret
+
+
+def _run_manifest(
+    prompt: str | list[str],
+    manifest_params: dict,
+    prompt_formatter: RajkumarFormatter,
+    manifest: Manifest,
+    stop_sequences: list[str] | None = None,
+) -> TextToSQLModelResponse:
+    """Run manifest for prompt format."""
+    logger.info(f"PARAMS: {manifest_params}")
+    if isinstance(prompt, list):
+        for p in prompt:
+            logger.info(f"PROMPT: {p['role']}: {p['content']}")
+    else:
+        logger.info(f"PROMPT: {prompt}")
+    start_time = time.time()
+    # Run result
+    response = cast(
+        Response,
+        manifest.run(
+            prompt,
+            return_response=True,
+            client_timeout=1800,
+            **manifest_params,  # type: ignore
+        ),
+    )
+    logger.info(f"TIME: {time.time() - start_time: .2f}")
+
+    response_usage = response.get_usage_obj()
+    summed_usage = Usage()
+    for usage in response_usage.usages:
+        summed_usage.completion_tokens += usage.completion_tokens
+        summed_usage.prompt_tokens += usage.prompt_tokens
+        summed_usage.total_tokens += usage.total_tokens
+    # This will restitch the query in the case we force it to start with SELECT
+    sql_query = prompt_formatter.format_model_output(
+        cast(str, response.get_response()), prompt
+    )
+
+    for token in stop_sequences:
+        sql_query = sql_query.split(token)[0]
+    logger.info(f"OUTPUT: {sql_query}")
+    model_response = TextToSQLModelResponse(
+        output=sql_query,
+        raw_output=cast(str, response.get_response()),
+        final_prompt=prompt,
+        usage=summed_usage,
+    )
+    return model_response

duckdb-nsql/manifest/.flake8

@@ -0,0 +1,11 @@
+# This is our code-style check. We currently allow the following exceptions:
+#   - E731: do not assign a lambda expression, use a def
+#   - E402: module level import not at top of file
+#   - W503: line break before binary operator
+#   - E203: whitespace before :
+
+[flake8]
+exclude = .git
+max-line-length = 88
+ignore = E731, E402, W503, E203, PAI100, PAI101, PAI201, PAI202, PAI203
+per-file-ignores = __init__.py:F401, version.py:D100

duckdb-nsql/manifest/.pre-commit-config.yaml

@@ -0,0 +1,23 @@
+repos:
+    - repo: https://github.com/pre-commit/pre-commit-hooks
+      rev: v3.2.0
+      hooks:
+          - id: trailing-whitespace
+          - id: end-of-file-fixer
+          - id: check-yaml
+          - id: check-toml
+          - id: check-merge-conflict
+          - id: check-added-large-files
+    - repo: https://github.com/timothycrosley/isort
+      rev: 5.13.2
+      hooks:
+          - id: isort
+    - repo: https://github.com/psf/black
+      rev: 22.3.0
+      hooks:
+          - id: black
+            language_version: python3
+    - repo: https://github.com/PyCQA/flake8
+      rev: 6.0.0
+      hooks:
+          - id: flake8

duckdb-nsql/manifest/CHANGELOG.rst

@@ -0,0 +1,93 @@
+0.1.10 - Unreleased
+---------------------
+
+0.1.9 - 2024-01-22
+---------------------
+Fixed
+^^^^^
+* Added trust code params HF models
+* Added LRU cache to HF model param calls to avoid extra calls
+* Fixed pydantic type issue HF model return
+* Support for Python 3.10-3.11
+
+0.1.8 - 2023-05-22
+---------------------
+Added
+^^^^^
+* Azure model support (completion and chat)
+* Google Vertex API model support (completion and chat)
+* Streaming responses for LM Completions (set stream=True)
+
+Fixed
+^^^^^
+* `run` with batches now acts the same as async run except not async. We will batch requests into appropriate batchs sizes.
+* Refactored client so unified preprocess and postprocess of requests and responses to better support model variants in request/response format.
+
+0.1.7 - 2023-05-17
+---------------------
+Fixed
+^^^^^
+* `_run_chat` fixed bug where not passing in kwargs
+
+0.1.6 - 2023-05-16
+---------------------
+Fixed
+^^^^^
+* Unified `run` and `run_chat` methods so it's just `run` now.
+* LLama HF models for eval
+
+0.1.5 - 2023-05-03
+---------------------
+Added
+^^^^^
+* Added chat input for chat models.
+
+0.1.4 - 2023-04-24
+---------------------
+Added
+^^^^^
+* Connection pools to swap between clients
+* Chunksize param for async runs
+
+Fixed
+^^^^^
+* Determine cache and response by request type, not client name
+* Refactor Response to use Pydantic types for Request and Response
+
+0.1.1
+---------------------
+Added
+^^^^^
+* Async support in arun_batch
+
+Fixed
+^^^^^
+* Batched runs now caches individual items
+* Score prompt does not truncate outside token
+
+Removed
+^^^^^
+* Deprecated chatGPT in favor of openaichat which uses OpenAI completions
+* Deprecated Sessions
+
+0.1.0 - 2022-01-31
+---------------------
+Added
+^^^^^
+* Batched inference support in `manifest.run`. No more separate `manifest.run_batch` method.
+* Standard request base model for all language inputs.
+* ChatGPT client. Requires CHATGPT_SESSION_KEY to be passed in.
+* Diffusion model support
+* Together model support
+
+Removed
+^^^^^^^
+* `Prompt` class
+* `OPT` client - OPT is now available in HuggingFace
+
+0.0.1 - 2022-11-08
+-------------------
+First major pip release of Manifest. Install via `pip install manifest-ml`.
+
+
+.. _@lorr1: https://github.com/lorr1

duckdb-nsql/manifest/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

duckdb-nsql/manifest/Makefile

@@ -0,0 +1,27 @@
+dev:
+	pip install -e .[all]
+	pre-commit install
+
+test: dev check
+	pytest tests
+
+format:
+	isort --atomic manifest/ tests/ web_app/
+	black manifest/ tests/ web_app/
+
+check:
+	isort -c manifest/ tests/ web_app/
+	black manifest/ tests/ web_app/ --check
+	flake8 manifest/ tests/ web_app/
+	mypy manifest/ tests/ web_app/
+
+clean:
+	pip uninstall -y manifest
+	rm -rf src/manifest.egg-info
+	rm -rf build/ dist/
+
+prune:
+	@bash -c "git fetch -p";
+	@bash -c "for branch in $(git branch -vv | grep ': gone]' | awk '{print $1}'); do git branch -d $branch; done";
+
+.PHONY: dev test clean check prune

duckdb-nsql/manifest/README.md

@@ -0,0 +1,304 @@
+# Manifest
+How to make prompt programming with Foundation Models a little easier.
+
+
+# Table of Contents
+- [Install](#install)
+- [Getting Started](#getting-started)
+- [Manifest](#manifest-components)
+- [Other Models Types](#other-models)
+    - [Local HuggingFace Models](#local-huggingface-models)
+    - [Chat Models](#chat-models)
+    - [Embedding Models](#embedding-models)
+- [Road Map](#road-map)
+- [Development](#development)
+- [Cite](#cite)
+
+
+# Install
+Install:
+```bash
+pip install manifest-ml
+```
+
+Install with diffusion support:
+```bash
+pip install manifest-ml[diffusers]
+```
+
+Install with HuggingFace local model support:
+```bash
+pip install manifest-ml[api]
+```
+
+Dev Install:
+```bash
+git clone [email protected]:HazyResearch/manifest.git
+cd manifest
+make dev
+```
+
+# Getting Started
+Running is simple to get started. If using OpenAI, set `export OPENAI_API_KEY=<OPENAIKEY>` (or pass key in through variable `client_connection`) then run
+
+```python
+from manifest import Manifest
+
+# Start a manifest session to OpenAI - default `engine=text-davinci-003`
+manifest = Manifest(
+    client_name = "openai",
+)
+manifest.run("Why is the grass green?")
+```
+
+## Examples
+We have example notebook and python scripts located at [examples](examples). These show how to use different models, model types (i.e. text, diffusers, or embedding models), and async running.
+
+# Manifest Components
+Manifest is meant to be a very light weight package to help with prompt design and iteration. Three key design decisions of Manifest are
+
+* All models are behind APIs
+* Supports caching of model inputs/outputs for iteration, reproducibility, and cost saving
+* Unified API to support generate, score, and embed
+
+## Models
+Manifest provides model clients for [OpenAI](https://openai.com/), [AI21](https://studio.ai21.com/), [Cohere](https://cohere.ai/), [Together](https://together.xyz/), and HuggingFace (see [below](#huggingface-models) for how to use locally hosted HuggingFace models). You can toggle between the models by changing `client_name` and `client_connection`. For example, if a HuggingFace model is loaded locally, run
+```python
+manifest = Manifest(
+    client_name = "huggingface",
+    client_connection = "http://127.0.0.1:5000",
+)
+```
+If you want to use Cohere, run
+```python
+manifest = Manifest(
+    client_name = "cohere",
+    client_connection = <COHERE_API_KEY>,
+)
+```
+You can also just set `export COHERE_API_KEY=<COHERE_API_KEY>` and not use `client_connection`.
+
+If you want to use AI21 Labs, run
+```python
+manifest = Manifest(
+    client_name = "ai21",
+    client_connection = <AI21_API_KEY>,
+)
+```
+
+You can see the model details and possible model inputs to `run()` via
+```python
+print(manifest.client_pool.get_current_client().get_model_params())
+print(manifest.client_pool.get_current_client().get_model_inputs())
+```
+
+## Global Cache
+We support having queries and results stored in a global cache that can be shared across users. We treat inputs and outputs as key value pairs and support SQLite or Redis backends. To start with global caching using SQLite, run
+
+```python
+manifest = Manifest(
+    client_name = "openai",
+    cache_name = "sqlite",
+    cache_connection = "mycache.sqlite",
+)
+```
+The cache will be saved in `mycache.sqlite`.
+
+We also support Redis backend.
+```python
+manifest = Manifest(
+    client_name = "openai",
+    cache_name = "redis",
+    cache_connection = "localhost:6379"
+)
+```
+As a hint, if you want to get Redis running, see the `docker run` command below under development.
+
+## Running Queries
+Once you have a session open, you can write and develop prompts.
+
+```python
+result = manifest.run("Hello, my name is Laurel")
+```
+
+You can also run over multiple examples if supported by the client.
+```python
+results = manifest.run(["Where are the cats?", "Where are the dogs?"])
+```
+
+We support async queries as well via
+```python
+import asyncio
+results = asyncio.run(manifest.arun_batch(["Where are the cats?", "Where are the dogs?"]))
+```
+
+If something doesn't go right, you can also ask to get a raw manifest Response.
+```python
+result_object = manifest.run(["Where are the cats?", "Where are the dogs?"], return_response=True)
+print(result_object.get_request_obj())
+print(result_object.is_cached())
+print(result_object.get_response_obj())
+```
+
+By default, we do not truncate results based on a stop token. You can change this by either passing a new stop token to a Manifest session or to a `run`.
+```python
+result = manifest.run(prompt, "Laurel", stop_token="and")
+```
+
+If you want to change default parameters to a model, we pass those as `kwargs` to the client.
+```python
+result = manifest.run(prompt, "Laurel", max_tokens=50)
+```
+
+## Streaming Queries
+Manifest also supports streaming the model response back, assuming it's supported by the underlying client. When calling `run`, pass `stream=True` to get a streaming iterator in response.
+
+```python
+result_iterator = manifest.run("Tell me a story. Once upon a time", max_tokens=100, stream=True)
+for res_text in result_iterator:
+    print(res_text)
+```
+Streaming responses are only supported for single string queries (not batch mode) for text completion models.
+
+## Model Pools
+Manifest supports querying multiple models with different schedulers. This is very much a work in progress effort, but Manifest will round robin select (or randomly select) the clients you want. You can use the same client multiple times with different connection strings (e.g. different API keys), or you can mix and match. The only requirement is that all clients are the same request type. I.e. you can't have a pool of generation models and embedding models.
+
+To query between a local model and OpenAI,
+```python
+from manifest.connections.client_pool import ClientConnection
+from manifest import Manifest
+
+client_connection1 = ClientConnection(
+    client_name="huggingface",
+    client_connection="http://127.0.0.1:5000",
+)
+client_connection2 = ClientConnection(client_name="openai", engine="text-ada-001")
+manifest = Manifest(
+    client_pool=[client_connection1, client_connection2],
+    cache_name="sqlite",
+    client_connection=sqlite_cache,
+)
+manifest.run(...)
+```
+
+The speed benefit comes in with async batched runs. When calling `arun_batch` with a list of prompts, Manifest supports a `chunk_size` param. This will break the prompts into `chunk_size` chunks to spread across the client pool. By default `chunk_size` is `-1` which means only one client will get all the prompts to run asynchronously. You must set `chunk_size > 1` to distribute across the pool. There is a further `batch_size` param which control the individual client `batch_size` to send to the model.
+
+```python
+responses = asyncio.run(manifest.arun_batch(prompts, max_tokens=30, chunk_size=20))
+```
+
+# Other Models
+
+## Local Huggingface Models
+To use a HuggingFace generative model, in `manifest/api` we have a Flask application that hosts the models for you.
+
+In a separate terminal or Tmux/Screen session, to load 6B parameters models, run
+```bash
+python3 -m manifest.api.app \
+    --model_type huggingface \
+    --model_name_or_path EleutherAI/gpt-j-6B \
+    --device 0
+```
+You will see the Flask session start and output a URL `http://127.0.0.1:5000`. Pass this in to Manifest. If you want to use a different port, set the `FLASK_PORT` environment variable.
+
+```python
+manifest = Manifest(
+    client_name = "huggingface",
+    client_connection = "http://127.0.0.1:5000",
+)
+```
+
+If you have a custom model you trained, pass the model path to `--model_name_or_path`.
+
+To help load larger models, we also support using `parallelize()` from HF, [accelerate](https://huggingface.co/docs/accelerate/index), [bitsandbytes](https://github.com/TimDettmers/bitsandbytes), and [deepspeed](https://github.com/microsoft/DeepSpeed). You will need to install these packages first via `pip install manifest-ml[api]`. We list the commands to load larger models below.
+
+* T0pp
+```bash
+python3 -m manifest.api.app \
+    --model_type huggingface \
+    --model_name_or_path bigscience/T0pp \
+    --use_hf_parallelize
+```
+
+* NeoX 20B (requires at least 60GB of GPU memory)
+```bash
+python3 -m manifest.api.app \
+    --model_type huggingface \
+    --model_name_or_path EleutherAI/gpt-neox-20b \
+    --use_accelerate_multigpu \
+    --percent_max_gpu_mem_reduction 0.75
+```
+* Bloom 175B (requires at least 240GB of GPU memory)
+```bash
+python3 -m manifest.api.app \
+    --model_type huggingface \
+    --model_name_or_path bigscience/bloom \
+    --use_bitsandbytes \
+    --percent_max_gpu_mem_reduction 0.85
+```
+
+## Chat Models
+Manifest has specific support for executing against chat models in the more standard "system" / "user" dialogue. To pass in a dialogue history to Manifest, use the `run` command with a list of dictionary inputs with `role` and `content` keys using an associated chat model such as `openaichat`.
+
+```python
+manifest = Manifest(client_name="openaichat")
+dialogue = [
+    {"role": "system", "content": "You are a helpful assistant who also responds in rhymes"},
+    {"role": "user", "content": "What is the date?"},
+]
+res = manifest.run(dialogue, max_tokens=100)
+```
+
+## Embedding Models
+Manifest also supports getting embeddings from models and available APIs. We do this all through changing the `client_name` argument. You still use `run` and `abatch_run`.
+
+To use OpenAI's embedding models, simply run
+```python
+manifest = Manifest(client_name="openaiembedding")
+embedding_as_np = manifest.run("Get me an embedding for a bunny")
+```
+
+As explained above, you can load local HuggingFace models that give you embeddings, too. If you want to use a standard generative model, load the model as above use use `client_name="huggingfaceembedding"`. If you want to use a standard embedding model, like those from SentenceTransformers, load your local model via
+```bash
+python3 -m manifest.api.app \
+    --model_type sentence_transformers \
+    --model_name_or_path all-mpnet-base-v2 \
+    --device 0
+```
+
+# Road Map
+Here's what's coming up next
+- [ ] Clients
+  - [ ] HuggingFace Hub
+  - [x] Azure OpenAI
+  - [x] Google Vertex
+  - [ ] Anthropic
+  - [x] Streaming Support Completions
+  - [ ] Streaming Support Chat Models
+- [ ] Data Types
+  - [ ] Diffusion Models
+- [x] Orchestration
+  - [x] Connection pools
+- [ ] Local Inference
+  - [ ] FlexGen
+
+# Development
+Before submitting a PR, run
+```bash
+export REDIS_PORT="6379"  # or whatever PORT local redis is running for those tests
+cd <REDIS_PATH>
+docker run -d -p 127.0.0.1:${REDIS_PORT}:6379 -v `pwd`:`pwd` -w `pwd` --name manifest_redis_test redis
+make test
+```
+
+# Cite
+Please cite Manifest if you used it for any publications. Thanks!!
+```
+@misc{orr2022manifest,
+  author = {Orr, Laurel},
+  title = {Manifest},
+  year = {2022},
+  publisher = {GitHub},
+  howpublished = {\url{https://github.com/HazyResearch/manifest}},
+}
+```

duckdb-nsql/manifest/examples/langchain_chatgpt.ipynb

@@ -0,0 +1,455 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "b253f4d5",
+   "metadata": {},
+   "source": [
+    "# ChatGPT Clone using TOMA GPT-JT-6B\n",
+    "(adopted from ChatGPT Clone [notebook](https://github.com/hwchase17/langchain/blob/master/docs/examples/chains/chatgpt_clone.ipynb))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b0302886",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "env: TOMA_URL=https://staging.together.xyz/api\n"
+     ]
+    }
+   ],
+   "source": [
+    "%env TOMA_URL=https://staging.together.xyz/api"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "93a18ea6",
+   "metadata": {},
+   "source": [
+    "Make sure you have langchain installed and manifest. For the most recent versions, run\n",
+    "```\n",
+    "pip install git+https://github.com/hwchase17/langchain.git\n",
+    "pip install git+https://github.com/HazyResearch/manifest.git\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "id": "a99acd89",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
+      "Prompt after formatting:\n",
+      "\u001b[32;1m\u001b[1;3mI am a classification model. It will try to classify your input.\n",
+      "\n",
+      "\n",
+      "Input: Classes are \"positive\" and \"negative\". For example given\n",
+      "Input: I love this product!\n",
+      "Output: positive.\n",
+      "I think this movie was one of the worst of the year. Script was boring!\n",
+      "Output:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished LLMChain chain.\u001b[0m\n",
+      "negative.\n"
+     ]
+    }
+   ],
+   "source": [
+    "from manifest import Manifest\n",
+    "from langchain.llms.manifest import ManifestWrapper\n",
+    "from langchain import ConversationChain, LLMChain, PromptTemplate\n",
+    "from langchain.chains.conversation.memory import ConversationalBufferWindowMemory\n",
+    "\n",
+    "\n",
+    "template = \"\"\"I am a classification model. It will try to classify your input.\n",
+    "\n",
+    "{history}\n",
+    "Input: {human_input}\n",
+    "Output:\"\"\"\n",
+    "\n",
+    "prompt = PromptTemplate(\n",
+    "    input_variables=[\"history\", \"human_input\"], \n",
+    "    template=template\n",
+    ")\n",
+    "\n",
+    "manifest = Manifest(\n",
+    "    client_name=\"toma\",\n",
+    "    engine=\"Together-gpt-JT-6B-v1\",\n",
+    "    max_tokens=150,\n",
+    "    top_p=0.9,\n",
+    "    top_k=40,\n",
+    "    stop_sequences=[\"\\n\"],\n",
+    ")\n",
+    "\n",
+    "chatgpt_chain = LLMChain(\n",
+    "    llm=ManifestWrapper(client=manifest), \n",
+    "    prompt=prompt, \n",
+    "    verbose=True, \n",
+    "    memory=ConversationalBufferWindowMemory(k=8),\n",
+    ")\n",
+    "\n",
+    "output = chatgpt_chain.predict(human_input=\"Classes are \\\"positive\\\" and \\\"negative\\\". For example given\\nInput: I love this product!\\nOutput: positive.\\nI think this movie was one of the worst of the year. Script was boring!\")\n",
+    "print(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "4ef711d6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
+      "Prompt after formatting:\n",
+      "\u001b[32;1m\u001b[1;3mI am a classification model. It will try to classify your input.\n",
+      "\n",
+      "Human: Classes are \"positive\" and \"negative\". For example given\n",
+      "Input: I love this product!\n",
+      "Output: positive.\n",
+      "I think this movie was one of the worst of the year. Script was boring!\n",
+      "AI: negative.\n",
+      "Input: So awesome! I wish I could have gone\n",
+      "Output:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished LLMChain chain.\u001b[0m\n",
+      "positive.\n"
+     ]
+    }
+   ],
+   "source": [
+    "output = chatgpt_chain.predict(human_input=\"So awesome! I wish I could have gone\")\n",
+    "print(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "a5d6dac2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
+      "Prompt after formatting:\n",
+      "\u001b[32;1m\u001b[1;3mI am a classification model. It will try to classify your input.\n",
+      "\n",
+      "Human: Classes are \"positive\" and \"negative\". For example given\n",
+      "Input: I love this product!\n",
+      "Output: positive.\n",
+      "I think this movie was one of the worst of the year. Script was boring!\n",
+      "AI: negative.\n",
+      "Human: So awesome! I wish I could have gone\n",
+      "AI: positive.\n",
+      "Input: Hate it.\n",
+      "Output:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished LLMChain chain.\u001b[0m\n",
+      "negative.\n"
+     ]
+    }
+   ],
+   "source": [
+    "output = chatgpt_chain.predict(human_input=\"Hate it.\")\n",
+    "print(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 43,
+   "id": "b9283077",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
+      "Prompt after formatting:\n",
+      "\u001b[32;1m\u001b[1;3mI am a classification model. It will try to classify your input.\n",
+      "\n",
+      "\n",
+      "Input: Classes are fruits \"apple\", \"banana\", \"orange\", \"pear\". For example given\n",
+      "Input: This fruit rippens off of the tree.\n",
+      "Output: banana.\n",
+      "Often comes in bosc and bartlett varieties.\n",
+      "Output:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished LLMChain chain.\u001b[0m\n",
+      "apple.\n"
+     ]
+    }
+   ],
+   "source": [
+    "chatgpt_chain.memory.clear()\n",
+    "output = chatgpt_chain.predict(human_input=\"Classes are fruits \\\"apple\\\", \\\"banana\\\", \\\"orange\\\", \\\"pear\\\". For example given\\nInput: This fruit rippens off of the tree.\\nOutput: banana.\\nOften comes in bosc and bartlett varieties.\")\n",
+    "print(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 44,
+   "id": "cd0a23d9",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
+      "Prompt after formatting:\n",
+      "\u001b[32;1m\u001b[1;3mI am a classification model. It will try to classify your input.\n",
+      "\n",
+      "Human: Classes are fruits \"apple\", \"banana\", \"orange\", \"pear\". For example given\n",
+      "Input: This fruit rippens off of the tree.\n",
+      "Output: banana.\n",
+      "Often comes in bosc and bartlett varieties.\n",
+      "AI: apple.\n",
+      "Input: Often associated with monkeys\n",
+      "Output:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished LLMChain chain.\u001b[0m\n",
+      "banana.\n"
+     ]
+    }
+   ],
+   "source": [
+    "output = chatgpt_chain.predict(human_input=\"Often associated with monkeys\")\n",
+    "print(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 45,
+   "id": "90db6eb2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
+      "Prompt after formatting:\n",
+      "\u001b[32;1m\u001b[1;3mI am a classification model. It will try to classify your input.\n",
+      "\n",
+      "Human: Classes are fruits \"apple\", \"banana\", \"orange\", \"pear\". For example given\n",
+      "Input: This fruit rippens off of the tree.\n",
+      "Output: banana.\n",
+      "Often comes in bosc and bartlett varieties.\n",
+      "AI: apple.\n",
+      "Human: Often associated with monkeys\n",
+      "AI: banana.\n",
+      "Input: Is the color red and often delicious.\n",
+      "Output:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished LLMChain chain.\u001b[0m\n",
+      "apple.\n"
+     ]
+    }
+   ],
+   "source": [
+    "output = chatgpt_chain.predict(human_input=\"Is the color red and often delicious.\")\n",
+    "print(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 48,
+   "id": "c3806f89",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
+      "Prompt after formatting:\n",
+      "\u001b[32;1m\u001b[1;3mI am a classification model. It will try to classify your input.\n",
+      "\n",
+      "\n",
+      "Input: Classes are colors \"red\", \"green\", \"blue\", \"yellow\". For example given\n",
+      "Input: The color of a school bus.\n",
+      "Output: yellow.\n",
+      "Is the color of the sky\n",
+      "Output:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished LLMChain chain.\u001b[0m\n",
+      "blue.\n"
+     ]
+    }
+   ],
+   "source": [
+    "chatgpt_chain.memory.clear()\n",
+    "output = chatgpt_chain.predict(human_input=\"Classes are colors \\\"red\\\", \\\"green\\\", \\\"blue\\\", \\\"yellow\\\". For example given\\nInput: The color of a school bus.\\nOutput: yellow.\\nIs the color of the sky\")\n",
+    "print(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "id": "f508f597",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
+      "Prompt after formatting:\n",
+      "\u001b[32;1m\u001b[1;3mI am a classification model. It will try to classify your input.\n",
+      "\n",
+      "Human: Classes are colors \"red\", \"green\", \"blue\", \"yellow\". For example given\n",
+      "Input: The color of a school bus.\n",
+      "Output: yellow.\n",
+      "Is the color of the sky\n",
+      "AI: blue.\n",
+      "Input: Color of a banana.\n",
+      "Output:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished LLMChain chain.\u001b[0m\n",
+      "yellow.\n"
+     ]
+    }
+   ],
+   "source": [
+    "output = chatgpt_chain.predict(human_input=\"Color of a banana.\")\n",
+    "print(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 50,
+   "id": "cbd607f4",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
+      "Prompt after formatting:\n",
+      "\u001b[32;1m\u001b[1;3mI am a classification model. It will try to classify your input.\n",
+      "\n",
+      "Human: Classes are colors \"red\", \"green\", \"blue\", \"yellow\". For example given\n",
+      "Input: The color of a school bus.\n",
+      "Output: yellow.\n",
+      "Is the color of the sky\n",
+      "AI: blue.\n",
+      "Human: Color of a banana.\n",
+      "AI: yellow.\n",
+      "Input: When someone is sick they are this color.\n",
+      "Output:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished LLMChain chain.\u001b[0m\n",
+      "green.\n"
+     ]
+    }
+   ],
+   "source": [
+    "output = chatgpt_chain.predict(human_input=\"When someone is sick they are this color.\")\n",
+    "print(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 51,
+   "id": "d33e0e28",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "\u001b[1m> Entering new LLMChain chain...\u001b[0m\n",
+      "Prompt after formatting:\n",
+      "\u001b[32;1m\u001b[1;3mI am a classification model. It will try to classify your input.\n",
+      "\n",
+      "Human: Classes are colors \"red\", \"green\", \"blue\", \"yellow\". For example given\n",
+      "Input: The color of a school bus.\n",
+      "Output: yellow.\n",
+      "Is the color of the sky\n",
+      "AI: blue.\n",
+      "Human: Color of a banana.\n",
+      "AI: yellow.\n",
+      "Human: When someone is sick they are this color.\n",
+      "AI: green.\n",
+      "Input: Color of anger.\n",
+      "Output:\u001b[0m\n",
+      "\n",
+      "\u001b[1m> Finished LLMChain chain.\u001b[0m\n",
+      "red.\n"
+     ]
+    }
+   ],
+   "source": [
+    "output = chatgpt_chain.predict(human_input=\"Color of anger.\")\n",
+    "print(output)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "bootleg",
+   "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.12 | packaged by conda-forge | (default, Jan 30 2022, 23:36:06) \n[Clang 11.1.0 ]"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "7a3f97ab0465937066e9b79893b779dfc8a12d73c41f9d98a7bf05133c798250"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

duckdb-nsql/manifest/examples/manifest_async.py

@@ -0,0 +1,27 @@
+import asyncio
+import time
+
+from manifest import Manifest
+
+
+def main():
+
+    manifest = Manifest(
+        client_name="openaichat",
+    )
+
+    print("Running in serial")
+    prompts = [f"Tell me something interesting about {i}" for i in range(50)]
+    st = time.time()
+    for pmt in prompts:
+        _ = manifest.run(pmt)
+    print(f"For loop: {time.time() - st :.2f}")
+
+    print("Running with async")
+    st = time.time()
+    _ = asyncio.run(manifest.arun_batch(prompts, max_tokens=30))
+    print(f"Async loop: {time.time() - st :.2f}")
+
+
+if __name__ == "__main__":
+    main()

duckdb-nsql/manifest/examples/manifest_azure.ipynb

@@ -0,0 +1,149 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "AZURE_KEY = \"API_KEY::URL\"\n",
+    "OPENAI_KEY = \"sk-XXX\""
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use Azure and OpenAI models"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from manifest import Manifest\n",
+    "from manifest.connections.client_pool import ClientConnection\n",
+    "from pathlib import Path\n",
+    "\n",
+    "cache_path = Path(\"manifest.db\")\n",
+    "if cache_path.exists():\n",
+    "    cache_path.unlink()\n",
+    "\n",
+    "\n",
+    "azure = ClientConnection(\n",
+    "    client_name=\"azureopenai\",\n",
+    "    client_connection=AZURE_KEY,\n",
+    "    engine=\"text-davinci-003\",\n",
+    ")\n",
+    "\n",
+    "manifest = Manifest(client_pool=[azure], \n",
+    "    cache_name=\"sqlite\",\n",
+    "    cache_connection=\"manifest.db\"\n",
+    ")\n",
+    "\n",
+    "\n",
+    "openai = ClientConnection(\n",
+    "    client_name=\"openai\",\n",
+    "    client_connection=OPENAI_KEY,\n",
+    "    engine=\"text-davinci-003\",\n",
+    ")\n",
+    "\n",
+    "manifest_openai_nocache = Manifest(client_pool=[openai])\n",
+    "\n",
+    "manifest_openai = Manifest(client_pool=[openai], \n",
+    "    cache_name=\"sqlite\",\n",
+    "    cache_connection=\"manifest.db\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Show caches are the same\n",
+    "text = \"What is the meaning of life?\"\n",
+    "res = manifest.run(text, max_tokens=100, temperature=0.7, return_response=True)\n",
+    "print(res.get_response())\n",
+    "print(res.is_cached())\n",
+    "res2 = manifest_openai.run(text, max_tokens=100, temperature=0.7, return_response=True)\n",
+    "print(res2.is_cached())\n",
+    "\n",
+    "assert res2.get_response() == res.get_response()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "azure_chat = ClientConnection(\n",
+    "    client_name=\"azureopenaichat\",\n",
+    "    client_connection=AZURE_KEY,\n",
+    "    engine=\"gpt-3.5-turbo\",\n",
+    ")\n",
+    "\n",
+    "manifest = Manifest(client_pool=[azure_chat])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(manifest.run(\"What do you think is the best food?\", max_tokens=100))\n",
+    "\n",
+    "chat_dict = [\n",
+    "    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
+    "    {\"role\": \"user\", \"content\": \"Who won the world series in 2020?\"},\n",
+    "    {\"role\": \"assistant\", \"content\": \"The Los Angeles Dodgers won the World Series in 2020.\"},\n",
+    "    {\"role\": \"user\", \"content\": \"Where was it played?\"}\n",
+    "]\n",
+    "print(manifest.run(chat_dict, max_tokens=100))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "manifest",
+   "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.10.4"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "fddffe4ac3b9f00470127629076101c1b5f38ecb1e7358b567d19305425e9491"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

duckdb-nsql/manifest/examples/manifest_chatgpt.ipynb

@@ -0,0 +1,101 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "OPENAI_KEY = \"sk-XXX\""
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use ChatOpenAI\n",
+    "\n",
+    "Set you `OPENAI_API_KEY` environment variable."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from manifest import Manifest\n",
+    "from manifest.connections.client_pool import ClientConnection\n",
+    "\n",
+    "openai_chat = ClientConnection(\n",
+    "    client_name=\"openaichat\",\n",
+    "    client_connection=OPENAI_KEY,\n",
+    "    engine=\"gpt-3.5-turbo\"\n",
+    ")\n",
+    "\n",
+    "manifest = Manifest(client_pool=[openai_chat])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Simple question\n",
+    "chat_dict = [\n",
+    "    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
+    "    {\"role\": \"user\", \"content\": \"Who won the world series in 2020?\"},\n",
+    "    {\"role\": \"assistant\", \"content\": \"The Los Angeles Dodgers won the World Series in 2020.\"},\n",
+    "    {\"role\": \"user\", \"content\": \"Where was it played?\"}\n",
+    "]\n",
+    "print(manifest.run(chat_dict, max_tokens=100))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "manifest",
+   "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.10.4"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "fddffe4ac3b9f00470127629076101c1b5f38ecb1e7358b567d19305425e9491"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

duckdb-nsql/manifest/examples/manifest_connection_pool.ipynb

@@ -0,0 +1,208 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "OPENAI_KEY1 = \"sk-XXX\"\n",
+    "OPENAI_KEY2 = \"sk-XX\""
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use OpenAI\n",
+    "\n",
+    "Set you `OPENAI_API_KEY` environment variable."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from manifest import Manifest\n",
+    "from manifest.connections.client_pool import ClientConnection\n",
+    "\n",
+    "openai_ada = ClientConnection(\n",
+    "    client_name=\"openai\",\n",
+    "    client_connection=OPENAI_KEY1,\n",
+    "    engine=\"text-ada-001\"\n",
+    ")\n",
+    "\n",
+    "openai_curie = ClientConnection(\n",
+    "    client_name=\"openai\",\n",
+    "    client_connection=OPENAI_KEY2,\n",
+    "    engine=\"text-curie-001\"\n",
+    ")\n",
+    "\n",
+    "manifest = Manifest(client_pool=[openai_ada, openai_curie], client_pool_schedule=\"round_robin\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0\n",
+      "I am a model.\n",
+      "1\n",
+      "I am a MacBook Pro with a retina\n"
+     ]
+    }
+   ],
+   "source": [
+    "res = manifest.run(\"What model are you?\", temperature=0.0)\n",
+    "print(manifest.client_pool.current_client_id)\n",
+    "print(res)\n",
+    "res = manifest.run(\"What model are you?\", temperature=0.0)\n",
+    "print(manifest.client_pool.current_client_id)\n",
+    "print(res)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## With Async"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import nest_asyncio\n",
+    "# This is required for asyncio.run(...) to work in Jupyter notebooks.\n",
+    "nest_asyncio.apply()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from manifest import Manifest\n",
+    "from manifest.connections.client_pool import ClientConnection\n",
+    "\n",
+    "openai_ada = ClientConnection(\n",
+    "    client_name=\"openai\",\n",
+    "    client_connection=OPENAI_KEY1,\n",
+    "    engine=\"text-ada-001\"\n",
+    ")\n",
+    "\n",
+    "openai_babbage = ClientConnection(\n",
+    "    client_name=\"openai\",\n",
+    "    client_connection=OPENAI_KEY2,\n",
+    "    engine=\"text-babbage-001\"\n",
+    ")\n",
+    "\n",
+    "openai_curie = ClientConnection(\n",
+    "    client_name=\"openai\",\n",
+    "    client_connection=OPENAI_KEY2,\n",
+    "    engine=\"text-curie-001\"\n",
+    ")\n",
+    "\n",
+    "manifest = Manifest(client_pool=[openai_ada, openai_babbage, openai_curie], client_pool_schedule=\"round_robin\")\n",
+    "manifest_single_client = Manifest(client_pool=[openai_babbage])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "For loop: 128.68\n",
+      "Running with async single client\n",
+      "Running 1 tasks across all clients.\n",
+      "Async loop: 4.02\n",
+      "Running with async two clients but not chunking\n",
+      "Running 1 tasks across all clients.\n",
+      "Async loop: 3.92\n",
+      "Running with async two clients and chunk size\n",
+      "Running 20 tasks across all clients.\n",
+      "Async loop: 1.44\n"
+     ]
+    }
+   ],
+   "source": [
+    "import time\n",
+    "import asyncio\n",
+    "\n",
+    "prompts = [f\"Tell me something interesting about {i}\" for i in range(400)]\n",
+    "st = time.time()\n",
+    "for pmt in prompts:\n",
+    "    _ = manifest_single_client.run(pmt, max_tokens=30)\n",
+    "print(f\"For loop: {time.time() - st :.2f}\")\n",
+    "\n",
+    "print(\"Running with async single client\")\n",
+    "st = time.time()\n",
+    "_ = asyncio.run(manifest_single_client.arun_batch(prompts, max_tokens=30, chunk_size=-1))\n",
+    "print(f\"Async loop: {time.time() - st :.2f}\")\n",
+    "\n",
+    "print(\"Running with async two clients but not chunking\")\n",
+    "st = time.time()\n",
+    "_ = asyncio.run(manifest.arun_batch(prompts, max_tokens=30, chunk_size=-1))\n",
+    "print(f\"Async loop: {time.time() - st :.2f}\")\n",
+    "\n",
+    "print(\"Running with async two clients and chunk size\")\n",
+    "st = time.time()\n",
+    "_ = asyncio.run(manifest.arun_batch(prompts, max_tokens=30, chunk_size=20))\n",
+    "print(f\"Async loop: {time.time() - st :.2f}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "manifest",
+   "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.10.4"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "fddffe4ac3b9f00470127629076101c1b5f38ecb1e7358b567d19305425e9491"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

duckdb-nsql/manifest/examples/manifest_embedding.ipynb

@@ -0,0 +1,156 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use OpenAI\n",
+    "\n",
+    "Set you `OPENAI_API_KEY` environment variable."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'model_name': 'openaiembedding', 'engine': 'text-embedding-ada-002'}\n"
+     ]
+    }
+   ],
+   "source": [
+    "from manifest import Manifest\n",
+    "\n",
+    "manifest = Manifest(client_name=\"openaiembedding\")\n",
+    "print(manifest.client_pool.get_next_client().get_model_params())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(1536,)\n"
+     ]
+    }
+   ],
+   "source": [
+    "emb = manifest.run(\"Is this an embedding?\")\n",
+    "print(emb.shape)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Using Locally Hosted Huggingface LM\n",
+    "\n",
+    "Run\n",
+    "```\n",
+    "python3 manifest/api/app.py --model_type huggingface --model_name_or_path EleutherAI/gpt-neo-125M --device 0\n",
+    "```\n",
+    "or\n",
+    "```\n",
+    "python3 manifest/api/app.py --model_type sentence_transformers --model_name_or_path all-mpnet-base-v2 --device 0\n",
+    "```\n",
+    "\n",
+    "in a separate `screen` or `tmux`. Make sure to note the port. You can change this with `export FLASK_PORT=<port>`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'model_name': 'all-mpnet-base-v2', 'model_path': 'all-mpnet-base-v2', 'client_name': 'huggingfaceembedding'}\n"
+     ]
+    }
+   ],
+   "source": [
+    "from manifest import Manifest\n",
+    "\n",
+    "# Local hosted GPT Neo 125M\n",
+    "manifest = Manifest(\n",
+    "    client_name=\"huggingfaceembedding\",\n",
+    "    client_connection=\"http://127.0.0.1:6000\",\n",
+    "    cache_name=\"sqlite\",\n",
+    "    cache_connection=\"my_sqlite_manifest.sqlite\"\n",
+    ")\n",
+    "print(manifest.client_pool.get_next_client().get_model_params())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(768,)\n",
+      "(768,) (768,)\n"
+     ]
+    }
+   ],
+   "source": [
+    "emb = manifest.run(\"Is this an embedding?\")\n",
+    "print(emb.shape)\n",
+    "\n",
+    "emb = manifest.run([\"Is this an embedding?\", \"Bananas!!!\"])\n",
+    "print(emb[0].shape, emb[1].shape)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "manifest",
+   "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.10.4"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "fddffe4ac3b9f00470127629076101c1b5f38ecb1e7358b567d19305425e9491"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

duckdb-nsql/manifest/examples/manifest_google.ipynb

@@ -0,0 +1,117 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "GOOGLE_KEY = \"KEY::PROJECT_ID\""
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use GoogleVertexAPI"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from manifest import Manifest\n",
+    "from manifest.connections.client_pool import ClientConnection\n",
+    "\n",
+    "google_bison = ClientConnection(\n",
+    "    client_name=\"google\",\n",
+    "    client_connection=GOOGLE_KEY\n",
+    ")\n",
+    "\n",
+    "manifest = Manifest(client_pool=[google_bison])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Simple question\n",
+    "print(manifest.run(\"What is your name\", max_tokens=40))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from manifest import Manifest\n",
+    "from manifest.connections.client_pool import ClientConnection\n",
+    "\n",
+    "google_bison = ClientConnection(\n",
+    "    client_name=\"googlechat\",\n",
+    "    client_connection=GOOGLE_KEY\n",
+    ")\n",
+    "\n",
+    "manifest = Manifest(client_pool=[google_bison])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chat_dict = [\n",
+    "    # {\"author\": \"bot\", \"content\": \"You are a helpful assistant.\"},\n",
+    "    {\"author\": \"user\", \"content\": \"Who won the world series in 2020?\"},\n",
+    "    {\"author\": \"bot\", \"content\": \"The Los Angeles Dodgers won the World Series in 2020.\"},\n",
+    "    {\"author\": \"user\", \"content\": \"Where was it played?\"}\n",
+    "]\n",
+    "print(manifest.run(chat_dict, max_tokens=8))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "manifest",
+   "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.10.4"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "fddffe4ac3b9f00470127629076101c1b5f38ecb1e7358b567d19305425e9491"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

duckdb-nsql/manifest/examples/manifest_openrouter.ipynb

@@ -0,0 +1,108 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "OPENROUTER_API_KEY = \"sk-...\""
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use ChatOpenAI\n",
+    "\n",
+    "Set you `OPENROUTER_API_KEY` environment variable."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from manifest import Manifest\n",
+    "from manifest.connections.client_pool import ClientConnection\n",
+    "\n",
+    "openai_chat = ClientConnection(\n",
+    "    client_name=\"openrouter\",\n",
+    "    client_connection=OPENROUTER_API_KEY,\n",
+    "    engine=\"meta-llama/codellama-70b-instruct\"\n",
+    ")\n",
+    "\n",
+    "manifest = Manifest(client_pool=[openai_chat])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2020 World Series was played at the Globe Life Field in Arlington, Texas.\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Simple question\n",
+    "chat_dict = [\n",
+    "    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
+    "    {\"role\": \"user\", \"content\": \"Who won the world series in 2020?\"},\n",
+    "    {\"role\": \"assistant\", \"content\": \"The Los Angeles Dodgers won the World Series in 2020.\"},\n",
+    "    {\"role\": \"user\", \"content\": \"Where was it played?\"}\n",
+    "]\n",
+    "print(manifest.run(chat_dict, max_tokens=100))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.11.5"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "fddffe4ac3b9f00470127629076101c1b5f38ecb1e7358b567d19305425e9491"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

duckdb-nsql/manifest/examples/manifest_streaming.ipynb

@@ -0,0 +1,105 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "OPENAI_KEY = \"sk-XXX\""
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Use ChatOpenAI\n",
+    "\n",
+    "Set you `OPENAI_API_KEY` environment variable."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from manifest import Manifest\n",
+    "from manifest.connections.client_pool import ClientConnection\n",
+    "\n",
+    "openai_chat = ClientConnection(\n",
+    "    client_name=\"openaichat\",\n",
+    "    client_connection=OPENAI_KEY,\n",
+    "    engine=\"gpt-3.5-turbo\"\n",
+    ")\n",
+    "\n",
+    "manifest = Manifest(client_pool=[openai_chat])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "manifest_iterator = manifest.run(\"Tell me a story about a fat cat.\\n\\nOnce upon a time\", max_tokens=200, stream=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "\n",
+    "cur_line_length = 0\n",
+    "# Iterate over stream\n",
+    "for res in manifest_iterator:\n",
+    "    sys.stdout.write(res)\n",
+    "    cur_line_length += len(res)\n",
+    "    if cur_line_length > 80:\n",
+    "        sys.stdout.write(\"\\n\")\n",
+    "        cur_line_length = 0"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "manifest",
+   "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.10.4"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "fddffe4ac3b9f00470127629076101c1b5f38ecb1e7358b567d19305425e9491"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

duckdb-nsql/manifest/examples/manifest_together.ipynb

@@ -0,0 +1,106 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "env: TOMA_URL=<TOMA_URL>\n"
+     ]
+    }
+   ],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "%env TOMA_URL=<TOMA_URL>"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from manifest import Manifest\n",
+    "\n",
+    "# The responses are not fast\n",
+    "manifest = Manifest(\n",
+    "    client_name=\"toma\",\n",
+    ")\n",
+    "\n",
+    "print(manifest.run(\"What is the color of an apple?\"))"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "With a cache"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from manifest import Manifest\n",
+    "\n",
+    "# The responses are not fast\n",
+    "manifest = Manifest(\n",
+    "    client_name=\"toma\",\n",
+    "    cache_name=\"sqlite\",\n",
+    "    cache_connection=\"my_manifest_cache.sqlite\",\n",
+    ")\n",
+    "\n",
+    "res = manifest.run(\"What is the color of an apple?\", return_response=True)\n",
+    "print(res.get_response())\n",
+    "print(\"Is Cached?\", res.is_cached())\n",
+    "\n",
+    "res = manifest.run(\"What is the color of an apple?\", return_response=True)\n",
+    "print(res.get_response())\n",
+    "print(\"Is Cached?\", res.is_cached())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "manifest",
+   "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.10.4"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "fddffe4ac3b9f00470127629076101c1b5f38ecb1e7358b567d19305425e9491"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

duckdb-nsql/manifest/manifest/__init__.py

@@ -0,0 +1,6 @@
+"""Manifest init."""
+from manifest.manifest import Manifest
+from manifest.request import Request
+from manifest.response import Response
+
+__all__ = ["Manifest", "Response", "Request"]

duckdb-nsql/manifest/manifest/api/__init__.py

@@ -0,0 +1 @@
+"""Api init."""

duckdb-nsql/manifest/manifest/api/app.py

@@ -0,0 +1,301 @@
+"""Flask app."""
+import argparse
+import io
+import json
+import logging
+import os
+import socket
+from typing import Dict
+
+import pkg_resources
+from flask import Flask, Response, request
+
+from manifest.api.models.diffuser import DiffuserModel
+from manifest.api.models.huggingface import (
+    MODEL_GENTYPE_REGISTRY,
+    CrossModalEncoderModel,
+    TextGenerationModel,
+)
+from manifest.api.models.sentence_transformer import SentenceTransformerModel
+from manifest.api.response import ModelResponse
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+logger = logging.getLogger(__name__)
+app = Flask(__name__)  # define app using Flask
+# Will be global
+model = None
+model_type = None
+PORT = int(os.environ.get("FLASK_PORT", 5000))
+MODEL_CONSTRUCTORS = {
+    "huggingface": TextGenerationModel,
+    "sentence_transformers": SentenceTransformerModel,
+    "huggingface_crossmodal": CrossModalEncoderModel,
+    "diffuser": DiffuserModel,
+}
+
+
+def parse_args() -> argparse.Namespace:
+    """Generate args."""
+    parser = argparse.ArgumentParser(description="Model args")
+    parser.add_argument(
+        "--model_type",
+        default=None,
+        type=str,
+        required=True,
+        help="Model type used for finding constructor.",
+        choices=MODEL_CONSTRUCTORS.keys(),
+    )
+    parser.add_argument(
+        "--model_generation_type",
+        default=None,
+        type=str,
+        help="Model generation type.",
+        choices=MODEL_GENTYPE_REGISTRY.keys(),
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        default=None,
+        type=str,
+        help="Name of model or path to model. Used in initialize of model class.",
+    )
+    parser.add_argument(
+        "--cache_dir", default=None, type=str, help="Cache directory for models."
+    )
+    parser.add_argument(
+        "--device", type=int, default=0, help="Model device. -1 for CPU."
+    )
+    parser.add_argument(
+        "--fp16", action="store_true", help="Force use fp16 for model params."
+    )
+    parser.add_argument(
+        "--percent_max_gpu_mem_reduction",
+        type=float,
+        default=0.85,
+        help="Used with accelerate multigpu. Scales down max memory.",
+    )
+    parser.add_argument(
+        "--use_bitsandbytes",
+        action="store_true",
+        help=("Use bits and bytes. " "This will override --device parameter."),
+    )
+    parser.add_argument(
+        "--use_accelerate_multigpu",
+        action="store_true",
+        help=(
+            "Use accelerate for multi gpu inference. "
+            "This will override --device parameter."
+        ),
+    )
+    parser.add_argument(
+        "--use_hf_parallelize",
+        action="store_true",
+        help=(
+            "Use HF parallelize for multi gpu inference. "
+            "This will override --device parameter."
+        ),
+    )
+    parser.add_argument(
+        "--use_deepspeed",
+        action="store_true",
+        help=("Use deepspeed. This will override --device parameter."),
+    )
+    args = parser.parse_args()
+    return args
+
+
+def is_port_in_use(port: int) -> bool:
+    """Check if port is in use."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        return s.connect_ex(("localhost", port)) == 0
+
+
+def main() -> None:
+    """Run main."""
+    kwargs = parse_args()
+    if is_port_in_use(PORT):
+        raise ValueError(f"Port {PORT} is already in use.")
+    global model_type
+    model_type = kwargs.model_type
+    model_gen_type = kwargs.model_generation_type
+    model_name_or_path = kwargs.model_name_or_path
+    if not model_name_or_path:
+        raise ValueError("Must provide model_name_or_path.")
+    if kwargs.use_accelerate_multigpu:
+        logger.info("Using accelerate. Overridding --device argument.")
+    if (
+        kwargs.percent_max_gpu_mem_reduction <= 0
+        or kwargs.percent_max_gpu_mem_reduction > 1
+    ):
+        raise ValueError("percent_max_gpu_mem_reduction must be in (0, 1].")
+    if (
+        sum(
+            [
+                kwargs.use_accelerate_multigpu,
+                kwargs.use_hf_parallelize,
+                kwargs.use_bitsandbytes,
+                kwargs.use_deepspeed,
+            ]
+        )
+        > 1
+    ):
+        raise ValueError(
+            "Only one of use_accelerate_multigpu, use_hf_parallelize, "
+            "use_bitsandbytes, and use_deepspeed can be set."
+        )
+    # Global model
+    global model
+    model = MODEL_CONSTRUCTORS[model_type](
+        model_name_or_path,
+        model_type=model_gen_type,
+        cache_dir=kwargs.cache_dir,
+        device=kwargs.device,
+        use_accelerate=kwargs.use_accelerate_multigpu,
+        use_parallelize=kwargs.use_hf_parallelize,
+        use_bitsandbytes=kwargs.use_bitsandbytes,
+        use_deepspeed=kwargs.use_deepspeed,
+        perc_max_gpu_mem_red=kwargs.percent_max_gpu_mem_reduction,
+        use_fp16=kwargs.fp16,
+    )
+    app.run(host="0.0.0.0", port=PORT)
+
+
+@app.route("/completions", methods=["POST"])
+def completions() -> Response:
+    """Get completions for generation."""
+    prompt = request.json["prompt"]
+    del request.json["prompt"]
+    generation_args = request.json
+
+    if not isinstance(prompt, (str, list)):
+        raise ValueError("Prompt must be a str or list of str")
+    try:
+        result_gens = []
+        for generations in model.generate(prompt, **generation_args):
+            result_gens.append(generations)
+        if model_type == "diffuser":
+            # Assign None logprob as it's not supported in diffusers
+            results = [
+                {"array": r[0], "logprob": None, "tokens": None, "token_logprobs": None}
+                for r in result_gens
+            ]
+            res_type = "image_generation"
+        else:
+            results = [
+                {"text": r[0], "logprob": r[1], "tokens": r[2], "token_logprobs": r[3]}
+                for r in result_gens
+            ]
+            res_type = "text_completion"
+        # transform the result into the openai format
+        return Response(
+            json.dumps(ModelResponse(results, response_type=res_type).__dict__()),
+            status=200,
+        )
+    except Exception as e:
+        logger.error(e)
+        return Response(
+            json.dumps({"message": str(e)}),
+            status=400,
+        )
+
+
+@app.route("/embed", methods=["POST"])
+def embed() -> Response:
+    """Get embed for generation."""
+    if "modality" in request.json:
+        modality = request.json["modality"]
+    else:
+        modality = "text"
+    if modality == "text":
+        prompts = request.json["prompt"]
+    elif modality == "image":
+        import base64
+
+        from PIL import Image
+
+        prompts = [
+            Image.open(io.BytesIO(base64.b64decode(data)))
+            for data in request.json["prompt"]
+        ]
+    else:
+        raise ValueError("modality must be text or image")
+
+    try:
+        results = []
+        embeddings = model.embed(prompts)
+        for embedding in embeddings:
+            results.append(
+                {
+                    "array": embedding,
+                    "logprob": None,
+                    "tokens": None,
+                    "token_logprobs": None,
+                }
+            )
+
+        return Response(
+            json.dumps(
+                ModelResponse(results, response_type="embedding_generation").__dict__()
+            ),
+            status=200,
+        )
+    except Exception as e:
+        logger.error(e)
+        return Response(
+            json.dumps({"message": str(e)}),
+            status=400,
+        )
+
+
+@app.route("/score_sequence", methods=["POST"])
+def score_sequence() -> Response:
+    """Get logprob of prompt."""
+    prompt = request.json["prompt"]
+    del request.json["prompt"]
+    generation_args = request.json
+
+    if not isinstance(prompt, (str, list)):
+        raise ValueError("Prompt must be a str or list of str")
+
+    try:
+        score_list = model.score_sequence(prompt, **generation_args)
+        results = [
+            {
+                "text": prompt if isinstance(prompt, str) else prompt[i],
+                "logprob": r[0],
+                "tokens": r[1],
+                "token_logprobs": r[2],
+            }
+            for i, r in enumerate(score_list)
+        ]
+        # transform the result into the openai format
+        return Response(
+            json.dumps(
+                ModelResponse(results, response_type="prompt_logit_score").__dict__()
+            ),
+            status=200,
+        )
+    except Exception as e:
+        logger.error(e)
+        return Response(
+            json.dumps({"message": str(e)}),
+            status=400,
+        )
+
+
+@app.route("/params", methods=["POST"])
+def params() -> Dict:
+    """Get model params."""
+    return model.get_init_params()
+
+
+@app.route("/")
+def index() -> str:
+    """Get index completion."""
+    fn = pkg_resources.resource_filename("metaseq", "service/index.html")
+    with open(fn) as f:
+        return f.read()
+
+
+if __name__ == "__main__":
+    main()