Upload 19 files

app.py

@@ -0,0 +1,189 @@
+import gradio as gr
+from PIL import Image, ImageFilter
+import os 
+from txt2sql_code3 import SQLPromptModel
+from qa_bot_chatgpt import QAInfer
+from gradio import Button
+import time
+import google.generativeai as genai
+
+image = Image.open(os.path.join(os.path.abspath(''), "house_excel_sheet.png"))
+
+def image_display(steps=0):
+    return image
+
+query = None
+rows = None
+columns = None
+user_choices = None
+data = [
+    [1, "Buckingham Palace", 27, 34, 0.12, "London", "United Kingdom", 601, 920, 105000, "Buckingham_palace.pdf"],
+    [2, "White House", 355, 67, 0.08, "Washington D.C.", "United States", 601, 1527, 80000, "White_House.pdf"],
+    [3, "Taj Mahal Palace", 455, 76, 0.15, "Mumbai", "India", 795, 748, 67000, "Taj_Mahal_palace.pdf"],
+    [4, "Versailles Palace", 455, 45, 0.1, "Versailles", "France", 731, 1800, 145000, "Palace_of_Versailes.pdf"],
+    [5, "Villa Leopolda", 223, 21, 0.05, "Villefranche-sur-Mer", "France", 680, 6886, 65000, "Villa_Leopolda.pdf"],
+    [6, "Antilia", 455, 70, 0.46, "Mumbai", "India", 612, 2520, 179000, "Antilia.pdf"],
+    [7, "The Biltmore Estate", 544, 93, 0.2, "Asheville", "United States", 639, 2040, 50000, "Biltmore_Estate.pdf"],
+    [8, "Hearst Castle", 303, 57, 0.08, "San Simeon", "United States", 731, 1050, 71600, "Hearst_castle.pdf"],
+    [9, "Villa Les Cèdres", 489, 88, 0.09, "Saint-Jean-Cap-Ferrat", "France", 730, 1092, 100000, "Villa_Les_Cedres.pdf"],
+    [10, "Istana Nurul Iman", 350, 34, 7.46, "Bandar Seri Begawan", "Brunei", 670, 5403, 300000, "Istana_Nurul_Iman.pdf"]
+]
+choices = [item[1] for item in data]
+
+
+
+def execute_sql_query(input_prompt):
+    global query, rows, columns, user_choices
+    
+    model_dir = "multi_table_demo/checkpoint-2600"
+    database = r"sql_pdf.db"
+    sql_model = SQLPromptModel(model_dir, database)
+
+    user_prompt = "Give complete details of properties in India"
+    for _ in range(3):  # Retry logic, try 3 times
+        try:
+            table_schema = sql_model.fetch_table_schema("sql_pdf")
+            if table_schema:
+                if input_prompt.strip():  
+                    # query = sql_model.text2sql(table_schema, user_prompt, input_prompt)
+                    # query = sql_model.text2sql_chatgpt(table_schema, user_prompt, input_prompt)
+                    query = sql_model.text2sql_gemini(table_schema, user_prompt, input_prompt)
+                else:
+                    # query = sql_model.text2sql(table_schema, user_prompt, user_prompt)
+                    # query = sql_model.text2sql_chatgpt(table_schema, user_prompt, user_prompt)
+                    query = sql_model.text2sql_gemini(table_schema, user_prompt, user_prompt)
+                rows, columns = sql_model.execute_query(query)
+                print(rows)
+                user_choices = []
+                for item in rows:
+                    if isinstance(item[0], str):
+                        user_choices.append(item[0])
+                    else:
+                        for val in item:
+                            if isinstance(val, str):
+                                user_choices.append(val)
+                                break
+                return rows
+            else:
+                print("Table schema not found.")
+                return None
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            print("Retrying...")
+            time.sleep(1)  # Wait for 1 second before retrying
+    return None  # Return None if all retries fail
+
+
+# def qa_infer_interface(row,query_on_pdf):
+#     qa_infer=QAInfer()
+#     qa_infer.qa_infer(query,rows,columns)
+user_choices = None
+def update_choices(nothing):
+    print("callback called")
+    print("choices",choices)
+    print("user_choices",user_choices)
+    if user_choices:
+        examples = [[user_choices[0], "Structure of the property"], [user_choices[0], "Property History"], [user_choices[0], "How many floors does the property have"]]
+        return gr.Dropdown(choices=user_choices, label="Property Choice",info="List of all properties",interactive=True)#,examples
+    else:
+        return gr.Dropdown(choices=[], label="Property Choice",info="List of all properties",interactive=True)
+    
+def update_examples(nothing):
+    if user_choices:
+        examples = [[user_choices[0], "Structure of the property"], [user_choices[0], "Property History"], [user_choices[0], "How many floors does the property have"]]
+        return examples
+    else :
+        examples=[["","Structure of the property "],[ ""," Property History "] ,["", " How many floors does the property have"]],
+
+
+def qa_infer_interface(property_choice, query_question):
+    qa_infer = QAInfer()
+    if not property_choice and user_choices:
+        property_choice = user_choices[0]
+    property_row = [row for row in data if row[1] == property_choice][0]  # Find the row corresponding to the selected property
+    if not query_question:
+        query_question = "area"
+    retries = 3
+    while retries > 0:
+        try:
+            print(property_row)
+            # answer = qa_infer.qa_infer_interface(property_row, query_question) #-->chatgpt
+            answer = qa_infer.qa_infer_interface_gemini(property_row, query_question)
+            return answer
+        except Exception as e:
+            print(f"Error occurred while inferring QA: {e}")
+            retries -= 1
+    print("Failed to infer QA after 3 retries.")
+    return None
+
+
+
+user_dropdown=gr.Dropdown(choices=[], label="Property Choice",info="List of all properties")
+properties_text=gr.components.Textbox(lines=2,label="User Database Query",placeholder="Click on an query from 'examples' below or write your own query based on the database above. Default : 'Properties in India'")
+interface_1_output=gr.Json(label="json")
+stage2_examples=[["","Structure of the property "],[ ""," Property History "] ,["", " How many floors does the property have"]]
+stage2_text=gr.components.Textbox(lines=2,label="Question on property",placeholder="Enter a question to know more about the properties , you can choose from one of the options below or write a own question Default: 'Area of the property'",)
+stage2_output="text"
+
+with gr.Blocks(title="House Excel Query") as demo: 
+
+    gr.Markdown("# House Excel Query")
+
+    generated_image = image_display()
+    gr.Image(generated_image)
+
+    gr.Markdown("""### The database provided contains information about different properties, including their fundamental details. Additional specifics about each property are stored in associated PDF files, which are referenced in the "PDF" column. You have the capability to query this database using various criteria. When a query is initiated, the system generates SQL queries and extracts relevant rows from the database in the backend.
+                \n ### Once the properties are retrieved based on the query, you can utilize the user interface (UI) below to perform question answering (QA). Simply select a property from the list of returned properties and compose a question pertaining to that property. You will receive an answer based on the available information.""")
+
+    interface_1 = gr.Interface(
+        execute_sql_query,
+        inputs=properties_text,
+        # "textbox",
+        outputs=interface_1_output,
+        # live=True,
+        # cache_examples=["Give me all details of properties from India"],
+        examples=["Properties in France "," Properties greater than a acre","Properties with more than 400 bedrooms"],
+    )
+    # print(interface_1.input_components[0])
+    
+
+    interface_2 = gr.Interface(
+        qa_infer_interface,
+        inputs=[user_dropdown,stage2_text],
+        # inputs=[gr.Dropdown.change(fn=update_choices),gr.components.Textbox(lines=2,label="Question on property",placeholder="Enter a question to know more about the properties")],
+        outputs=stage2_output,
+        # examples=stage2_examples,
+        # live=True,
+        # gr.Button("Next"),
+        # Button.click(next,value="Next"),
+        
+    )
+
+    gr.Examples(["How many floors does the property have "," Total square feet of the property " ," Total area of the property"],inputs=stage2_text,outputs=stage2_output,fn=qa_infer_interface)
+    
+
+    properties_text.change(update_choices,inputs=[properties_text],outputs=[user_dropdown])
+    interface_1_output.change(update_choices,inputs=[interface_1_output],outputs=[user_dropdown])
+
+
+    # user_dropdown.change(update_examples, inputs=[user_dropdown], outputs=[stage2_examples])
+
+    # properties_text.change(update_choices,inputs=[stage2_examples],outputs=[interface_2.examples])
+    # interface_1_output.change(update_choices,inputs=[stage2_examples],outputs=[interface_2.examples])
+
+    # user_dropdown.change(update_choices, inputs=[user_dropdown], outputs=[user_dropdown, interface_2])
+
+    # user_dropdown.change(fn=update_choices,inputs=[user_dropdown],outputs=[user_dropdown])
+
+    # with gr.Row():
+    #             save_btn = gr.Button("Next")
+
+    # Button.click(next,value="Next",),
+
+if __name__ == "__main__":
+    demo.launch(share=True)
+
+
+## download pdf buttons
+## upload pdf 
+## dynamic selection

