Spaces:
Sleeping
Sleeping
Create txt2sql_code3.py
Browse files- txt2sql_code3.py +158 -0
txt2sql_code3.py
ADDED
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import sqlite3
|
2 |
+
from sqlite3 import Error
|
3 |
+
from peft import AutoPeftModelForCausalLM
|
4 |
+
from transformers import AutoTokenizer, BitsAndBytesConfig
|
5 |
+
from transformers import AutoModelForCausalLM
|
6 |
+
from openai import OpenAI
|
7 |
+
import google.generativeai as genai
|
8 |
+
|
9 |
+
class SQLPromptModel:
|
10 |
+
def __init__(self, model_dir, database):
|
11 |
+
self.model_dir = model_dir
|
12 |
+
self.database = database
|
13 |
+
# peft_model_dir = self.model_dir
|
14 |
+
bnb_config = BitsAndBytesConfig(
|
15 |
+
load_in_4bit=True,
|
16 |
+
bnb_4bit_quant_type="nf4",
|
17 |
+
bnb_4bit_compute_dtype="float16",
|
18 |
+
bnb_4bit_use_double_quant=True,
|
19 |
+
)
|
20 |
+
# self.model = AutoPeftModelForCausalLM.from_pretrained(
|
21 |
+
# peft_model_dir, low_cpu_mem_usage=True, quantization_config=bnb_config
|
22 |
+
# )
|
23 |
+
# self.tokenizer = AutoTokenizer.from_pretrained(peft_model_dir)
|
24 |
+
# self.model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
|
25 |
+
# self.tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
|
26 |
+
self.chatgpt_client = OpenAI(api_key="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")
|
27 |
+
self.genai = genai
|
28 |
+
self.genai.configure(api_key="AIzaSyAFG94rVbm9eWepO5uPGsMha8XJ-sHbMdA")
|
29 |
+
self.genai_model = genai.GenerativeModel('gemini-pro')
|
30 |
+
|
31 |
+
self.conn = sqlite3.connect(self.database)
|
32 |
+
|
33 |
+
def fetch_table_schema(self, table_name):
|
34 |
+
"""Fetch the schema of a table from the database."""
|
35 |
+
cursor = self.conn.cursor()
|
36 |
+
cursor.execute(f"PRAGMA table_info({table_name})")
|
37 |
+
schema = cursor.fetchall()
|
38 |
+
if schema:
|
39 |
+
return schema
|
40 |
+
else:
|
41 |
+
print(f"Table {table_name} does not exist or has no schema.")
|
42 |
+
return None
|
43 |
+
|
44 |
+
def text2sql(self, schema, user_prompt, inp_prompt=None):
|
45 |
+
"""Generate SQL query based on user prompt and table schema.inp_prompt is for gradio purpose"""
|
46 |
+
table_columns = ', '.join([f"{col[1]} {col[2]}" for col in schema])
|
47 |
+
|
48 |
+
prompt = f"""Below are SQL table schemas paired with instructions that describe a task.
|
49 |
+
Using valid SQLite, write a response that appropriately completes the request for the provided tables.
|
50 |
+
### Instruction: {user_prompt} ###
|
51 |
+
Input: CREATE TABLE sql_pdf({table_columns});
|
52 |
+
### Response: (Return only query , nothing extra)"""
|
53 |
+
|
54 |
+
if inp_prompt is not None :
|
55 |
+
prompt = prompt.replace(user_prompt, inp_prompt + " ")
|
56 |
+
else:
|
57 |
+
inp_prompt = input("Press Enter for default question or Enter user prompt without newline characters: ").strip()
|
58 |
+
if inp_prompt:
|
59 |
+
prompt = prompt.replace(user_prompt, inp_prompt + " ")
|
60 |
+
|
61 |
+
"""Text to SQL query generation"""
|
62 |
+
input_ids = self.tokenizer(
|
63 |
+
prompt, return_tensors="pt", truncation=True
|
64 |
+
).input_ids.to(next(self.model.parameters()).device) # Move input to the device of the model
|
65 |
+
outputs = self.model.generate(input_ids=input_ids, max_new_tokens=200)
|
66 |
+
response = self.tokenizer.batch_decode(
|
67 |
+
outputs.detach().cpu().numpy(), skip_special_tokens=True
|
68 |
+
)[0][:]
|
69 |
+
return response[len(prompt):]
|
70 |
+
|
71 |
+
def text2sql_chatgpt(self, schema, user_prompt, inp_prompt=None):
|
72 |
+
table_columns = ', '.join([f"{col[1]} {col[2]}" for col in schema])
|
73 |
+
|
74 |
+
prompt = f"""Below are SQL table schemas paired with instructions that describe a task.
|
75 |
+
Using valid SQLite, write a response that appropriately completes the request for the provided tables.
|
76 |
+
### Instruction: {user_prompt} ###
|
77 |
+
Input: CREATE TABLE sql_pdf({table_columns});
|
78 |
+
### Response: (Return only generated query based on user_prompt , nothing extra)"""
|
79 |
+
|
80 |
+
if inp_prompt is not None :
|
81 |
+
prompt = prompt.replace(user_prompt, inp_prompt + " ")
|
82 |
+
else:
|
83 |
+
inp_prompt = input("Press Enter for default question or Enter user prompt without newline characters: ").strip()
|
84 |
+
if inp_prompt:
|
85 |
+
prompt = prompt.replace(user_prompt, inp_prompt + " ")
|
86 |
+
print(prompt)
|
87 |
+
completion = self.chatgpt_client.chat.completions.create(
|
88 |
+
model="gpt-3.5-turbo",
|
89 |
+
messages=[
|
90 |
+
{"role": "system", "content": "You are a expert SQL developer , generate a sql query and return it"},
|
91 |
+
{"role": "user", "content": prompt }
|
92 |
+
]
|
93 |
+
)
|
94 |
+
return completion.choices[0].message.content
|
95 |
+
|
96 |
+
def text2sql_gemini(self, schema, user_prompt, inp_prompt=None):
|
97 |
+
table_columns = ', '.join([f"{col[1]} {col[2]}" for col in schema])
|
98 |
+
|
99 |
+
prompt = f"""Below are SQL table schemas paired with instructions that describe a task.
|
100 |
+
Using valid SQLite, write a response that appropriately completes the request for the provided tables.
|
101 |
+
### Instruction: {user_prompt} ###
|
102 |
+
Input: CREATE TABLE sql_pdf({table_columns});
|
103 |
+
### Response: (Return only generated query based on user_prompt , nothing extra)"""
|
104 |
+
|
105 |
+
if inp_prompt is not None :
|
106 |
+
prompt = prompt.replace(user_prompt, inp_prompt + " ")
|
107 |
+
else:
|
108 |
+
inp_prompt = input("Press Enter for default question or Enter user prompt without newline characters: ").strip()
|
109 |
+
if inp_prompt:
|
110 |
+
prompt = prompt.replace(user_prompt, inp_prompt + " ")
|
111 |
+
print(prompt)
|
112 |
+
completion = self.genai_model.generate_content(prompt)
|
113 |
+
generated_query=completion.text
|
114 |
+
start_index = generated_query.find("SELECT")
|
115 |
+
end_index = generated_query.find(";", start_index) + 1
|
116 |
+
print(start_index,end_index)
|
117 |
+
if start_index != -1 and end_index != 0:
|
118 |
+
return generated_query[start_index:end_index]
|
119 |
+
else:
|
120 |
+
return generated_query
|
121 |
+
|
122 |
+
|
123 |
+
|
124 |
+
def execute_query(self, query):
|
125 |
+
"""Executing the query on database and returning rows and columns."""
|
126 |
+
print(query)
|
127 |
+
cur = self.conn.cursor()
|
128 |
+
cur.execute(query)
|
129 |
+
col = [header[0] for header in cur.description]
|
130 |
+
dash = "-" * sum(len(col_name) + 4 for col_name in col)
|
131 |
+
print(tuple(col))
|
132 |
+
print(dash)
|
133 |
+
rows = []
|
134 |
+
for member in cur:
|
135 |
+
rows.append(member)
|
136 |
+
print(member)
|
137 |
+
cur.close()
|
138 |
+
self.conn.commit()
|
139 |
+
# print(rows)
|
140 |
+
return rows, col
|
141 |
+
|
142 |
+
if __name__ == "__main__":
|
143 |
+
model_dir = "multi_table_demo/checkpoint-2600"
|
144 |
+
database = r"sql_pdf.db"
|
145 |
+
sql_model = SQLPromptModel(model_dir, database)
|
146 |
+
user_prompt = "Give complete details of properties in India"
|
147 |
+
while True:
|
148 |
+
table_schema = sql_model.fetch_table_schema("sql_pdf")
|
149 |
+
if table_schema:
|
150 |
+
# query = sql_model.text2sql(table_schema, user_prompt)
|
151 |
+
# query = sql_model.text2sql_chatgpt(table_schema, user_prompt)
|
152 |
+
query = sql_model.text2sql_gemini(table_schema, user_prompt)
|
153 |
+
print(query)
|
154 |
+
sql_model.execute_query(query)
|
155 |
+
|
156 |
+
sql_model.conn.close()
|
157 |
+
|
158 |
+
|