| 113 | } |
| 114 | |
| 115 | def predict_dataset(model, tokenizer, dataloader, device, output_file=None): |
| 116 | model.eval() |
| 117 | results = [] |
| 118 | |
| 119 | for batch in tqdm(dataloader, desc="Predicting"): |
| 120 | input_ids = batch["input_ids"].to(device) |
| 121 | attention_mask = batch["attention_mask"].to(device) |
| 122 | offset_mapping = batch["offset_mapping"] # [B, L, 2] |
| 123 | texts = batch["texts"] |
| 124 | |
| 125 | with torch.no_grad(): |
| 126 | logits = model(input_ids=input_ids, attention_mask=attention_mask).logits # [B, L, num_labels] |
| 127 | |
| 128 | # Get the predicted class (0 or 1) for each token |
| 129 | pred_labels = torch.argmax(logits, dim=-1) # [B, L] |
| 130 | |
| 131 | for i in range(input_ids.size(0)): |
| 132 | text = texts[i] |
| 133 | # Get offsets for the current sample. Note: offsets include special tokens |
| 134 | sample_offsets = offset_mapping[i][:len(input_ids[i])].cpu().tolist() |
| 135 | # Get the predicted label sequence for the current sample |
| 136 | # predicted_label_sequence[j] predicts the class of input_ids[j+1] |
| 137 | sample_pred_labels = pred_labels[i][:len(input_ids[i])].cpu().tolist() |
| 138 | |
| 139 | reconstructed_pred_results = [] |
| 140 | current_span_start_char = -1 # Record the character start position of the current entity span |
| 141 | last_token_end_char = -1 # Record the character end position of the last token in the current entity span |
| 142 | |
| 143 | # Iterate through input_ids, starting from the second token (the first actual token after [CLS]) |
| 144 | # Because predicted_label_sequence[j] corresponds to the label of input_ids[j+1] |
| 145 | # So we start from j=0 (predicting the label of input_ids[1]) to the second-to-last token (predicting the label of input_ids[L-1]) |
| 146 | # This way j+1 will not exceed the valid range of input_ids |
| 147 | |
| 148 | # This loop iterates through the label sequence, pred_labels[j] is the prediction for input_ids[j+1] |
| 149 | for j in range(len(sample_pred_labels) - 1): # The predicted label sequence is one shorter than input_ids (the last token has no label) |
| 150 | # Get the information of the token corresponding to the current prediction (input_ids[j+1]) |
| 151 | token_idx_in_input_ids = j + 1 |
| 152 | |
| 153 | # Skip special tokens (e.g., [SEP] token, whose offsets are usually (0,0)) |
| 154 | # And ensure token_idx_in_input_ids is within the valid range of sample_offsets |
| 155 | if token_idx_in_input_ids >= len(sample_offsets): |
| 156 | break # Should not happen, but just in case |
| 157 | |
| 158 | token_char_start, token_char_end = sample_offsets[token_idx_in_input_ids] |
| 159 | |
| 160 | # Ignore special tokens, their offsets are often (0,0) and are not in the actual text range |
| 161 | # If both token_char_start and token_char_end are 0, it usually indicates a special token |
| 162 | if token_char_start == 0 and token_char_end == 0 and token_idx_in_input_ids != 0: # Exclude the case where the first token is CLS |
| 163 | # If there is an ongoing span and we encounter a special token, we should end the span |
| 164 | if current_span_start_char != -1: |
| 165 | if last_token_end_char != -1 and current_span_start_char < last_token_end_char: |
| 166 | reconstructed_pred_results.append({ |
| 167 | "text": text[current_span_start_char:last_token_end_char], |
| 168 | "start": current_span_start_char, |
| 169 | "end": last_token_end_char |
| 170 | }) |
| 171 | current_span_start_char = -1 |
| 172 | last_token_end_char = -1 |