Transformers部署sqlcoder-70b-alpha实战：环境搭建、模型下载与文本到SQL推理测试

发表于 | 分类于 | 753 | 4 分钟

一、sqlcoder-70b-alpha测试

sqlcoder环境安装

text
1
2
3
4
5
6
git clone https://github.com/defog-ai/sqlcoder.git
conda create -n sqlcoder_py311 python=3.11
pip install -r requirements.txt

pip install accelerate
pip install bitsandbytes

模型下载

text
1
2
https://hf-mirror.com/defog/sqlcoder-70b-alpha
https://huggingface.co/defog/sqlcoder-70b-alpha

修改inference.py文件

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
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import argparse


def generate_prompt(question, prompt_file="prompt.md", metadata_file="metadata.sql"):
    with open(prompt_file, "r") as f:
        prompt = f.read()

    with open(metadata_file, "r") as f:
        table_metadata_string = f.read()

    prompt = prompt.format(
        user_question=question, table_metadata_string=table_metadata_string
    )
    return prompt


def get_tokenizer_model(model_name):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        # torch_dtype=torch.float16,
        load_in_8bits=True,
        device_map="auto",
        use_cache=True,
    )
    return tokenizer, model


def run_inference(question, prompt_file="prompt.md", metadata_file="metadata.sql"):
    tokenizer, model = get_tokenizer_model("/mnt/lustre/data/sqlcoder-70b-alpha")
    prompt = generate_prompt(question, prompt_file, metadata_file)

    # make sure the model stops generating at triple ticks
    # eos_token_id = tokenizer.convert_tokens_to_ids(["```"])[0]
    eos_token_id = tokenizer.eos_token_id
    pipe = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=300,
        do_sample=False,
        return_full_text=False,  # added return_full_text parameter to prevent splitting issues with prompt
        num_beams=5,  # do beam search with 5 beams for high quality results
    )
    generated_query = (
            pipe(
                prompt,
                num_return_sequences=1,
                eos_token_id=eos_token_id,
                pad_token_id=eos_token_id,
            )[0]["generated_text"]
            .split(";")[0]
            .split("```")[0]
            .strip()
            + ";"
    )
    return generated_query


if __name__ == "__main__":
    # Parse arguments
    _default_question = "Do we get more sales from customers in New York compared to customers in San Francisco? Give me the total sales for each city, and the difference between the two."
    parser = argparse.ArgumentParser(description="Run inference on a question")
    parser.add_argument("-q", "--question", type=str, default=_default_question, help="Question to run inference on")
    args = parser.parse_args()
    question = args.question
    print("Loading a model and generating a SQL query for answering your question...")
    print(f"Answer:{run_inference(question)}")

执行脚本

text
1
python inference.py -q "Question about the sample database goes here"

模型部署返回结果

按照官网示例，部署（4卡）的int8量化后的，显存占用81.28GB，进行一次问答需要15分钟（70b效果个人感觉还行，但耗时过长）

sqlcoder接口部署

修改sqlcoder_70b_alpha.py

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
import argparse
import json
import os
import sys

import torch
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM

os.environ['CUDA_VISIBLE_DEVICES'] = "2,3,4,5"

import uvicorn
from fastapi import FastAPI
from loguru import logger

project_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
sys.path.insert(0, project_path)

from config.message import SqlCoderRequest
from config.paths import metadata_sql_path
from utils.time_util import time_it

system_prompt = """
### Task
Generate a SQL query to answer [QUESTION]{user_question}[/QUESTION]

### Instructions
- If you cannot answer the question with the available database schema, return 'I do not know'

### Database Schema
The query will run on a database with the following schema:
{table_metadata_string}

### Answer
Given the database schema, here is the SQL query that answers [QUESTION]{user_question}[/QUESTION]
[SQL]
"""


def generate_prompt(question, metadata_file=metadata_sql_path):
    with open(metadata_file, "r") as f:
        table_metadata_string = f.read()

    prompt = system_prompt.format(
        user_question=question, table_metadata_string=table_metadata_string
    )
    return prompt


def start_llm_service(args):
    """
    启动服务
    :param args:
    :return:
    """
    chat_app = FastAPI()

    model_path = "/mnt/lustre/data/sqlcoder-70b-alpha/"
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        trust_remote_code=True,
        # torch_dtype=torch.float16,
        load_in_8bit=True,
        device_map="auto",
        use_cache=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    eos_token_id = tokenizer.eos_token_id
    pipeline = transformers.pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=300,
        do_sample=False,
        return_full_text=False,  # added return_full_text parameter to prevent splitting issues with prompt
        num_beams=3,  # do beam search with 5 beams for high quality results
    )

    @time_it
    @chat_app.post("/chat")
    def chat(request: SqlCoderRequest):
        ans = None
        state = True
        try:
            request = json.loads(request.json())
            prompt = generate_prompt(request.get("prompt"), metadata_sql_path)
            logger.info("Receive message: {} ".format(prompt))
            ans = (
                    pipeline(
                        prompt,
                        num_return_sequences=1,
                        eos_token_id=eos_token_id,
                        pad_token_id=eos_token_id,
                    )[0]["generated_text"]
                    .split(";")[0]
                    .split("```")[0]
                    .strip()
                    + ";"
            )
            logger.info("Sending message: {}".format(ans))
        except Exception as e:
            state = False
            logger.error(e)
        return {
            "state": state,
            "response": ans
        }

    uvicorn.run(chat_app,
                host=args.host,
                port=int(args.port))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="sqlcoder_70b_alpha server.")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, required=True)
    start_llm_service(parser.parse_args())

执行脚本

text
1
python sqlcoder_70b_alpha.py --port 9002

启动服务

模型服务调用

服务地址：https://10.103.190.9:9002/chat
入参：

1
2
3
{
	"prompt": "What was our revenue by product in the New York region last month?"
}

postman方式调用

日志返回

本文结束 感谢您的阅读