Chat2DB

/

Chat2DB-SQL-7B

like 38

Follow

Chat2DB 4

Setting pad_token_id to eos_token_id:2 for open-end generation.

1. 需要安装依赖包

!pip install accelerate
# 这个包需要重启内核生效
import os
os._exit(00)

2. 示例代码中多了一个空行，会导致语法错误。

3. 调用 pipeline 的时候需要加个参数：
   pipe(prompt) 改为 pipe(prompt,pad_token_id=pipe.tokenizer.eos_token_id)
   否则会报错

Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.

4. 用示例里面给的提示词是生成不了SQL的，里面的建表语句不知道为啥搞个方括号还有单引号括起来，也不像JSON语法，
   我把这个方括号还有单引号删了就能生成SQL了。

总体看来开发者不怎么用心，随便搞了个说明应付了事。

以下是修改后的代码，供参考

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
model_path = "Chat2DB/Chat2DB-SQL-7B" # This can be replaced with your local model path
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", trust_remote_code=True, torch_dtype=torch.float16, use_cache=True)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, return_full_text=False, max_new_tokens=100)
prompt = """
### Database Schema

CREATE TABLE "stadium" (
"Stadium_ID" int,
"Location" text,
"Name" text,
"Capacity" int,
"Highest" int,
"Lowest" int,
"Average" int,
PRIMARY KEY ("Stadium_ID")
);

CREATE TABLE "singer" (
"Singer_ID" int,
"Name" text,
"Country" text,
"Song_Name" text,
"Song_release_year" text,
"Age" int,
"Is_male" bool,
PRIMARY KEY ("Singer_ID")
);

CREATE TABLE "concert" (
"concert_ID" int,
"concert_Name" text,
"Theme" text,
"Stadium_ID" text,
"Year" text,
PRIMARY KEY ("concert_ID"),
FOREIGN KEY ("Stadium_ID") REFERENCES "stadium"("Stadium_ID")
);

CREATE TABLE "singer_in_concert" (
"concert_ID" int,
"Singer_ID" text,
PRIMARY KEY ("concert_ID","Singer_ID"),
FOREIGN KEY ("concert_ID") REFERENCES "concert"("concert_ID"),
FOREIGN KEY ("Singer_ID") REFERENCES "singer"("Singer_ID")
);


### Task 

Based on the provided database schema information, How many singers do we have?[SQL]

"""
response = pipe(prompt,pad_token_id=pipe.tokenizer.eos_token_id)[0]["generated_text"]
print(response)

thanks， i try it and get success.