Views
No views yet

| Model and Dataset | Download Latest |
|---|---|
| SynSQL-2.5M | ✨ Modelscope, 🤗 HuggingFace |
| OmniSQL-7B | ✨ Modelscope, 🤗 HuggingFace |
| OmniSQL-14B | ✨ Modelscope, 🤗 HuggingFace |
| OmniSQL-32B | ✨ Modelscope, 🤗 HuggingFace |
<database, question, SQL query, chain-of-thought solution> quad.simple, moderate, complex, highly complex, from single-table queries to advanced multi-table joins, functions, and common table expressions.formal, colloquial, imperative, interrogative, descriptive, concise, vague, metaphorical, and conversational.
1input_prompt_template = '''Task Overview:
2You are a data science expert. Below, you are provided with a database schema and a natural language question. Your task is to understand the schema and generate a valid SQL query to answer the question.
3
4Database Engine:
5SQLite
6
7Database Schema:
8{db_details}
9This schema describes the database's structure, including tables, columns, primary keys, foreign keys, and any relevant relationships or constraints.
10
11Question:
12{question}
13
14Instructions:
15- Make sure you only output the information that is asked in the question. If the question asks for a specific column, make sure to only include that column in the SELECT clause, nothing more.
16- The generated query should return all of the information asked in the question without any missing or extra information.
17- Before generating the final SQL query, please think through the steps of how to write the query.
18
19Output Format:
20In your answer, please enclose the generated SQL query in a code block:
21```
22-- Your SQL query
23```
24
25Take a deep breath and think step by step to find the correct SQL query.'''CREATE TABLE statements (i.e., DDL) of tables in the database. You can add database values and column descriptions in DDLs with SQL comments. External knowledge can be concatenated with the natural language question and placed in the "question" placeholder. OmniSQL currently supports only SQLite, as the SQL queries in SynSQL-2.5M are synthesized using the SQLite dialect.examples folder.1from vllm import LLM, SamplingParams
2from transformers import AutoTokenizer
3
4prompt = input_prompt_template.format(db_details = "...", question = "...")
5model_path = "seeklhy/OmniSQL-7B"
6tokenizer = AutoTokenizer.from_pretrained(model_path)
7sampling_params = SamplingParams(
8 temperature = 0,
9 max_tokens = 2048,
10 n = 1
11)
12
13llm = LLM(
14 model = model_path,
15 dtype = "float16",
16 tensor_parallel_size = 1,
17 max_model_len = 8192,
18 gpu_memory_utilization = 0.92,
19 swap_space = 8,
20 enforce_eager = True,
21 disable_custom_all_reduce = True,
22 trust_remote_code = True
23)
24
25chat_prompt = tokenizer.apply_chat_template(
26 [{"role": "user", "content": prompt}],
27 add_generation_prompt = True, tokenize = False
28)
29
30outputs = llm.generate([chat_prompt], sampling_params)
31
32for output in outputs:
33 responses = [o.text for o in output.outputs]
34 print(responses[0])1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4prompt = input_prompt_template.format(db_details = "...", question = "...")
5model_path = "seeklhy/OmniSQL-7B"
6tokenizer = AutoTokenizer.from_pretrained(model_path)
7model = AutoModelForCausalLM.from_pretrained(
8 model_path,
9 torch_dtype=torch.bfloat16
10).to("cuda:0")
11
12chat_prompt = tokenizer.apply_chat_template(
13 [{"role": "user", "content": prompt}],
14 add_generation_prompt = True, tokenize = False
15)
16
17inputs = tokenizer([chat_prompt], return_tensors="pt")
18inputs = inputs.to(model.device)
19
20output_ids = model.generate(
21 **inputs,
22 eos_token_id = tokenizer.eos_token_id,
23 max_new_tokens = 2048
24)
25
26input_len = len(inputs.input_ids[0])
27output_ids = output_ids[0][input_len:]
28
29response = tokenizer.batch_decode([output_ids], skip_special_tokens = True)[0]
30print(response)