make_table_cache.py

import glob
import re
import queue
#from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
from models.qwen2 import Qwen2ModifiedForCausalLM
from utils.cache_manager import FIFO_cache
import torch
fa = []
block_size = []
def GetFa(x):
    global fa
    if x == fa[x]:
        return x
    fa[x] = GetFa(fa[x])
    return fa[x]

def stack_past_key_values(past_key_values_list):
    num_layers = len(past_key_values_list[0])
    batch_past_key_values = []
    for layer in range(num_layers):
        #print([past_key_values[layer][0] for past_key_values in past_key_values_list])
        keys = torch.cat([past_key_values[layer][0] for past_key_values in past_key_values_list], dim=2)
        values = torch.cat([past_key_values[layer][1] for past_key_values in past_key_values_list], dim=2)
        batch_past_key_values.append((keys, values))
    return tuple(batch_past_key_values)

output_path = ""
cache_manager = FIFO_cache(1024,output_path)

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

def calc_and_save_kv(need_set,y,_id):
    concated_table = ""
    past_key_and_value = []
    for x in need_set:
        #print(x)
        #if not (x in calced):
        chunk_id = str(x[0])
        table = x[1]
        concated_table += table
        cache_manager.insert(chunk_id)
        kv_value = cache_manager.query(chunk_id)
        #print("kv_value:",kv_value)
        #exit()
        past_key_and_value.append(kv_value)
    if len(past_key_and_value) == 0:
        stacked_past_kv = None
    else:
        stacked_past_kv = stack_past_key_values(past_key_and_value)
    #print(y)
    #print("len_of_past_kv:",len(past_key_and_value))
    #print(stacked_past_kv)
    concated_table = concated_table
    concated_table = concated_table + y
    inputs = tokenizer(concated_table, return_tensors="pt").to(model.device)
    inputs_1 = tokenizer.encode(y,return_tensors = "pt")
    input_ids = tokenizer.encode(concated_table, return_tensors="pt")
    full_inputs = tokenizer(concated_table, return_tensors="pt").to(model.device)
    y_inputs = tokenizer(y, return_tensors="pt").to(model.device)
    
    full_tokens = full_inputs.input_ids[0]
    y_tokens = y_inputs.input_ids[0]
    y_start_idx = None
    for i in range(len(full_tokens) - len(y_tokens) + 1):
        #if "bond" in y:
        #    print(full_tokens[i:i+len(y_tokens)],y_tokens)
        if torch.equal(full_tokens[i:i+len(y_tokens)], y_tokens):
            y_start_idx = i
            break
    if y_start_idx is None:
        print(concated_table)
        print(y)
        raise ValueError("Could not find y tokens in the full sequence")
    with torch.no_grad():
        outputs = model(**inputs,use_cache=True)
    last_past_kv_cache = []
    past_key_values = outputs.past_key_values
    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:, :]
            layer_kv_cache.append(last_tokens_kv_cache)
        last_past_kv_cache.append(tuple(layer_kv_cache))
    
    #past_key_values = outputs.past_key_values[ : , -len(inputs_1):]
    last_past_kv_cache = tuple(last_past_kv_cache)
    kvcache_file_path = f'{output_path}/kvcache_chunk_{_id}.pt'
    torch.save(last_past_kv_cache, kvcache_file_path)

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

table_ids = {}
table_names = {}
cnt = 0
file_names = glob.glob("./inst_datas_bird/*.txt")
for file in file_names:
    with open(file,"r",encoding = "utf-8") as f:
        content = f.read()
    table_name = content.split("(")[0].strip()
    try:
        table_name = table_name.split("TABLE")[1].strip()
    except:
        print(table_name)
        continue
    if not (table_name in table_ids):
        origin_path = "./inst_datas_bird_clean/"
        txt_name = str(cnt) + ".txt"
        content = content
        with open(origin_path + txt_name,"w",encoding = "utf-8") as fw:
            fw.write(content)
        table_ids[table_name] = (cnt,content)
        table_names[cnt] = content
        cnt += 1

for i in range(cnt):
    fa.append(i)
    block_size.append(1)
origin_blocks = cnt
mx_block_size = 0
vec = {}
in_d = [0]*100
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
        #print(_id,"\t",_id_to)
        fx , fy = GetFa(_id) , GetFa(_id_to)
        if not (fx == fy):
            fa[fx] = fy
            block_size[fy] += block_size[fx]
            block_size[fx] = 0
            origin_blocks -= 1
        mx_block_size = max(block_size[fy],mx_block_size)

print("max_size:",mx_block_size)
print("blocks:",origin_blocks)
print("aver_size",cnt / origin_blocks)

rep_data = []
for i in range(cnt):
    if fa[i] == i:
        rep_data.append(i)

need_set = {}



for x in rep_data:
    print(x)
    point_set = []
    topo = queue.Queue()
    print("graph:")
    topo_sort_list = []
    for i in range(cnt):
        if GetFa(i) == x:
            point_set.append(i)
            if in_d[i] == 0:
                topo.put(i)
                topo_sort_list.append((i,table_names[i]))
                calc_and_save_kv([], table_names[i] ,str(i))
                #print("start_point",i)
            if not (i in vec):
                continue
            for y in vec[i]:
                print(i,y)
    #print("size:",len(point_set))
    while (not topo.empty()):
        Now = topo.get()
        #print("now point:",Now)
        if not Now in vec:
            continue
        for y in vec[Now]:
            in_d[y] -= 1
            if not y in need_set:
                need_set[y] = []
            need_set[y].append((Now,table_names[Now]))
            if in_d[y] == 0:
                print(y)
                topo.put(y)
                calc_and_save_kv(topo_sort_list,table_names[y],str(y))
                topo_sort_list.append((y,table_names[y]))