app2.py

@@ -0,0 +1,7 @@
+import gradio as gr
+
+def greet(name):
+    return "Hello! " + name + "api" + api
+
+iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+iface.launch()

qa_bot_chatgpt.py

@@ -0,0 +1,148 @@
+import sqlite3
+from sqlite3 import Error
+from PyPDF2 import PdfReader
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
+import os
+import torch
+from huggingface_hub import login
+import ast
+from openai import OpenAI
+import google.generativeai as genai
+class QAInfer:
+    def __init__(self):
+        torch.cuda.empty_cache()
+       
+        # self.chatgpt_client = OpenAI(api_key="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")
+
+        self.chatgpt_client = OpenAI(api_key="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")
+        self.genai = genai
+        self.genai.configure(api_key="AIzaSyAFG94rVbm9eWepO5uPGsMha8XJ-sHbMdA")
+        self.genai_model = genai.GenerativeModel('gemini-pro')
+
+    def extract_text_from_pdf(self, pdf_path):
+        """Extract text from a PDF file."""
+        reader = PdfReader(pdf_path)
+        text = ''
+        for page in reader.pages:
+                text += page.extract_text()
+        return text
+    
+    def qa_infer(self, query, rows, col):
+        """QA inference function."""
+        print(query)
+       
+        print(tuple(col))   
+        file_index = -1
+
+        if "additional_info" not in col:
+            pass
+        else:
+            file_index = [i for i in range(len(col)) if col[i] == "additional_info"][0]
+            initiate_qa = input("\nDo you wish to ask questions about the properties [y/n]?: ").lower()
+
+            if initiate_qa in ['y', 'yes']:
+                for row in rows:
+                    pdf_text = self.extract_text_from_pdf(row[file_index])
+                    print("Extracted text from PDF", os.path.basename(row[file_index]))
+
+                    while True:
+                        user_question = input("\nWhat do you want to know about this property? (Press Enter to exit): ").strip()
+                        if not user_question:
+                            break
+                        # Construct QA prompt directly here
+                        question = user_question if user_question else "Who is lashkar e taiba"
+                        prompt = f"""Below is a question and context, search the context to find the answer for the question and return the response ###question:{question} ###context:{pdf_text} ###response:"""
+
+                        # Run the language model to generate a response
+                        inputs = self.qa_tokenizer(prompt, return_tensors='pt', truncation=True, max_length=512)
+                        pipe = pipeline(
+                            "text-generation", 
+                            model=self.qa_model, 
+                            tokenizer=self.qa_tokenizer, 
+                            torch_dtype=torch.bfloat16, 
+                            device_map="auto"
+                        )
+                        sequences = pipe(
+                            prompt,
+                            do_sample=True,
+                            max_new_tokens=200, 
+                            temperature=0.7, 
+                            top_k=50, 
+                            top_p=0.95,
+                            num_return_sequences=1,
+                        )
+                        answer = sequences[0]['generated_text']
+                        print("Answer:", answer)
+            else:
+                continue_to_next = input("Do you want to continue with the next property? [y/n]: ").lower()
+                if continue_to_next != 'y':
+                    return
+    
+    def qa_infer_interface(self, row, query_question):
+        """This method is used for gradio interface only"""
+        file_path = row[-1]  # Assuming the last element in row contains the PDF file path
+        pdf_text = self.extract_text_from_pdf(file_path)
+        
+        # prompt = f"""Below is a question and context, search the context to find the answer for the question and return the response , if related answer cannot be found return "Answer not in the context" ###question:{query_question} ###context:{pdf_text} ###response:"""
+        prompt = f"""You have been provided with a question and a corresponding context. Your task is to search the context to find the answer to the question. If the answer is found, return the response. If the answer cannot be found in the context, please respond with "Answer not found in the context".
+
+        === Question ===
+        {query_question}
+
+        === Context ===
+        {pdf_text}
+
+        === Response ===
+         Try mostly to answer from given pdf , if related answer is not found return 'Information not present in the pdf' and below it provide something related to the question .Note: return only answer dont include terms like 'Response','###','Answer'"""
+
+    
+
+        print(prompt)
+        completion = self.chatgpt_client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[
+                {"role": "system", "content": "You are a expert PDF parser , go through the pdf and answer the question properly , if related answer is not found return 'Information not present in the pdf' and below it provide something related to the question"},
+                {"role": "user", "content": prompt }
+            ]
+        )
+        return completion.choices[0].message.content
+
+    def qa_infer_interface_gemini(self, row, query_question):
+        """This method is used for gradio interface only"""
+        file_path = row[-1]  # Assuming the last element in row contains the PDF file path
+        pdf_text = self.extract_text_from_pdf(file_path)
+        
+        # prompt = f"""Below is a question and context, search the context to find the answer for the question and return the response , if related answer cannot be found return "Answer not in the context" ###question:{query_question} ###context:{pdf_text} ###response:"""
+        prompt = f"""You have been provided with a question and a corresponding context. Your task is to search the context to find the answer to the question. If the answer is found, return the response. If the answer cannot be found in the context, please respond with "Answer not found in the context".
+
+        === Question ===
+        {query_question}
+
+        === Context ===
+        {pdf_text}
+
+        === Response ===
+         If related answer is not found return 'Information not present in the pdf' and below it provide something related to the question"""
+
+        print(prompt)
+        completion = self.genai_model.generate_content(prompt)
+        generated_answer=completion.text
+
+        return generated_answer
+
+
+if __name__ == '__main__':
+    qa_infer = QAInfer()
+    query = 'SELECT * FROM sql_pdf WHERE country = "India" '
+    rows = [
+        (3, 'Taj Mahal Palace', 455, 76, 0.15, 'Mumbai', 'India', 795, 748, 67000, 'pdf_files/pdf/Taj_Mahal_palace.pdf'), 
+        (6, 'Antilia', 455, 70, 0.46, 'Mumbai', 'India', 612, 2520, 179000, 'pdf_files/pdf/Antilia.pdf')
+    ]
+    col = [
+        "property_id", "name", "bed", "bath", "acre_lot", "city", "country", 
+        "zip_code", "house_size", "price", "additional_info"
+    ]
+    qa_infer.qa_infer(query, rows, col)
+
+
+

