| 24 | from onmt.utils.misc import tile |
| 25 | |
| 26 | class FtDecodingWeights(object): |
| 27 | def __init__(self, layer_num, hidden_dim, onmtcheckpoint, max_step_for_pe=2048): |
| 28 | self.max_step_for_pe = max_step_for_pe |
| 29 | self.hidden_dim = hidden_dim |
| 30 | self.w = [] |
| 31 | prefix = 'decoder.transformer_layers.' |
| 32 | self.w.append(torch.stack( |
| 33 | [onmtcheckpoint['model'][prefix + str(i) + '.layer_norm_1.weight'] for i in range(layer_num)], |
| 34 | 0).contiguous()) |
| 35 | self.w.append(torch.stack( |
| 36 | [onmtcheckpoint['model'][prefix + str(i) + '.layer_norm_1.bias'] for i in range(layer_num)], |
| 37 | 0).contiguous()) |
| 38 | self.w.append(torch.stack( |
| 39 | [torch.stack([onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_query.weight'].transpose(-1, -2), |
| 40 | onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_keys.weight'].transpose(-1, -2), |
| 41 | onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_values.weight'].transpose(-1, -2)], -2) |
| 42 | for i in range(layer_num)], 0).contiguous()) |
| 43 | self.w.append(torch.stack( |
| 44 | [torch.stack([onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_query.bias'], |
| 45 | onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_keys.bias'], |
| 46 | onmtcheckpoint['model'][prefix + str(i) + '.self_attn.linear_values.bias']], -2) for i in range(layer_num)], |
| 47 | 0).contiguous()) |
| 48 | self.w.append(torch.stack( |
| 49 | [onmtcheckpoint['model'][prefix + str(i) + '.self_attn.final_linear.weight'].transpose(-1, -2) for i in range(layer_num)], |
| 50 | 0).contiguous()) |
| 51 | self.w.append(torch.stack( |
| 52 | [onmtcheckpoint['model'][prefix + str(i) + '.self_attn.final_linear.bias'] for i in range(layer_num)], |
| 53 | 0).contiguous()) |
| 54 | self.w.append(torch.stack( |
| 55 | [onmtcheckpoint['model'][prefix + str(i) + '.layer_norm_2.weight'] for i in range(layer_num)], |
| 56 | 0).contiguous()) |
| 57 | self.w.append(torch.stack( |
| 58 | [onmtcheckpoint['model'][prefix + str(i) + '.layer_norm_2.bias'] for i in range(layer_num)], |
| 59 | 0).contiguous()) |
| 60 | self.w.append(torch.stack( |
| 61 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_query.weight'].transpose(-1, -2) for i in range(layer_num)], |
| 62 | 0).contiguous()) |
| 63 | self.w.append(torch.stack( |
| 64 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_keys.weight'].transpose(-1, -2) for i in range(layer_num)], |
| 65 | 0).contiguous()) |
| 66 | self.w.append(torch.stack( |
| 67 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_values.weight'].transpose(-1, -2) for i in range(layer_num)], |
| 68 | 0).contiguous()) |
| 69 | self.w.append(torch.stack( |
| 70 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_query.bias'] for i in range(layer_num)], |
| 71 | 0).contiguous()) |
| 72 | self.w.append(torch.stack( |
| 73 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_keys.bias'] for i in range(layer_num)], |
| 74 | 0).contiguous()) |
| 75 | self.w.append(torch.stack( |
| 76 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.linear_values.bias'] for i in range(layer_num)], |
| 77 | 0).contiguous()) |
| 78 | self.w.append(torch.stack( |
| 79 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.final_linear.weight'].transpose(-1, -2) for i in range(layer_num)], |
| 80 | 0).contiguous()) |
| 81 | self.w.append(torch.stack( |
| 82 | [onmtcheckpoint['model'][prefix + str(i) + '.context_attn.final_linear.bias'] for i in range(layer_num)], |
| 83 | 0).contiguous()) |
no outgoing calls
no test coverage detected