File size: 6,348 Bytes
dc82c71
 
 
 
 
 
 
 
 
 
 
 
 
2a9178f
 
dc82c71
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
import os
import re
import torch
from transformers import AutoModel, AutoTokenizer
import gradio as gr
import mdtex2html
from transformers import AutoTokenizer, AutoModel
from utility.utils import config_dict
from utility.loggers import logger
from sentence_transformers import util
from local_database import db_operate
from prompt import table_schema, embedder,corpus_embeddings, corpus,In_context_prompt, query_template

tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4",trust_remote_code=True).float()
model = model.eval()


"""Override Chatbot.postprocess"""

def postprocess(self, y):
    if y is None:
        return []
    for i, (message, response) in enumerate(y):
        y[i] = (
            None if message is None else mdtex2html.convert((message)),
            None if response is None else mdtex2html.convert(response),
        )
    return y

gr.Chatbot.postprocess = postprocess

def parse_text(text):
    """copy from https://github.com/GaiZhenbiao/ChuanhuChatGPT/"""
    lines = text.split("\n")
    lines = [line for line in lines if line != ""]
    count = 0
    for i, line in enumerate(lines):
        if "```" in line:
            count += 1
            items = line.split('`')
            if count % 2 == 1:
                lines[i] = f'<pre><code class="language-{items[-1]}">'
            else:
                lines[i] = f'<br></code></pre>'
        else:
            if i > 0:
                if count % 2 == 1:
                    line = line.replace("`", "\`")
                    line = line.replace("<", "&lt;")
                    line = line.replace(">", "&gt;")
                    line = line.replace(" ", "&nbsp;")
                    line = line.replace("*", "&ast;")
                    line = line.replace("_", "&lowbar;")
                    line = line.replace("-", "&#45;")
                    line = line.replace(".", "&#46;")
                    line = line.replace("!", "&#33;")
                    line = line.replace("(", "&#40;")
                    line = line.replace(")", "&#41;")
                    line = line.replace("$", "&#36;")
                lines[i] = "<br>"+line
    text = "".join(lines)
    return text


def obtain_sql(response):
    response = re.split("```|\n\n", response)
    for text in response:
        if "SELECT" in text:
            response = text
            break
    else:
        response = response[0]
    response = response.replace("\n", " ").replace("``", "").replace("`", "").strip()
    response = re.sub(' +',' ', response)
    return response


def predict(input, chatbot, history):
    max_length = 2048
    top_p = 0.7
    temperature = 0.2
    top_k = 3
    dboperate = db_operate(config_dict['db_path'])
    logger.info(f"query:{input}")
    chatbot_prompt = """
你是一个文本转SQL的生成器,你的主要目标是尽可能的协助用户将输入的文本转换为正确的SQL语句。
上下文开始
生成的表名和表字段均来自以下表:
"""
    query_embedding = embedder.encode(input, convert_to_tensor=True) # 与6张表的表名和输入的问题进行相似度计算
    cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0] 
    top_results = torch.topk(cos_scores, k=top_k) # 拿到topk=3的表名
    # 组合Prompt
    table_nums = 0 
    for score, idx in zip(top_results[0], top_results[1]):
        # 阈值过滤
        if score > 0.45:
            table_nums += 1
            chatbot_prompt += table_schema[corpus[idx]]
        chatbot_prompt += "上下文结束\n"
    # In-Context Learning
    if table_nums >= 2 and not history: # 如果表名大于等于2个,且没有历史记录,就加上In-Context Learning
        chatbot_prompt += In_context_prompt
    #  加上查询模板
    chatbot_prompt += query_template
    query = chatbot_prompt.replace("<user_input>", input)
    chatbot.append((parse_text(input), ""))
    # 流式输出
    # for response, history in model.stream_chat(tokenizer, query, history, max_length=max_length, top_p=top_p,
    #                                            temperature=temperature):
    #     chatbot[-1] = (parse_text(input), parse_text(response))
    response, history = model.chat(tokenizer, query, history=history, max_length=max_length, top_p=top_p,temperature=temperature)
    chatbot[-1] = (parse_text(input), parse_text(response))
    # chatbot[-1] = (chatbot[-1][0], chatbot[-1][1])
    # 获取结果中的SQL语句
    response = obtain_sql(response)
    # 查询结果
    if "SELECT" in response:
        try:
            sql_stauts = "sql语句执行成功,结果如下:"
            sql_result = dboperate.query_data(response)
            sql_result = str(sql_result)
        except Exception as e:
            sql_stauts = "sql语句执行失败"
            sql_result = str(e)
        chatbot[-1] = (chatbot[-1][0], 
                       chatbot[-1][1] + "\n\n"+ "===================="+"\n\n" + sql_stauts + "\n\n" + sql_result)
    return chatbot, history


def reset_user_input():
    return gr.update(value='')


def reset_state():
    return [], []

with gr.Blocks() as demo:
    gr.HTML("""<h1 align="center">🤖ChatSQL</h1>""")

    chatbot = gr.Chatbot()
    with gr.Row():
        with gr.Column(scale=4):
            with gr.Column(scale=12):
                user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10).style(
                    container=False)
            with gr.Column(min_width=32, scale=1):
                submitBtn = gr.Button("Submit", variant="primary")
        with gr.Column(scale=1):
            emptyBtn = gr.Button("Clear History")
            # max_length = gr.Slider(0, 4096, value=2048, step=1.0, label="Maximum length", interactive=True)
            # top_p = gr.Slider(0, 1, value=0.7, step=0.01, label="Top P", interactive=True)
            # temperature = gr.Slider(0, 1, value=0.95, step=0.01, label="Temperature", interactive=True)

    history = gr.State([])

    submitBtn.click(predict, [user_input, chatbot, history], [chatbot, history],
                    show_progress=True)
    submitBtn.click(reset_user_input, [], [user_input])

    emptyBtn.click(reset_state, outputs=[chatbot, history], show_progress=True)

demo.queue().launch(share=False, inbrowser=True)