| 128 | return pe.cuda().contiguous() |
| 129 | |
| 130 | class CustomDecoding(nn.Module): |
| 131 | def __init__(self, head_num, head_size, |
| 132 | inter_size, mem_hidden_dim, layer_num, vocab_size, start_id, end_id, |
| 133 | beam_search_diversity_rate, top_k, top_p, temperature, |
| 134 | len_penalty, repetition_penalty, weights, args=None): |
| 135 | super().__init__() |
| 136 | self.end_id = end_id |
| 137 | self.args = args |
| 138 | torch.classes.load_library(os.path.abspath(args.decoding_ths_path)) |
| 139 | try: |
| 140 | self.decoding = torch.classes.FasterTransformer.Decoding(head_num, head_size, |
| 141 | inter_size, mem_hidden_dim, layer_num, vocab_size, start_id, end_id, |
| 142 | beam_search_diversity_rate, top_k, top_p, temperature, |
| 143 | len_penalty, repetition_penalty, *weights.w) |
| 144 | except: |
| 145 | # legacy ths for 20.03 image |
| 146 | self.decoding = torch.classes.FasterTransformerDecoding(head_num, head_size, |
| 147 | inter_size, mem_hidden_dim, layer_num, vocab_size, start_id, end_id, |
| 148 | beam_search_diversity_rate, top_k, top_p, temperature, |
| 149 | len_penalty, repetition_penalty, *weights.w) |
| 150 | self.is_clean_cache = False |
| 151 | def forward(self, batch_size, beam_size, seq_len, memory, memory_seq_lens): |
| 152 | if self.is_clean_cache == False: |
| 153 | torch.cuda.empty_cache() |
| 154 | self.is_clean_cache = True |
| 155 | |
| 156 | extended_memory = tile(memory, beam_size) |
| 157 | extended_memory_seq_lens = tile(memory_seq_lens, beam_size) |
| 158 | output_ids, parent_ids, out_seq_lens = self.decoding.forward(beam_size, seq_len, extended_memory, extended_memory_seq_lens) |
| 159 | output_ids = output_ids.reshape([seq_len, memory.size(0), beam_size]) |
| 160 | output_ids = output_ids.permute(1, 2, 0) |
| 161 | return output_ids, out_seq_lens |
| 162 | |
| 163 | class ArgHelper(object): |
| 164 | def __init__(self, model_type=None, data_type=None, ths_path=None): |
no outgoing calls
no test coverage detected