import csv import os import numpy as np import pandas as pd import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader from transformers import AutoTokenizer, AutoModel from sklearn.metrics import f1_score, roc_auc_score, accuracy_score, precision_recall_fscore_support import itertools import warnings import random def set_seed(seed=13): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" set_seed(13) warnings.filterwarnings("ignore", category=FutureWarning) # --- CONFIG --- param_grid = { "learning_rate": [1e-5, 2e-5, 3e-5, 4e-5, 5e-5], "batch_size": [16, 32, 64], "optimizer": ["Adam"], "lambda": [0.2, 0.3, 0.4, 0.5, 0.6, 0.7] } num_epochs = 7 max_length = 128 model_name = "bert-base-multilingual-cased" num_labels = 3 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # --- LOAD DATA --- train_df = pd.read_csv("train.csv") val_df = pd.read_csv("val.csv") valid_labels = {"Negative": 0, "Neutral": 1, "Positive": 2} train_df = train_df[train_df["final_label"].isin(valid_labels.keys())] val_df = val_df[val_df["final_label"].isin(valid_labels.keys())] if train_df.empty: raise ValueError("Train dataset empty after filtering.") if val_df.empty: raise ValueError("Validation dataset empty after filtering.") # --- INITIALIZE TOKENIZER & ADD EMOJIS --- tokenizer = AutoTokenizer.from_pretrained(model_name) emoji_path = "emoji.csv" # adjust path if needed if os.path.exists(emoji_path): emoji_df = pd.read_csv(emoji_path) emoji_list = emoji_df["emoji"].dropna().astype(str).str.strip().tolist() existing_vocab = set(tokenizer.get_vocab().keys()) emoji_set = set(emoji_list) - existing_vocab if emoji_set: tokenizer.add_tokens(list(emoji_set)) print(f"Added {len(emoji_set)} new emoji tokens to the tokenizer.") else: print("No new emojis to add.") else: print(f"Emoji file not found at: {emoji_path}") # --- FUNCTIONS --- def generate_attention_vectors_from_rationales(df, tokenizer, epsilon=1e-8): attention_vectors = [] for _, row in df.iterrows(): text = str(row["Content"]) final_label = str(row["final_label"]).strip() encoding = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) offsets = encoding["offset_mapping"] num_tokens = len(offsets) avg_vector = np.zeros(num_tokens, dtype=np.float32) annotations = str(row.get("Annotations", "")).split("|") rationales = str(row.get("Rationale", "")).split("|") annot_vectors = [] for annot_label, annot_rationale in zip(annotations, rationales): if not annot_label: continue if annot_label.split("-")[0].strip() != final_label: continue spans = [s.strip() for s in annot_rationale.split(",") if s.strip()] if not spans: continue vec = np.zeros(num_tokens, dtype=np.float32) for span_text in spans: start = 0 while True: idx = text.find(span_text, start) if idx < 0: break span_start, span_end = idx, idx + len(span_text) for i, (tok_start, tok_end) in enumerate(offsets): if tok_end > span_start and tok_start < span_end: vec[i] = 1.0 start = idx + 1 if vec.sum() > 0: annot_vectors.append(vec) if annot_vectors: avg_vector = np.mean(annot_vectors, axis=0) avg_vector = np.where(avg_vector == 0, epsilon, avg_vector) attn_str = " ".join(f"{v:.8f}" for v in avg_vector) attention_vectors.append(attn_str) df["embert_attention"] = attention_vectors return df class RationaleDataset(Dataset): def __init__(self, df, tokenizer, max_length=128, label_mapping=None): self.df = df self.tokenizer = tokenizer self.max_length = max_length self.label_mapping = label_mapping def __len__(self): return len(self.df) def __getitem__(self, idx): row = self.df.iloc[idx] text = row["Content"] label = self.label_mapping[row["final_label"]] encoding = self.tokenizer( text, padding="max_length", truncation=True, max_length=self.max_length, return_tensors="pt" ) rationale_raw = [float(x) for x in row["embert_attention"].split()] \ if pd.notna(row["embert_attention"]) and row["embert_attention"].strip() else [] rationale_vector = np.concatenate([ np.array([0.0], dtype=np.float32), np.array(rationale_raw, dtype=np.float32), np.array([0.0], dtype=np.float32) ]) rationale_vector = rationale_vector[:self.max_length] if len(rationale_vector) < self.max_length: rationale_vector = np.pad(rationale_vector, (0, self.max_length - len(rationale_vector)), constant_values=0.0) rationale_tensor = torch.tensor(rationale_vector, dtype=torch.float32) if torch.sum(rationale_tensor) == 0.0: has_rationale = False rationale_probs = torch.ones(self.max_length, dtype=torch.float32) / self.max_length else: has_rationale = True rationale_probs = torch.softmax(rationale_tensor, dim=0) return ( encoding["input_ids"].squeeze(0), encoding["attention_mask"].squeeze(0), rationale_probs, torch.tensor(label, dtype=torch.long), torch.tensor(has_rationale, dtype=torch.bool) ) class RationaleModel(nn.Module): def __init__(self, model_name, num_labels): super().__init__() self.bert = AutoModel.from_pretrained(model_name, output_attentions=True) self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels) def forward(self, input_ids, attention_mask): outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) cls_output = outputs.last_hidden_state[:, 0, :] logits = self.classifier(cls_output) last_layer_attn = outputs.attentions[-1] # (batch, heads, seq, seq) cls_attn = last_layer_attn[:, :, 0, :] # (batch, heads, seq) cls_attn_avg = cls_attn.mean(dim=1) # (batch, seq) return logits, cls_attn_avg def evaluate_model(model, val_loader, criterion_cls, device): model.eval() total_val_loss = 0.0 all_preds = [] all_labels = [] all_probs = [] with torch.no_grad(): for batch in val_loader: input_ids, attention_mask, _, labels, _ = [b.to(device) for b in batch] logits, _ = model(input_ids, attention_mask) loss = criterion_cls(logits, labels) total_val_loss += loss.item() probs = torch.softmax(logits, dim=1) preds = torch.argmax(probs, dim=1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) all_probs.extend(probs.cpu().numpy()) avg_val_loss = total_val_loss / len(val_loader) # Overall metrics accuracy = accuracy_score(all_labels, all_preds) f1_macro = f1_score(all_labels, all_preds, average="macro") try: y_true_oh = np.eye(num_labels)[all_labels] auroc_ovr = roc_auc_score(y_true_oh, all_probs, multi_class="ovr") except Exception: auroc_ovr = -1.0 # Class-wise metrics class_wise_metrics = {} target_names = sorted(valid_labels, key=valid_labels.get) precision, recall, f1_per_class, _ = precision_recall_fscore_support(all_labels, all_preds, average=None, labels=[valid_labels[label_name] for label_name in target_names]) for i, label_name in enumerate(target_names): class_wise_metrics[f"{label_name}_precision"] = precision[i] class_wise_metrics[f"{label_name}_recall"] = recall[i] class_wise_metrics[f"{label_name}_f1"] = f1_per_class[i] # Per-class accuracy: of true class samples, how many were predicted correctly idx = np.array(all_labels) == valid_labels[label_name] if idx.sum() > 0: acc = (np.array(all_preds)[idx] == valid_labels[label_name]).sum() / idx.sum() else: acc = -1.0 class_wise_metrics[f"{label_name}_accuracy"] = acc # Class-wise AUROC try: binary_labels = (np.array(all_labels) == valid_labels[label_name]).astype(int) class_probs = np.array(all_probs)[:, valid_labels[label_name]] if len(np.unique(binary_labels)) > 1: class_wise_metrics[f"{label_name}_auroc"] = roc_auc_score(binary_labels, class_probs) else: class_wise_metrics[f"{label_name}_auroc"] = -1.0 except Exception: class_wise_metrics[f"{label_name}_auroc"] = -1.0 return avg_val_loss, accuracy, f1_macro, auroc_ovr, class_wise_metrics def train_model(model, train_loader, val_loader, num_epochs, device, lambda_attn=1.0, optimizer=None, learning_rate=2e-5, results_writer=None, results_file_handle=None, params=None): criterion_cls = nn.CrossEntropyLoss() criterion_kl = nn.KLDivLoss(reduction="batchmean") if optimizer is None: optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) for epoch in range(num_epochs): model.train() total_train_loss = 0.0 for batch in train_loader: input_ids, attention_mask, rationale_probs, labels, has_rationale = [b.to(device) for b in batch] optimizer.zero_grad() logits, model_attention = model(input_ids, attention_mask) loss_cls = criterion_cls(logits, labels) loss = loss_cls if has_rationale.any(): model_attn_batch = model_attention[has_rationale] rationale_batch = rationale_probs[has_rationale] log_model_attn = torch.log(model_attn_batch + 1e-8) loss_kl = criterion_kl(log_model_attn, rationale_batch) loss += lambda_attn * loss_kl loss.backward() optimizer.step() total_train_loss += loss.item() avg_train_loss = total_train_loss / len(train_loader) val_loss, val_acc, val_f1_macro, val_auroc_ovr, class_wise_metrics = evaluate_model(model, val_loader, criterion_cls, device) print(f"Epoch {epoch+1} | Train Loss: {avg_train_loss:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f} | Val F1 (Macro): {val_f1_macro:.4f} | Val AUROC (OvR): {val_auroc_ovr:.4f}") sorted_labels = sorted(valid_labels, key=valid_labels.get) for label_name in sorted_labels: print(f" {label_name}: P={class_wise_metrics[f'{label_name}_precision']:.4f}, R={class_wise_metrics[f'{label_name}_recall']:.4f}, F1={class_wise_metrics[f'{label_name}_f1']:.4f}, Acc={class_wise_metrics[f'{label_name}_accuracy']:.4f}, AUROC={class_wise_metrics[f'{label_name}_auroc']:.4f}") if results_writer and results_file_handle: row_data = [ params["learning_rate"], params["batch_size"], params["optimizer"], params["lambda"], epoch + 1, avg_train_loss, val_loss, val_acc, val_f1_macro, val_auroc_ovr ] for label_name in sorted_labels: row_data.extend([ class_wise_metrics[f"{label_name}_precision"], class_wise_metrics[f"{label_name}_recall"], class_wise_metrics[f"{label_name}_f1"], class_wise_metrics[f"{label_name}_accuracy"], class_wise_metrics[f"{label_name}_auroc"] ]) results_writer.writerow(row_data) results_file_handle.flush() os.fsync(results_file_handle.fileno()) # --- PREPARE DATASETS --- print("Generating attention vectors for training data...") train_df = generate_attention_vectors_from_rationales(train_df, tokenizer) print("Generating attention vectors for validation data...") val_df = generate_attention_vectors_from_rationales(val_df, tokenizer) train_dataset = RationaleDataset(train_df, tokenizer, max_length, label_mapping=valid_labels) val_dataset = RationaleDataset(val_df, tokenizer, max_length, label_mapping=valid_labels) # --- GRID SEARCH LOOP --- keys, values = zip(*param_grid.items()) param_combinations = [dict(zip(keys, v)) for v in itertools.product(*values)] results_file = "grid_results_detailed.csv" headers = ["learning_rate", "batch_size", "optimizer", "lambda", "epoch", "train_loss", "val_loss", "val_accuracy", "val_f1_macro", "val_auroc_ovr"] sorted_labels = sorted(valid_labels, key=valid_labels.get) for label_name in sorted_labels: headers.extend([f"{label_name}_precision", f"{label_name}_recall", f"{label_name}_f1", f"{label_name}_accuracy", f"{label_name}_auroc"]) with open(results_file, mode="w", newline="") as f: writer = csv.writer(f) writer.writerow(headers) for params in param_combinations: print("\nRunning:", params) learning_rate = params["learning_rate"] batch_size = params["batch_size"] optimizer_type = params["optimizer"] lambda_attn = params["lambda"] model = RationaleModel(model_name=model_name, num_labels=num_labels).to(device) if 'emoji_set' in locals() and len(emoji_set) > 0: model.bert.resize_token_embeddings(len(tokenizer)) if optimizer_type == "Adam": optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) else: raise ValueError("Unsupported optimizer") train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, generator=torch.Generator().manual_seed(13)) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) train_model( model=model, train_loader=train_loader, val_loader=val_loader, num_epochs=num_epochs, device=device, lambda_attn=lambda_attn, optimizer=optimizer, learning_rate=learning_rate, results_writer=writer, results_file_handle=f, params=params ) print("Grid search complete. Results saved to:", results_file)