systex_phase2/text_to_sql.py

import re
from dotenv import load_dotenv
load_dotenv()

from langchain_community.utilities import SQLDatabase
import os
URI: str =  os.environ.get('SUPABASE_URI')
db = SQLDatabase.from_uri(URI)

# print(db.dialect)
# print(db.get_usable_table_names())
# db.run('SELECT * FROM "2022 清冊數據(GHG)" LIMIT 10;')

context = db.get_context()
# print(list(context))
# print(context["table_info"])

from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate
from langchain.chains import create_sql_query_chain
from langchain_community.llms import Ollama

from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool
from operator import itemgetter

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough

# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer,pipeline
import torch
from langchain_huggingface import HuggingFacePipeline

# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLM
# model_id = "defog/llama-3-sqlcoder-8b"
# tokenizer = AutoTokenizer.from_pretrained(model_id)
# sql_llm = HuggingFacePipeline.from_model_id(
#     model_id=model_id,
#     task="text-generation",
#     model_kwargs={"torch_dtype": torch.bfloat16},
#     pipeline_kwargs={"return_full_text": False},
#     device=0, device_map='cuda')


model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)

llm = HuggingFacePipeline.from_model_id(
    model_id=model_id,
    task="text-generation",
    model_kwargs={"torch_dtype": torch.bfloat16},
    pipeline_kwargs={"return_full_text": False,
        "max_new_tokens": 512},
    device=0, device_map='cuda')
print(llm.pipeline)
llm.pipeline.tokenizer.pad_token_id = llm.pipeline.model.config.eos_token_id[0]
# model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True)

# pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=500, top_k=50, temperature=0.1, 
#                 model_kwargs={"torch_dtype": torch.bfloat16, "return_full_text": False})
#, device="auto", load_in_4bit=True
# llm = HuggingFacePipeline(pipeline=pipe)

# llm = HuggingFacePipeline(pipeline=pipe)

# llm = Ollama(model = "llama3-groq-tool-use:latest", num_gpu=1)
def get_examples():
    examples = [
        {
            "input": "去年的固定燃燒總排放量是多少?",
            "query": 'SELECT SUM("高雄總部及運通廠" + "台北辦事處" + "昆山廣興廠" + "北海建準廠" + "北海立準廠" + "菲律賓建準廠" + "Inc" + "SAS" + "India") AS "固定燃燒總排放量"\nFROM "2023 清冊數據(GHG)"\nWHERE "排放源" = \'固定燃燒\'',
        },
        {
            "input": "建準廣興廠去年的類別1總排放量是多少?",
            "query": 'SELECT SUM("昆山廣興廠") AS "建準廣興廠類別1總排放量"\nFROM "2023 清冊數據(GHG)"\nWHERE "類別" like \'%類別1%\'',
        },
        {
            "input": "建準廣興廠去年的直接排放總排放量是多少?",
            "query": 'SELECT SUM("昆山廣興廠") AS "建準廣興廠直接排放總排放量"\nFROM "2023 清冊數據(GHG)"\nWHERE "類別" like \'%直接排放%\'',
        },

    ]

    return examples

def write_query_chain(db):

    template = """<|begin_of_text|><|start_header_id|>user<|end_header_id|>

    Generate a SQL query to answer this question: `{input}`

    You are a PostgreSQL expert in ESG field. Given an input question, first create a syntactically correct PostgreSQL query to run, 
    then look at the results of the query and return the answer to the input question.\n\
    Unless the user specifies in the question a specific number of examples to obtain, query for at most {top_k} results using the LIMIT clause as per PostgreSQL. 
    You can order the results to return the most informative data in the database.\n\
    Never query for all columns from a table. You must query only the columns that are needed to answer the question. 
    Wrap each column name in  Quotation Mark (") to denote them as delimited identifiers.\n\
    
    ***Pay attention to only return query for PostgreSQL WITHOUT "```sql", And DO NOT content any other words.\n\
    ***Pay attention to only return PostgreSQL query.\n\

    DDL statements:
    {table_info}<|eot_id|><|start_header_id|>assistant<|end_header_id|>

    The following SQL query best answers the question `{input}`:
    ```sql
    """
    # prompt_template = PromptTemplate.from_template(template)

    example_prompt = PromptTemplate.from_template("User input: {input}\nSQL query: {query}")
    prompt = FewShotPromptTemplate(
        examples=get_examples(),
        example_prompt=example_prompt,
        prefix=template,
        suffix="User input: {input}\nSQL query: ",
        input_variables=["input", "top_k", "table_info"],
    )

    # llm = Ollama(model = "mannix/defog-llama3-sqlcoder-8b", num_gpu=1)
    # llm = HuggingFacePipeline(pipeline=pipe)
    
    
    write_query = create_sql_query_chain(llm, db, prompt)


    return write_query

def sql_to_nl_chain():
    # llm = Ollama(model = "llama3.1", num_gpu=1)
    # llm = Ollama(model = "llama3.1:8b-instruct-q2_K", num_gpu=1)
    # llm = Ollama(model = "llama3-groq-tool-use:latest", num_gpu=1)
    answer_prompt = PromptTemplate.from_template(
        """
        <|begin_of_text|><|start_header_id|>system<|end_header_id|>
        Given the following user question, corresponding SQL query, and SQL result, answer the user question.
        給定以下使用者問題、對應的 SQL 查詢和 SQL 結果，以繁體中文回答使用者問題。

        For example
        Question: 建準廣興廠去年的類別1總排放量是多少?
        SQL Query: SELECT SUM("昆山廣興廠") AS "建準廣興廠類別1總排放量"\nFROM "2023 清冊數據(GHG)"\nWHERE "類別" like \'%類別1%\'
        SQL Result: [(1102.3712,)]
        Answer: 建準廣興廠去年的類別1總排放量是1102.3712

        Question: {question}
        SQL Query: {query}
        SQL Result: {result}
        Answer: """
        )

    chain = answer_prompt | llm | StrOutputParser()

    return chain

def run(db, question, selected_table):

    write_query = write_query_chain(db)
    query = write_query.invoke({"question": question, 'table_names_to_use': selected_table, "top_k": 1000, "table_info":context["table_info"]})
    
    query = re.split('SQL query: ', query)[-1]
    print(query)

    execute_query = QuerySQLDataBaseTool(db=db)
    result = execute_query.invoke(query)
    print(result)

    chain = sql_to_nl_chain()
    answer = chain.invoke({"question": question, "query": query, "result": result})

    return query, result, answer


if __name__ == "__main__":
    import time
    
    start = time.time()
    
    selected_table = ['2022 清冊數據(GHG)', '2022 清冊數據(ISO)', '2023 清冊數據(GHG)', '2023 清冊數據(ISO)']
    question = "去年的固定燃燒總排放量是多少?"
    query, result, answer = run(db, question, selected_table)
    print("query: ", query)
    print("result: ", result)
    print("answer: ", answer)
    
    print(time.time()-start)