ignitariumcloud/knowledge_model

1import sqlite32from sqlite3 import Error3from peft import AutoPeftModelForCausalLM4from transformers import AutoTokenizer, BitsAndBytesConfig5from transformers import AutoModelForCausalLM6from openai import OpenAI7import google.generativeai as genai8 9class SQLPromptModel:10    def __init__(self, model_dir, database):11        self.model_dir = model_dir12        self.database = database13        # peft_model_dir = self.model_dir14        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_config22        # )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 = genai28        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 schema40        else:41            print(f"Table {table_name} does not exist or has no schema.")42            return None43        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        Select all columns unless specified in specific.51        Example row :1,Michael,michael@ignitarium.com,59,Female,Headphones,2023-01-03,General inquiry,Email,44 hours,88 hours,4,Technical Issue,Server crashes due to memory leaks in custom-developed software.,Closed,"Debug and optimize the software code to identify and fix memory leaks, and implement regular monitoring for early detection.",Medium52        ### Instruction: {user_prompt} ### 53        Input: CREATE TABLE ticket_dataset({table_columns});54        ### Response: (Return only query , nothing extra)"""55 56        if inp_prompt is not None :57            prompt = prompt.replace(user_prompt, inp_prompt + " ")58        else:59            inp_prompt = input("Press Enter for default question or Enter user prompt without newline characters: ").strip()60            if inp_prompt:61                prompt = prompt.replace(user_prompt, inp_prompt + " ")62 63        """Text to SQL query generation"""64        input_ids = self.tokenizer(65            prompt, return_tensors="pt", truncation=True66        ).input_ids.to(next(self.model.parameters()).device)  # Move input to the device of the model67        outputs = self.model.generate(input_ids=input_ids, max_new_tokens=200)68        response = self.tokenizer.batch_decode(69            outputs.detach().cpu().numpy(), skip_special_tokens=True70        )[0][:]71        return response[len(prompt):]72 73    def text2sql_chatgpt(self, schema, user_prompt, inp_prompt=None):74        table_columns = ', '.join([f"{col[1]} {col[2]}" for col in schema])75 76        prompt = f"""Below are SQL table schemas paired with instructions that describe a task.77        Using valid SQLite, write a response that appropriately completes the request for the provided tables.78        Select all columns unless specified in specific.79        Example row :1,Michael,michael@ignitarium.com,59,Female,Headphones,2023-01-03,General inquiry,Email,44 hours,88 hours,4,Technical Issue,Server crashes due to memory leaks in custom-developed software.,Closed,"Debug and optimize the software code to identify and fix memory leaks, and implement regular monitoring for early detection.",Medium80        ### Instruction: {user_prompt} ### 81        Input: CREATE TABLE ticket_dataset({table_columns});82        ### Response: (Return only generated query based on user_prompt , nothing extra)"""83 84        if inp_prompt is not None :85            prompt = prompt.replace(user_prompt, inp_prompt + " ")86        else:87            inp_prompt = input("Press Enter for default question or Enter user prompt without newline characters: ").strip()88            if inp_prompt:89                prompt = prompt.replace(user_prompt, inp_prompt + " ")90        print(prompt)91        completion = self.chatgpt_client.chat.completions.create(92            model="gpt-3.5-turbo",93            messages=[94                {"role": "system", "content": "You are a expert SQL developer , generate a sql query and return it"},95                {"role": "user", "content": prompt }96            ]97        )98        return completion.choices[0].message.content99 100    def text2sql_gemini(self, schema, user_prompt, inp_prompt=None):101        table_columns = ', '.join([f"{col[1]} {col[2]}" for col in schema])102 103        prompt = f"""Below are SQL table schemas paired with instructions that describe a task.104        Using valid SQLite, write a response that appropriately completes the request for the provided tables.105        Select all columns unless specified in specific.106        Example row :1,Michael,michael@ignitarium.com,59,Female,Headphones,2023-01-03,General inquiry,Email,44 hours,88 hours,4,Technical Issue,Server crashes due to memory leaks in custom-developed software.,Closed,"Debug and optimize the software code to identify and fix memory leaks, and implement regular monitoring for early detection.",Medium107        ### Instruction: {user_prompt} ### 108        Input: CREATE TABLE ticket_dataset({table_columns});109        ### Response: (Return only generated query based on user_prompt , nothing extra)"""110 111        if inp_prompt is not None :112            prompt = prompt.replace(user_prompt, inp_prompt + " ")113        else:114            inp_prompt = input("Press Enter for default question or Enter user prompt without newline characters: ").strip()115            if inp_prompt:116                prompt = prompt.replace(user_prompt, inp_prompt + " ")117        print(prompt)118        completion = self.genai_model.generate_content(prompt)119        generated_query=completion.text120        start_index = generated_query.find("SELECT")121        end_index = generated_query.find(";", start_index) + 1122        print(start_index,end_index)123        if start_index != -1 and end_index != 0:124            return generated_query[start_index:end_index]125        else:126            return generated_query127        128 129 130    def execute_query(self, query):131        """Executing the query on database and returning rows and columns."""132        print(query)133        cur = self.conn.cursor()134        cur.execute(query)135        col = [header[0] for header in cur.description]136        dash = "-" * sum(len(col_name) + 4 for col_name in col)137        print(tuple(col))138        print(dash)139        rows = []140        for member in cur:141            rows.append(member)142            print(member)143        cur.close()144        self.conn.commit()145        # print(rows)146        return rows, col147 148if __name__ == "__main__":149    model_dir = "multi_table_demo/checkpoint-2600"150    database = r"ticket_dataset.db"151    sql_model = SQLPromptModel(model_dir, database)152    user_prompt = "Give complete details of properties in India"153    while True:154        table_schema = sql_model.fetch_table_schema("ticket_dataset")155        if table_schema:156            # query = sql_model.text2sql(table_schema, user_prompt)157            # query = sql_model.text2sql_chatgpt(table_schema, user_prompt)158            query = sql_model.text2sql_gemini(table_schema, user_prompt)159            print(query)160            sql_model.execute_query(query)161            162    sql_model.conn.close()163