File size: 2,832 Bytes
bc56514 5071ec6 bc56514 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import numpy as np
from sklearn.metrics import precision_recall_fscore_support
import torch
from torch.utils.data import Dataset
# Custom Dataset class
class AbuseDataset(Dataset):
def __init__(self, texts, labels, tokenizer):
self.encodings = tokenizer(texts, truncation=True, padding=True, max_length=512)
self.labels = labels
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
item["labels"] = torch.tensor(self.labels[idx], dtype=torch.float)
return item
def __getitem__(self, idx):
item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
item["labels"] = torch.tensor(self.labels[idx], dtype=torch.float)
return item
# Label map used across modules
label_map = {
0.0: "no",
0.5: "plausibly",
1.0: "yes"
}
# Function to map probabilities to 3 classes
# (0.0, 0.5, 1.0) based on thresholds
def map_to_3_classes(prob_array, low, high):
"""Map probabilities to 0.0, 0.5, 1.0 using thresholds."""
mapped = np.zeros_like(prob_array)
mapped[(prob_array > low) & (prob_array <= high)] = 0.5
mapped[prob_array > high] = 1.0
return mapped
def convert_to_label_strings(array):
"""Convert float label array to list of strings."""
return [label_map[val] for val in array.flatten()]
def tune_thresholds(probs, true_labels, verbose=True):
"""Search for best (low, high) thresholds by macro F1 score."""
best_macro_f1 = 0.0
best_low, best_high = 0.0, 0.0
for low in np.arange(0.2, 0.5, 0.05):
for high in np.arange(0.55, 0.8, 0.05):
if high <= low:
continue
pred_soft = map_to_3_classes(probs, low, high)
pred_str = convert_to_label_strings(pred_soft)
true_str = convert_to_label_strings(true_labels)
_, _, f1, _ = precision_recall_fscore_support(
true_str, pred_str,
labels=["no", "plausibly", "yes"],
average="macro",
zero_division=0
)
if verbose:
print(f"low={low:.2f}, high={high:.2f} -> macro F1={f1:.3f}")
if f1 > best_macro_f1:
best_macro_f1 = f1
best_low, best_high = low, high
return best_low, best_high, best_macro_f1
# Convert label values to soft scores: "yes" = 1.0, "plausibly" = 0.5, others = 0.0
def label_row_soft(row, label_columns):
labels = []
for col in label_columns:
val = str(row[col]).strip().lower()
if val == "yes":
labels.append(1.0)
elif val == "plausibly":
labels.append(0.5)
else:
labels.append(0.0)
return labels
|