aggregator_predict.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. import json
  2. import torch
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5. from torch.autograd import Variable
  6. import numpy as np
  7. from sqlnet.model.modules.net_utils import run_lstm, col_name_encode
  8. class AggPredictor(nn.Module):
  9. def __init__(self, N_word, N_h, N_depth, use_ca):
  10. super(AggPredictor, self).__init__()
  11. self.use_ca = use_ca
  12. self.agg_lstm = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
  13. if use_ca:
  14. print ("Using column attention on aggregator predicting")
  15. self.agg_col_name_enc = nn.LSTM(input_size=N_word, hidden_size=int(N_h/2), num_layers=N_depth, batch_first=True, dropout=0.3, bidirectional=True)
  16. self.agg_att = nn.Linear(N_h, N_h)
  17. else:
  18. print ("Not using column attention on aggregator predicting")
  19. self.agg_att = nn.Linear(N_h, 1)
  20. self.agg_out = nn.Sequential(nn.Linear(N_h, N_h), nn.Tanh(), nn.Linear(N_h, 6))
  21. self.softmax = nn.Softmax(dim=-1)
  22. self.agg_out_K = nn.Linear(N_h, N_h)
  23. self.col_out_col = nn.Linear(N_h, N_h)
  24. def forward(self, x_emb_var, x_len, col_inp_var=None, col_name_len=None,
  25. col_len=None, col_num=None, gt_sel=None, gt_sel_num=None):
  26. B = len(x_emb_var)
  27. max_x_len = max(x_len)
  28. e_col, _ = col_name_encode(col_inp_var, col_name_len, col_len, self.agg_col_name_enc)
  29. h_enc, _ = run_lstm(self.agg_lstm, x_emb_var, x_len)
  30. col_emb = []
  31. for b in range(B):
  32. cur_col_emb = torch.stack([e_col[b,x] for x in gt_sel[b]] + [e_col[b,0]] * (4-len(gt_sel[b])))
  33. col_emb.append(cur_col_emb)
  34. col_emb = torch.stack(col_emb)
  35. att_val = torch.matmul(self.agg_att(h_enc).unsqueeze(1), col_emb.unsqueeze(3)).squeeze() # .transpose(1,2))
  36. for idx, num in enumerate(x_len):
  37. if num < max_x_len:
  38. att_val[idx, num:] = -100
  39. att = self.softmax(att_val.view(B*4, -1)).view(B, 4, -1)
  40. K_agg = (h_enc.unsqueeze(1) * att.unsqueeze(3)).sum(2)
  41. agg_score = self.agg_out(self.agg_out_K(K_agg) + self.col_out_col(col_emb)).squeeze()
  42. return agg_score