spider dev dataset evaluation result

[chuangchuang-wang]

i use example/sqlite.ipynb to test spider dev dataset, only get execute accuracy=0.5，is well below publish 0.75。Are there any test details I missed on spider

[senwu]

Please use torch.float16 or torch.bfloat16 and set max_new_tokens greater than 300 tokens. The incomplete response might be because 1) different dtype 2) not enough max_new_token.

[chuangchuang-wang]

thank you for your response. I have used torch.bfloat16 and max_new_tokens=1024, get ea=0.593, em=0.585, still bellow 0.75, my inference code as fllows:

from examples.db_connectors import SQLiteConnector
from examples.prompt_formatters import RajkumarFormatter
from examples.db_connectors import PostgresConnector
        
CUDA_DEV = torch.device("cuda:0")

def get_prompt(db_path, question):
    # Get the connector and formatter
    TABLES = []
    sqlite_connector = SQLiteConnector(
        database_path=db_path
    )
    sqlite_connector.connect()
    if len(TABLES) <= 0:
        TABLES.extend(sqlite_connector.get_tables())

    print(f"Loading tables: {TABLES}")

    db_schema = [sqlite_connector.get_schema(table) for table in TABLES]
    formatter = RajkumarFormatter(db_schema)
    prompt = formatter.format_prompt(question)
    return prompt

def get_tokenizer_model_nsql(model_name):
    tokenizer = AutoTokenizer.from_pretrained("NumbersStation/nsql-llama-2-7B", cache_dir='pretrain')
    model = AutoModelForCausalLM.from_pretrained("NumbersStation/nsql-llama-2-7B", torch_dtype=torch.bfloat16, cache_dir='pretrain').to(CUDA_DEV)
    return tokenizer, model

def run_inference_nsql(question, tokenizer, model, prompt_file="prompt_nsql.md", metadata_file="metadata_nsql.sql"):
    # prompt = generate_prompt(question, prompt_file, metadata_file)
    prompt = get_prompt(metadata_file, question)
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(CUDA_DEV)
    generated_ids = model.generate(input_ids, max_length=1024)
    generated_query = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return 'SELECT' + generated_query.split('SELECT')[-1]

if __name__ == "__main__":
    checkpoint = "NumbersStation/nsql-llama-2-7B"
    db_dir = "/home/notebook/data/personal/data/dataset/spider/database"
    devPath = "/home/notebook/data/personal/data/dataset/spider/dev.json"
    genSql = "./nsql_pred_a100_maxlen1024_bf16_modelinit.txt"
    genFile = open(genSql, "a")
    with open(devPath) as f:
        devs=json.load(f)
        tokenizer, model = get_tokenizer_model_nsql(checkpoint)
        for dev in tqdm(devs):
            question = dev["question"]
            dbname = dev["db_id"]
            # metadata_file = os.path.join(db_dir, dbname, "schema.sql")
            # if not os.path.exists(metadata_file):
            #     metadata_file = os.path.join(db_dir, dbname, dbname+".sql")
            metadata_file = os.path.join(db_dir, dbname, dbname+".sqlite")
            generateSql=run_inference_nsql(question, tokenizer, model,metadata_file=metadata_file).replace("\n", " ").replace("\t", " ")
            # try:
            #     generateSql=run_inference_nsql(question, tokenizer, model,metadata_file=metadata_file)
            # except Exception as e:
            #     generateSql= "Error:" + str(e).replace("\n", " ")
            genFile.write(generateSql + "\n")

    genFile.close()