Code Generation Model for Multilingual Program Synthesis.
| 25 | |
| 26 | |
| 27 | class CodeGeeXModel(MegatronModule): |
| 28 | """Code Generation Model for Multilingual Program Synthesis.""" |
| 29 | |
| 30 | def __init__(self, num_tokentypes=0, parallel_output=False): |
| 31 | super(CodeGeeXModel, self).__init__() |
| 32 | args = get_args() |
| 33 | |
| 34 | self.parallel_output = parallel_output |
| 35 | self.fp16_lm_cross_entropy = args.fp16_lm_cross_entropy |
| 36 | |
| 37 | self.language_model, self._language_model_key = get_language_model( |
| 38 | num_tokentypes=num_tokentypes, |
| 39 | add_pooler=False, |
| 40 | init_method=init_method_normal(args.init_method_std), |
| 41 | scaled_init_method=scaled_init_method_normal(args.init_method_std, |
| 42 | args.num_layers)) |
| 43 | |
| 44 | def set_input_tensor(self, input_tensor): |
| 45 | """See megatron.model.transformer.set_input_tensor()""" |
| 46 | self.language_model.set_input_tensor(input_tensor) |
| 47 | |
| 48 | def forward( |
| 49 | self, |
| 50 | input_ids, |
| 51 | position_ids, |
| 52 | attention_mask, |
| 53 | labels=None, |
| 54 | tokentype_ids=None, |
| 55 | layer_past=None, |
| 56 | get_key_value=False, |
| 57 | forward_method_parallel_output=None, |
| 58 | prompt_length=None, |
| 59 | context_length=None, |
| 60 | ): |
| 61 | |
| 62 | # Language model. |
| 63 | lm_output = self.language_model(input_ids, |
| 64 | position_ids, |
| 65 | attention_mask, |
| 66 | tokentype_ids=tokentype_ids, |
| 67 | layer_past=layer_past, |
| 68 | get_key_value=get_key_value, |
| 69 | prompt_length=prompt_length, |
| 70 | context_length=context_length) |
| 71 | |
| 72 | if get_key_value: |
| 73 | lm_output, presents = lm_output |
| 74 | |
| 75 | lm_output = torch.add(lm_output, 0) |
| 76 | # Output. |
| 77 | parallel_output = self.parallel_output |
| 78 | if forward_method_parallel_output is not None: |
| 79 | parallel_output = forward_method_parallel_output |
| 80 | output = parallel_lm_logits( |
| 81 | lm_output, |
| 82 | self.language_model.embedding.word_embeddings.weight, |
| 83 | parallel_output) |
| 84 |
no outgoing calls
no test coverage detected