import csv import os import numpy as np import pandas as pd import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from transformers import AutoTokenizer, AutoModelForSequenceClassification from sklearn.metrics import f1_score, roc_auc_score, accuracy_score, precision_recall_fscore_support 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 --- model_name = "bert-base-multilingual-cased" # Set your model name here num_epochs = 4 max_length = 128 num_labels = 3 learning_rate = 2e-5 batch_size = 64 optimizer_type = "Adam" lambda_attn = 0.6 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") test_df = pd.read_csv("test.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())] test_df = test_df[test_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.") # --- 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 = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels, output_attentions=True) def forward(self, input_ids, attention_mask): outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) logits = outputs.logits 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, valid_labels, num_labels): 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) all_labels_np = np.array(all_labels) all_preds_np = np.array(all_preds) all_probs_np = np.array(all_probs) accuracy = accuracy_score(all_labels_np, all_preds_np) f1_macro = f1_score(all_labels_np, all_preds_np, average="macro") try: y_true_oh = np.eye(num_labels)[all_labels_np] auroc_ovr = roc_auc_score(y_true_oh, all_probs_np, multi_class="ovr") except: auroc_ovr = -1.0 class_wise_metrics = {} target_names = sorted(valid_labels, key=valid_labels.get) label_indices = [valid_labels[label_name] for label_name in target_names] precision, recall, f1_per_class, support = precision_recall_fscore_support( all_labels_np, all_preds_np, labels=label_indices, average=None) for i, label_name in enumerate(target_names): label_id = valid_labels[label_name] 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] label_mask = all_labels_np == label_id correct_preds = np.sum((all_preds_np == label_id) & label_mask) total_label = np.sum(label_mask) if total_label > 0: class_wise_metrics[f"{label_name}_accuracy"] = correct_preds / total_label else: class_wise_metrics[f"{label_name}_accuracy"] = -1.0 try: binary_labels = (all_labels_np == label_id).astype(int) class_probs = all_probs_np[:, label_id] 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: 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): 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, valid_labels, num_labels) 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 = [ learning_rate, batch_size, optimizer_type, lambda_attn, 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()) # --- OUTPUT FOLDERS --- csv_output_dir = "csv_outputs" os.makedirs(csv_output_dir, exist_ok=True) results_file = os.path.join(csv_output_dir, "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 in sorted_labels: headers.extend([f"{label}_precision", f"{label}_recall", f"{label}_f1", f"{label}_accuracy", f"{label}_auroc"]) # --- INITIALIZE TOKENIZER & ADD EMOJIS --- tokenizer = AutoTokenizer.from_pretrained(model_name) emoji_path = "emoji.csv" 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}") # --- PREPARE DATASETS --- print("Generating attention vectors for training data...") train_df_model = generate_attention_vectors_from_rationales(train_df.copy(), tokenizer) print("Generating attention vectors for validation data...") val_df_model = generate_attention_vectors_from_rationales(val_df.copy(), tokenizer) train_dataset = RationaleDataset(train_df_model, tokenizer, max_length, label_mapping=valid_labels) val_dataset = RationaleDataset(val_df_model, tokenizer, max_length, label_mapping=valid_labels) 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) # --- CSV Setup --- with open(results_file, mode="w", newline="") as f: writer = csv.writer(f) writer.writerow(headers) 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)) optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) 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 ) # Save final model and tokenizer model.bert.save_pretrained("model_outputs") tokenizer.save_pretrained("model_outputs") print(f"Final model and tokenizer saved to model_outputs")