requirements.txt

@@ -0,0 +1,15 @@
+bitsandbytes
+git+https://github.com/huggingface/transformers.git
+git+https://github.com/huggingface/peft.git
+git+https://github.com/huggingface/accelerate.git
+datasets
+evaluate
+trl==0.7.1
+jupyter
+scipy
+gradio
+python-dotenv
+openpyxl
+PyPDF2
+llama-parse
+google-generativeai

sql.py

@@ -0,0 +1,128 @@
+import sqlite3
+from sqlite3 import Error
+import csv
+import pandas as pd
+import os
+
+
+def create_connection(db_file):
+    """ create a database connection to a database that resides
+        in the memory
+    """
+    conn = None;
+    try:
+        conn = sqlite3.connect(db_file)
+        return conn
+    
+    except Error as e:
+        print(e)
+
+    return conn
+
+
+def create_table(conn, create_table_sql, table_name):
+    """ create a table from the create_table_sql statement
+    :param conn: Connection object
+    :param create_table_sql: a CREATE TABLE statement
+    :return:
+    """
+    try:
+        c = conn.cursor()
+        c.execute(f"""DROP TABLE IF EXISTS {table_name}""")
+        c.execute(create_table_sql)
+    except Error as e:
+        print(e)
+
+
+def insert_values(conn, task, sql):
+
+    cur = conn.cursor()
+    cur.execute(sql, task)
+    conn.commit()
+    return cur.lastrowid
+
+
+def populate(csv_file, db_file, table_insert):
+    
+    conn = create_connection(db_file)
+    with conn:
+        
+        
+        with open(csv_file, mode ='r')as file:
+            csvfile = csv.reader(file)
+
+            for n,lines in enumerate(csvfile):
+                
+                if n>0:
+                    lines = tuple(i for i in lines)
+                    insert_values(conn, lines, table_insert)
+
+                else:pass
+
+def main():
+
+    name = "sql_pdf.xlsx"
+    excel_file = (pd.read_excel(name))
+    csv_file = f"""{name.split(".")[0]}.csv"""
+    excel_file.to_csv(csv_file,
+                    index=None,
+                    header=True)
+    column = [x for x in excel_file.columns]
+    column_type = {}
+    type_map = {
+        "<class 'str'>": "TEXT",
+        "<class 'int'>": "INTEGER",
+        "<class 'float'>": "REAL",
+    }
+
+
+    for i in range(len(column)):
+        datatype = {}
+        for j in excel_file.values:
+            if type(j[i]) not in list(datatype.keys()):datatype[type(j[i])] = 1
+            else: datatype[type(j[i])] += 1
+        
+        ma_x = 0
+        max_type = "<class 'str'>"
+        
+        for k in list(datatype.keys()):
+            if ma_x < datatype[k]:max_type = str(k)
+
+        try:
+            column_type[column[i]] = type_map[max_type]
+
+        except KeyError:
+            column_type[column[i]] = "TEXT"
+
+    print(column_type)
+
+    table_construct = f"""CREATE TABLE IF NOT EXISTS {name.split(".")[0]}( """
+    table_insert = f"""INSERT INTO {name.split(".")[0]}("""
+    table_values = f"""VALUES ("""
+    for l in list(column_type.keys()):
+        table_construct += f"""{l} {column_type[l]}, """
+        table_insert += f"""{l}, """
+        table_values += "?, "
+    table_construct = f"""{table_construct[:-2]});"""
+    table_values = f"""{table_values[:-2]})"""
+    table_insert = f"""{table_insert[:-2]})\n{table_values}"""
+
+    print(table_construct)
+    print("\n\n", table_insert)
+
+    database = f"""{name.split(".")[0]}.db"""
+    
+    conn = create_connection(database)
+
+    # create tables
+    if conn is not None:
+        # create projects table
+        create_table(conn, table_construct, name.split(".")[0])
+    else:
+        print("Error! cannot create the database connection.")
+
+
+    populate(csv_file, database, table_insert)
+
+if __name__ == '__main__':
+    main()

