make_spider_cache.py

import json
import re
import glob
import re
import queue
#from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer , AutoModelForCausalLM
from models.qwen2 import Qwen2ModifiedForCausalLM
from utils.cache_manager import FIFO_cache
import torch

output_path = ""

def extract_foreign_key_table_names(string):
    pattern = r"REFERENCES (\w*)"
    matches = re.findall(pattern, string)
    return matches

device = "cuda"
model = Qwen2ModifiedForCausalLM.from_pretrained("",torch_dtype=torch.bfloat16).to(device)
tokenizer = AutoTokenizer.from_pretrained("")

def extract_table_schemas(input_text):
    table_pattern = r'CREATE TABLE \w+ \([\s\S]*?\);'
    tables = re.findall(table_pattern, input_text)
    
    return tables

def remove_create_tables(sql: str) -> str:
    pattern = re.compile(
        r"CREATE\s+(?:TEMPORARY\s+|TEMP\s+)?TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?"
        r"(?:[a-zA-Z_][a-zA-Z0-9_]*\.)?"
        r"(?:`[^`]+`|\"[^\"]+\"|\[[^\]]+\]|[a-zA-Z_][a-zA-Z0-9_]*)"
        r"\s*\(.*?\);",
        re.IGNORECASE | re.DOTALL
    )
    cleaned = pattern.sub("", sql)
    cleaned = re.sub(r"\n\s*\n+", "\n", cleaned).strip()
    return cleaned

with open("golden.txt","r",encoding = "utf-8") as f:
    lines = f.readlines()
cnt = 0
db_name = {}

for line in lines:
    db_name[cnt] = line.split("\t")[1]
    cnt += 1

with open('./eval/dev_spider.json',"r",encoding = "utf-8") as f:
    lines = json.load(f)

cnt_tables = 0
ori_pos = 0

for i in range(len(lines)):
    x = lines[i]
    if i == len(lines) - 1 or not(db_name[i] == db_name[i+1]) :
        #print(i)
        st_pos = i
        table_list = extract_table_schemas(x['input_seq'])
        cnt = 0
        table_names = {}
        table_ids = {}
        vec = {}
        in_d = [0]*100
        for x in table_list:
            content = x
            table_name = content.split("(")[0].strip()
            table_name = table_name.split("TABLE")[1].strip()
            if not (table_name in table_ids):
                table_ids[table_name] = (cnt,content)
                table_names[cnt] = content
                cnt += 1
        for x , y in table_ids.items():
            _id , table = y
            refs = extract_foreign_key_table_names(table)
            refs = list(set(refs))
            for to in refs:
                _id_to , table_to = table_ids[to]
                if not (_id_to in vec):
                    vec[_id_to] = []
                vec[_id_to].append(_id)
                in_d[_id] += 1
        topo = queue.Queue()
        topo_list = []
        vis = [0] * cnt
        topo_sort_list = []
        for i in range(cnt):
            if in_d[i] == 0:
                topo.put(i)
                topo_sort_list.append(table_names[i])
                vis[i] = 1
        while (not topo.empty()):
            Now = topo.get()
            if not Now in vec:
                continue
            for y in vec[Now]:
                in_d[y] -= 1
            if in_d[y] == 0:
                topo_sort_list.append(table_names[y])
                vis[y] = 1
                topo.put(y)
        
        concated_table = ""
        
        for x in topo_sort_list:
            concated_table += x

        for data in lines[ori_pos + 1 : st_pos + 1]:
            now_data = data
            prompt = now_data['input_seq']
            gt = now_data['output_seq']
            other_prompt = remove_create_tables(prompt)
            new_inputs = concated_table + "\n" + other_prompt
            new_data = {"input_seq":new_inputs,"output_seq":gt}
        ori_pos = st_pos
        

        full_inputs = tokenizer(concated_table, return_tensors="pt").to(model.device)
        full_tokens = full_inputs.input_ids[0]
        with torch.no_grad():
            outputs = model(**full_inputs,use_cache=True)
        past_key_values = outputs.past_key_values
        for x in topo_sort_list:
            last_past_kv_cache = []
            y = x
            y_inputs = tokenizer(y, return_tensors="pt").to(model.device)
            y_tokens = y_inputs.input_ids[0]
            y_start_idx = None
            for i in range(len(full_tokens) - len(y_tokens) + 1):
                if torch.equal(full_tokens[i:i+len(y_tokens)], y_tokens):
                    y_start_idx = i
                    break
            if y_start_idx is None:
                raise ValueError("Could not find y tokens in the full sequence")
            for layer_past_key_values in past_key_values:
                layer_kv_cache = []
                for head_kv in layer_past_key_values:
                    last_tokens_kv_cache = head_kv[: , :, y_start_idx:y_start_idx + len(y_tokens), :]
                    layer_kv_cache.append(last_tokens_kv_cache)
                last_past_kv_cache.append(tuple(layer_kv_cache))
            last_past_kv_cache = tuple(last_past_kv_cache)
            _id = str(cnt_tables)
            kvcache_file_path = f'{output_path}/kvcache_chunk_{_id}.pt'
            torch.save(last_past_kv_cache, kvcache_file_path)
            cnt_tables += 1