| 136 | ) |
| 137 | |
| 138 | class InvalidScoreLogitsProcessor(LogitsProcessor): |
| 139 | def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: |
| 140 | if torch.isnan(scores).any() or torch.isinf(scores).any(): |
| 141 | scores.zero_() |
| 142 | scores[..., 20005] = 1e4 |
| 143 | return scores |
| 144 | |
| 145 | |
| 146 | def load_tf_weights_in_chatglm_6b(model, config, tf_checkpoint_path): |