txt2sql_code3.py

@@ -0,0 +1,158 @@
+import sqlite3
+from sqlite3 import Error
+from peft import AutoPeftModelForCausalLM
+from transformers import AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM
+from openai import OpenAI
+import google.generativeai as genai
+
+class SQLPromptModel:
+    def __init__(self, model_dir, database):
+        self.model_dir = model_dir
+        self.database = database
+        # peft_model_dir = self.model_dir
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype="float16",
+            bnb_4bit_use_double_quant=True,
+        )
+        # self.model = AutoPeftModelForCausalLM.from_pretrained(
+        #     peft_model_dir, low_cpu_mem_usage=True, quantization_config=bnb_config
+        # )
+        # self.tokenizer = AutoTokenizer.from_pretrained(peft_model_dir)
+        # self.model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
+        # self.tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
+        self.chatgpt_client = OpenAI(api_key="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")
+        self.genai = genai
+        self.genai.configure(api_key="AIzaSyAFG94rVbm9eWepO5uPGsMha8XJ-sHbMdA")
+        self.genai_model = genai.GenerativeModel('gemini-pro')
+
+        self.conn = sqlite3.connect(self.database)
+
+    def fetch_table_schema(self, table_name):
+        """Fetch the schema of a table from the database."""
+        cursor = self.conn.cursor()
+        cursor.execute(f"PRAGMA table_info({table_name})")
+        schema = cursor.fetchall()
+        if schema:
+            return schema
+        else:
+            print(f"Table {table_name} does not exist or has no schema.")
+            return None
+        
+    def text2sql(self, schema, user_prompt, inp_prompt=None):
+        """Generate SQL query based on user prompt and table schema.inp_prompt is for gradio purpose"""
+        table_columns = ', '.join([f"{col[1]} {col[2]}" for col in schema])
+
+        prompt = f"""Below are SQL table schemas paired with instructions that describe a task.
+        Using valid SQLite, write a response that appropriately completes the request for the provided tables.
+        ### Instruction: {user_prompt} ### 
+        Input: CREATE TABLE sql_pdf({table_columns});
+        ### Response: (Return only query , nothing extra)"""
+
+        if inp_prompt is not None :
+            prompt = prompt.replace(user_prompt, inp_prompt + " ")
+        else:
+            inp_prompt = input("Press Enter for default question or Enter user prompt without newline characters: ").strip()
+            if inp_prompt:
+                prompt = prompt.replace(user_prompt, inp_prompt + " ")
+
+        """Text to SQL query generation"""
+        input_ids = self.tokenizer(
+            prompt, return_tensors="pt", truncation=True
+        ).input_ids.to(next(self.model.parameters()).device)  # Move input to the device of the model
+        outputs = self.model.generate(input_ids=input_ids, max_new_tokens=200)
+        response = self.tokenizer.batch_decode(
+            outputs.detach().cpu().numpy(), skip_special_tokens=True
+        )[0][:]
+        return response[len(prompt):]
+
+    def text2sql_chatgpt(self, schema, user_prompt, inp_prompt=None):
+        table_columns = ', '.join([f"{col[1]} {col[2]}" for col in schema])
+
+        prompt = f"""Below are SQL table schemas paired with instructions that describe a task.
+        Using valid SQLite, write a response that appropriately completes the request for the provided tables.
+        ### Instruction: {user_prompt} ### 
+        Input: CREATE TABLE sql_pdf({table_columns});
+        ### Response: (Return only generated query based on user_prompt , nothing extra)"""
+
+        if inp_prompt is not None :
+            prompt = prompt.replace(user_prompt, inp_prompt + " ")
+        else:
+            inp_prompt = input("Press Enter for default question or Enter user prompt without newline characters: ").strip()
+            if inp_prompt:
+                prompt = prompt.replace(user_prompt, inp_prompt + " ")
+        print(prompt)
+        completion = self.chatgpt_client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[
+                {"role": "system", "content": "You are a expert SQL developer , generate a sql query and return it"},
+                {"role": "user", "content": prompt }
+            ]
+        )
+        return completion.choices[0].message.content
+
+    def text2sql_gemini(self, schema, user_prompt, inp_prompt=None):
+        table_columns = ', '.join([f"{col[1]} {col[2]}" for col in schema])
+
+        prompt = f"""Below are SQL table schemas paired with instructions that describe a task.
+        Using valid SQLite, write a response that appropriately completes the request for the provided tables.
+        ### Instruction: {user_prompt} ### 
+        Input: CREATE TABLE sql_pdf({table_columns});
+        ### Response: (Return only generated query based on user_prompt , nothing extra)"""
+
+        if inp_prompt is not None :
+            prompt = prompt.replace(user_prompt, inp_prompt + " ")
+        else:
+            inp_prompt = input("Press Enter for default question or Enter user prompt without newline characters: ").strip()
+            if inp_prompt:
+                prompt = prompt.replace(user_prompt, inp_prompt + " ")
+        print(prompt)
+        completion = self.genai_model.generate_content(prompt)
+        generated_query=completion.text
+        start_index = generated_query.find("SELECT")
+        end_index = generated_query.find(";", start_index) + 1
+        print(start_index,end_index)
+        if start_index != -1 and end_index != 0:
+            return generated_query[start_index:end_index]
+        else:
+            return generated_query
+        
+
+
+    def execute_query(self, query):
+        """Executing the query on database and returning rows and columns."""
+        print(query)
+        cur = self.conn.cursor()
+        cur.execute(query)
+        col = [header[0] for header in cur.description]
+        dash = "-" * sum(len(col_name) + 4 for col_name in col)
+        print(tuple(col))
+        print(dash)
+        rows = []
+        for member in cur:
+            rows.append(member)
+            print(member)
+        cur.close()
+        self.conn.commit()
+        # print(rows)
+        return rows, col
+
+if __name__ == "__main__":
+    model_dir = "multi_table_demo/checkpoint-2600"
+    database = r"sql_pdf.db"
+    sql_model = SQLPromptModel(model_dir, database)
+    user_prompt = "Give complete details of properties in India"
+    while True:
+        table_schema = sql_model.fetch_table_schema("sql_pdf")
+        if table_schema:
+            # query = sql_model.text2sql(table_schema, user_prompt)
+            # query = sql_model.text2sql_chatgpt(table_schema, user_prompt)
+            query = sql_model.text2sql_gemini(table_schema, user_prompt)
+            print(query)
+            sql_model.execute_query(query)
+            
+    sql_model.conn.close()
+
+