[Droid] Update Natural SQL 7B Model and Dependencies

natural-sql-7b/README.md

@@ -0,0 +1,57 @@
+# Natural SQL 7B Truss
+
+Natural SQL 7B, developed by ChatDB and benchmarked using the SQL-Eval framework by the Defog team, is a groundbreaking model for converting natural language queries into sequential query language (SQL) with unparalleled accuracy. This model represents a significant advancement in natural language processing for databases. For more detailed information and to reference the source, visit the [Natural SQL 7B on Hugging Face](https://huggingface.co/chatdb/natural-sql-7b).
+
+## Deploying Natural SQL 7B on Baseten
+
+To deploy the Natural SQL 7B model on Baseten:
+
+1. Ensure you have cloned the truss-examples repository:
+    ```
+    git clone https://github.com/basetenlabs/truss-examples/
+    ```
+2. Navigate to the `natural-sql-7b` directory within the cloned repository.
+3. Deploy the truss to Baseten using the following command:
+    ```
+    truss push natural-sql-7b
+    ```
+
+## Input
+
+The model expects a JSON formatted input specifying 'question' and 'schema'. 'question' should be a natural language query, and 'schema' should describe the relevant database schema in string format. For example:
+```json
+{
+  "question": "Show me the day with the most users joining",
+  "schema": "CREATE TABLE users (user_id SERIAL PRIMARY KEY, username VARCHAR(50) NOT NULL, email VARCHAR(100) NOT NULL, password_hash TEXT NOT NULL, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP);"
+}
+```
+
+## Output
+
+The output is provided in JSON format. It includes the 'question' and a 'sql_query' – the SQL query representing the answer to the natural language question. For instance:
+```json
+{
+  "sql_query": "SELECT created_at::date AS join_date, COUNT(*) AS user_count FROM users GROUP BY join_date ORDER BY user_count DESC LIMIT 1;"
+}
+```
+
+## Example Usage
+
+An example command to use the model with Baseten is provided below:
+```shell
+truss predict --input '{"question": "Show me the day with the most users joining", "schema": "CREATE TABLE users (user_id SERIAL PRIMARY KEY, username VARCHAR(50) NOT NULL, email VARCHAR(100) NOT NULL, password_hash TEXT NOT NULL, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP);"}'
+```
+
+Expected output from the above command:
+```json
+{
+  "question": "Show me the day with the most users joining",
+  "sql_query": "SELECT created_at::date AS join_date, COUNT(*) AS user_count FROM users GROUP BY join_date ORDER BY user_count DESC LIMIT 1;"
+}
+```
+
+This quick example underlines the model's capability to accurately transform compendious natural language questions into executable SQL queries.
+
+## License
+
+This model and its underlying codebase are made available under the Apache 2.0 license. The model itself is distributed under the CC BY-SA 4.0 license, ensuring that it can be freely used and shared within the community under the same license terms.

natural-sql-7b/config.yaml

@@ -0,0 +1,19 @@
+description: Configuration for the Natural SQL 7B model following truss documentation conventions.
+model_name: Natural SQL 7B
+python_version: py311
+requirements_file: ./natural-sql-7b/requirements.txt
+model_metadata:
+  example_model_input: '{"question": "Show me the day with the most users joining", "schema": "CREATE TABLE users (user_id SERIAL PRIMARY KEY, username VARCHAR(50) NOT NULL, email VARCHAR(100) NOT NULL, password_hash TEXT NOT NULL, created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP);"}'
+resources:
+  use_gpu: true
+  accelerator: A10G
+model_cache:
+  repo_id: chatdb/natural-sql-7b
+  revision: main
+  allow_patterns: []
+  ignore_patterns: []
+requirements:
+  - transformers==4.35.2
+  - accelerate>=0.25.0
+  - sqlparse>=0.4.2
+external_package_dirs: []

natural-sql-7b/model.py

@@ -0,0 +1,59 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class Model:
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained("chatdb/natural-sql-7b")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            "chatdb/natural-sql-7b",
+            device_map="auto",
+            torch_dtype=torch.float16,
+        )
+
+    def predict(self, input_dict):
+        # Ensure that the input_dict contains the requisite keys for 'questions'
+        if "questions" not in input_dict:
+            raise ValueError("Input dictionary must contain a 'questions' key.")
+
+        questions = input_dict["questions"]
+
+        outputs = []
+        for question in questions:
+            prompt = f"""
+            ### Task
+
+            Generate a SQL query to answer the following question: `{question}`
+
+            ### Database Schema
+            The query will run on a database with the following schema:
+            ```schema
+            # ... include the schema here ...
+            ```
+
+            ### Answer
+            Here is the SQL query that answers the question: `{question}`
+            """
+            inputs = self.tokenizer(
+                prompt, return_tensors="pt", padding=True, truncation=True
+            )
+
+            # Check if CUDA is available and move tensors accordingly
+            if torch.cuda.is_available():
+                inputs = inputs.to("cuda")
+                self.model.to("cuda")
+
+            generated_ids = self.model.generate(
+                inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                max_length=512,
+                num_beams=5,
+                early_stopping=True,
+            )
+
+            # Post-process the output to extract the SQL query
+            output = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+            sql_query = output.split("```sql")[-1].strip()
+            outputs.append(sql_query)
+
+        return outputs

natural-sql-7b/model/model.py

@@ -0,0 +1,53 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class Model:
+    def __init__(self):
+        self.tokenizer = None
+        self.model = None
+
+    def load(self):
+        self.tokenizer = AutoTokenizer.from_pretrained("chatdb/natural-sql-7b")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            "chatdb/natural-sql-7b", device_map="auto", torch_dtype=torch.float16
+        )
+        # Ensure the model is optimized for CUDA if available
+        if torch.cuda.is_available():
+            self.model = self.model.to("cuda")
+
+    def predict(self, input_dict):
+        if "question" not in input_dict or "schema" not in input_dict:
+            raise ValueError(
+                "Input dictionary must contain 'question' and 'schema' keys."
+            )
+
+        question = input_dict["question"]
+        schema = input_dict["schema"]
+        prompt = f"""
+        ### Task
+        Generate a SQL query to answer the following question: `{question}`
+
+        ### PostgreSQL Database Schema
+        {schema}
+
+        ### Answer
+        Here is the SQL query that answers the question: `{question}`
+        """
+        # Ensure the tokenizer generates tensors directly on the model's device
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        generated_ids = self.model.generate(
+            **inputs,
+            num_return_sequences=1,
+            eos_token_id=self.tokenizer.eos_token_id,
+            pad_token_id=self.tokenizer.pad_token_id,
+            max_new_tokens=400,
+            do_sample=False,
+            num_beams=1,
+        )
+        result = (
+            self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+            .split("### Answer")[1]
+            .strip()
+        )
+        return {"question": question, "sql_query": result}

natural-sql-7b/requirements.txt

@@ -0,0 +1,3 @@
+transformers==4.35.2
+accelerate>=0.25.0
+sqlparse>=0.4.2