aiops/RCRank/pretrain/alignment_new/pretrain.py

import json
import torch.optim as optim
import sys
sys.path.append("../..")
from model.modules.LogModel.log_model import LogModel
import torch.nn as nn
import torch
from model.modules.QueryFormer.utils import *
from model.modules.QueryFormer.QueryFormer import QueryFormer
from transformers import BertTokenizer, BertModel
import pickle

class Predict(nn.Module):
    def __init__(self, input_dim, model_dim, num_heads, ff_dim, dropout=0.1):
        super(Predict, self).__init__()
        encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=num_heads, dim_feedforward=ff_dim, dropout=dropout,batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
        self.input_fc = nn.Linear(input_dim,1100,bias=False)  
        self.input_fc1 = nn.Linear(1100,model_dim,bias=False)  
    def forward(self, x):
        x = self.input_fc(x)
        x = self.input_fc1(x)
        x = self.transformer_encoder(x)
        return x
    
class Alignment(nn.Module):
    def __init__(self,device):
        super(Alignment, self).__init__()
        self.flatten = nn.Flatten()
        self.plan_model = QueryFormer(pred_hid=32)
        self.sql_model = BertModel.from_pretrained("./bert-base-uncased")
        encoder_layer = nn.TransformerEncoderLayer(d_model=768, nhead=8, dim_feedforward=2048, dropout=0.1, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=2)
        self.tokenizer = BertTokenizer.from_pretrained("./bert-base-uncased")
        self.log_model = LogModel(input_dim=13, hidden_dim = 64, output_dim = 32)

        self.concat_dim_mask_plan = 17600
        self.predict_mask_plan = Predict(input_dim=self.concat_dim_mask_plan, model_dim=1024, num_heads=8, ff_dim=2048)
        self.Linear_mask_plan = nn.Linear(1024, 1067)

        self.concat_dim_mask_sql = 17600
        self.predict_mask_sql= Predict(input_dim=self.concat_dim_mask_sql, model_dim=768, num_heads=8, ff_dim=2048)
        self.Linear_mask_sql= nn.Linear(768, 768)

        self.device = device
    def forward(self, plan,sql, log,dic,mod):
        plan = self.plan_model(plan)
        sql = self.tokenizer(sql, return_tensors="pt", padding=True, truncation=True, max_length=512).to(self.device)
        sql = self.sql_model(**sql).last_hidden_state
        sql = self.transformer_encoder(sql).mean(dim=1)
        log = self.log_model(log)
        plan = self.flatten(plan)
        concatenated_vector = torch.cat((plan, sql, log, dic), dim=1).unsqueeze(1)
        if mod == 'mask_plan':
          transformer_output = self.predict_mask_plan(concatenated_vector)  
          transformer_output = transformer_output[:, 0, :]
          predicted_vector = self.Linear_mask_plan(transformer_output)
        elif mod == 'mask_sql':
          transformer_output = self.predict_mask_sql(concatenated_vector)
          transformer_output = transformer_output[:, 0, :]
          predicted_vector = self.Linear_mask_sql(transformer_output)
        pass
        return predicted_vector
      
  
if __name__ == "__main__":
    
    f1=open('mask_table_plan/encoding_1w.pickle','rb')
    encoding_mask_plan=json.dumps(pickle.load(f1).idx2table)
    f2=open('mask_table_sql/encoding1w.pickle','rb')
    encoding_mask_sql=json.dumps(pickle.load(f2).idx2table)
   
    device = 'cuda:2'
    tokenizer = BertTokenizer.from_pretrained("./bert-base-uncased")
    bert = BertModel.from_pretrained("./bert-base-uncased").to(device)    
    encoding_mask_plan = tokenizer(encoding_mask_plan, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
    encoding_mask_sql = tokenizer(encoding_mask_sql, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)  
    encoding={'mask_plan':encoding_mask_plan,'mask_sql':encoding_mask_sql}
    model = Alignment(device).to(device)
    model.load_state_dict(torch.load('./model6.pth'))
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    epochs=500
    saved = 0
    best=0.0
    best_val_loss = float('inf') 
    batch_size=18
    batch_loss={'mask_sql':0.0,'mask_plan':0.0}
    for epoch in range(epochs):
       
        if epoch <= 6:
          print(epoch)
        else:
          i=0
          size=0
          model.train()  
          epoch_loss = 0.0
          file = {}
          file['mask_plan'] = open('mask_table_plan/plan_mask_1w.txt','r')
          file['mask_sql'] = open('mask_table_sql/sql_mask1w.txt','r')
          batch_list={}
          while i <= 10110:
            for mod in ['mask_sql','mask_plan']:
              for _ in range(batch_size):
                  line = file[mod].readline()
                  line=line.strip()
                  if not line:  
                      break
                  (query,x1,attn_bias,rel_pos,height,log_all,predict) = json.loads(line)
                  size=size+1
                  if len(batch_list.keys()) == 0:
                      batch_list['query']=[query]
                      plan={}
                      plan['x']= torch.tensor(x1).to(torch.float32).to(device) 
                      plan["attn_bias"] = torch.tensor(attn_bias).to(device) 
                      plan["rel_pos"] = torch.tensor(rel_pos).to(device)
                      plan["heights"] = torch.tensor(height).to(device) 
                      batch_list['plan']=plan
                      batch_list['log'] = torch.tensor(log_all).unsqueeze(0).to(device) 
                      batch_list['predict'] =  torch.tensor(predict).to(device).unsqueeze(0)
                  else:
                      batch_list['query'].append(query)
                      batch_list['plan']['x'] = torch.cat((batch_list['plan']['x'],torch.tensor(x1).to(torch.float32).to(device)),dim=0)
                      batch_list['plan']["attn_bias"]= torch.cat([batch_list['plan']['attn_bias'],torch.tensor(attn_bias).to(device)],dim=0) 
                      batch_list['plan']["rel_pos"]= torch.cat([batch_list['plan']['rel_pos'],torch.tensor(rel_pos).to(device)],dim=0)
                      batch_list['plan']["heights"]= torch.cat([batch_list['plan']['heights'],torch.tensor(height).to(device)],dim=0 )
                      batch_list['log'] =torch.cat([batch_list['log'],torch.tensor(log_all).unsqueeze(0).to(device)],dim=0 )
                      batch_list['predict'] =torch.cat([batch_list['predict'], torch.tensor(predict).unsqueeze(0).to(device)],dim=0)
              
              optimizer.zero_grad()
              sql, plan, log,predict = batch_list["query"], batch_list["plan"],batch_list["log"],batch_list['predict']
              dic = bert(**encoding[mod]).pooler_output
              dic = dic.repeat(len(sql), 1)
              output = model(plan,sql,log,dic,mod)
              loss = criterion(output, predict)
              loss.backward()
              optimizer.step()
              batch_loss[mod] = loss.item()
              batch_list={}
              if mod == 'mask_sql':
                i=i+batch_size
          


          if  (epoch+1) %3 == 0:
              torch.save(model.state_dict(), f'model{(epoch+1)}.pth')