import os import sys import time import json import joblib import argparse from functools import partial from typing import Union, List, Dict, Optional, Any import torch import torch.nn as nn import torch.distributed as dist from torch.utils.data import DataLoader import numpy as np import pandas as pd from tqdm import tqdm from transformers import ( set_seed, BertTokenizerFast, GPT2TokenizerFast, GPT2LMHeadModel, GPT2Model, AutoConfig, BertModel, ) from transformers.models.bert.modeling_bert import BertEncoder from datasets import DatasetDict from datasets import Dataset as Dataset2 from ScanDL2.scandl_module.original_scandl.sp_rounding import denoised_fn_round from ScanDL2.scandl_module.original_scandl.utils import dist_util, logger from ScanDL2.scandl_module.original_scandl.utils.nn import * from ScanDL2.scandl2_utils import text_dataset_loader, FixdurDataset from ScanDL2.scandl_module.scripts.sp_load_celer_zuco import _collate_batch_helper from ScanDL2.scandl_module.scripts.sp_basic_utils import ( load_defaults_config, create_model_and_diffusion, add_dict_to_argparser, args_to_dict, ) from ScanDL2.fix_dur_module.model_seq2seq import Seq2SeqModel from ScanDL2.fix_dur_module.utils_data import aggregate_input_embeddings, padding_and_mask_seq2seq from ScanDL2.PATHS import ( SENT_SCANDL_MODULE, SENT_FIXDUR_MODULE, PAR_SCANDL_MODULE, PAR_FIXDUR_MODULE, ) class ScanDL2(nn.Module): def __init__( self, text_type: str = "sentence", # sentence, paragraph bsz: Optional[int] = 2, save: Optional[str] = None, filename: Optional[str] = None, ): super(ScanDL2, self).__init__() self.save = save self.filename = filename # initialize the ScanDL module and the Fixdur Module self.scandl_module = ScanDLModule( text_type=text_type, bsz=bsz, ) self.fixdur_module = FixdurModule( text_type=text_type, bsz=bsz, ) def forward( self, texts: Union[str, List[str]], ): # check if the input is in the correct format self._validate_inputs(texts=texts) # get the fixation location predictions from the ScanDL module scandl_module_output = self.scandl_module(texts=texts) # get the fixation duration predictions from the Fixdur module fixdur_module_output = self.fixdur_module(scandl_module_output=scandl_module_output) if self.save is not None: filename = self.filename if self.filename is not None else f"scandl2_outputs.json" filename = f"{filename}.json" if not filename.endswith(".json") else filename self._save_results(results=fixdur_module_output, filename=filename) return fixdur_module_output def _save_results(self, results, filename): if not os.path.exists(self.save): os.makedirs(self.save) with open(os.path.join(self.save, filename), "w") as f: json.dump(results, f) print(f"--- ScanDL 2.0 outputs saved to {os.path.join(self.save, filename)}.") def _validate_inputs(self, texts: Union[str, List[str]]) -> None: if not isinstance(texts, (str, list)) or ( isinstance(texts, list) and not all(isinstance(t, str) for t in texts) ): raise TypeError("Invalid input: 'texts' must be of type 'str' or 'List[str]'.") class ScanDLModule(nn.Module): def __init__( self, text_type: str, # sentence, paragraph bsz: int, ): super(ScanDLModule, self).__init__() base_path = os.path.dirname(__file__) if text_type == "paragraph": self.path_to_config = os.path.join(base_path, "config_emtec.json") self.path_to_scandl_module = PAR_SCANDL_MODULE elif text_type == "sentence": self.path_to_config = os.path.join(base_path, "config.json") self.path_to_scandl_module = SENT_SCANDL_MODULE else: raise NotImplementedError(f"Text type {text_type} not implemented.") # get the args self.args = self._get_args() self.args.batch_size = bsz # seting up the environment dist_util.setup_dist() logger.configure() self.world_size = dist.get_world_size() or 1 self.rank = dist.get_rank() or 0 # set_seed(self.args.seed2) # load the tokenizer self.tokenizer = self._load_tokenizer() # load the ScanDL module and the Diffusion self.scandl_module, self.diffusion = self._load_scandl_module( path_to_scandl_module=self.path_to_scandl_module ) self.sn_sp_repr_embedding = self._get_sn_sp_repr_emb() def forward( self, texts: Union[str, List[str]], ) -> Dict[str, Union[List[List[str]], List[List[int]], List[str]]]: data_loader = self._preprocess_text(texts=texts) predicted_sp_words, predicted_sp_ids = [], [] original_sn = [] print("\t\t### ScanDL Module generates fixation locations ...") unique_idx = list() idx_ctr = 0 for batch_idx, batch in tqdm(enumerate(data_loader)): mask = batch["mask"].to(dist_util.dev()) sn_sp_repr = batch["sn_sp_repr"].to(dist_util.dev()) sn_input_ids = batch["sn_input_ids"].to(dist_util.dev()) indices_pos_enc = batch["indices_pos_enc"].to(dist_util.dev()) sn_repr_len = batch["sn_repr_len"].to(dist_util.dev()) words_for_mapping = batch["words_for_mapping"] sn_sp_emb, pos_enc, sn_input_ids_emb = self.scandl_module.get_embeds( sn_sp_repr=sn_sp_repr, sn_input_ids=sn_input_ids, indices_pos_enc=indices_pos_enc, ) x_start = sn_sp_emb noise = torch.randn_like(x_start) mask = torch.broadcast_to(mask.unsqueeze(dim=-1), x_start.shape).to(dist_util.dev()) x_noised = torch.where(mask == 0, x_start, noise) self.args.use_ddim = False step_gap = 1 sample_fn = ( self.diffusion.p_sample_loop if not self.args.use_ddim else self.diffusion.ddim_sample_loop ) sample_shape = (x_start.shape[0], self.args.seq_len, self.args.hidden_dim) subwords = [self.tokenizer.convert_ids_to_tokens(i) for i in sn_input_ids] samples = sample_fn( model=self.scandl_module, shape=sample_shape, noise=x_noised, sn_input_ids_emb=sn_input_ids_emb, pos_enc=pos_enc, mask_sn_padding=None, mask_transformer_att=None, clip_denoised=self.args.clip_denoised, denoised_fn=partial(denoised_fn_round, self.args, self.sn_sp_repr_embedding), model_kwargs=None, top_p=self.args.top_p, clamp_step=self.args.clamp_step, clamp_first=self.args.clamp_first_bool, mask=mask, x_start=x_start, gap=step_gap, ) sample = samples[-1] logits = self.scandl_module.get_logits(sample) cands = torch.topk(logits, k=1, dim=-1) for instance_idx, (pred_seq, orig_words, sn_len) in enumerate( zip(cands.indices, words_for_mapping, sn_repr_len) ): pred_seq_sp = pred_seq[sn_len:] words_split = orig_words.split() predicted_sp = [words_split[i] for i in pred_seq_sp] pred_sp_ids = [e.item() for e in pred_seq_sp] # cut off trailing pad tokens while len(predicted_sp) > 1 and predicted_sp[-1] == "[PAD]": predicted_sp.pop() while len(pred_sp_ids) > 1 and pred_sp_ids[-1] == self.args.seq_len - 1: pred_sp_ids.pop() while len(words_split) > 1 and words_split[-1] == "[PAD]": words_split.pop() # remove CLS and SEP tokens from predictions if predicted_sp[0] == "[CLS]": predicted_sp = predicted_sp[1:] pred_sp_ids = pred_sp_ids[1:] if predicted_sp[-1] == "[SEP]": predicted_sp = predicted_sp[:-1] pred_sp_ids = pred_sp_ids[:-1] words_split = words_split[1:-1] # filter out erroneously predicted PAD tokens (they will raise an error in the fixdur module) pred_sp_ids, predicted_sp = self._remove_special_tokens( predicted_sp_ids=pred_sp_ids, predicted_sp_words=predicted_sp, token="[PAD]", ) pred_sp_ids, predicted_sp = self._remove_special_tokens( predicted_sp_ids=pred_sp_ids, predicted_sp_words=predicted_sp, token="[CLS]", ) pred_sp_ids, predicted_sp = self._remove_special_tokens( predicted_sp_ids=pred_sp_ids, predicted_sp_words=predicted_sp, token="[SEP]", ) predicted_sp_words.append(predicted_sp) predicted_sp_ids.append(pred_sp_ids) original_sn.append(words_split) idx_ctr += 1 unique_idx.append(idx_ctr) predictions = { "predicted_sp_words": predicted_sp_words, "predicted_sp_ids": predicted_sp_ids, "original_sn": original_sn, "unique_idx": unique_idx, } return predictions def _remove_special_tokens( self, predicted_sp_ids: List[int], predicted_sp_words: List[str], token: str, # '[CLS]' or '[SEP]' or '[PAD]' ): filtered_sp_ids, filtered_sp_words = [], [] for sp_word, sp_id in zip(predicted_sp_words, predicted_sp_ids): if sp_word != token: filtered_sp_ids.append(sp_id) filtered_sp_words.append(sp_word) return filtered_sp_ids, filtered_sp_words def _preprocess_text( self, texts: Union[str, List[str]], ): data = { "mask": [], "sn_sp_repr": [], "sn_input_ids": [], "indices_pos_enc": [], "words_for_mapping": [], "sn_repr_len": [], } if isinstance(texts, str): texts = [texts] for sn_idx, sn in enumerate(texts): if sn.startswith("[CLS]") and sn.endswith("[SEP]"): sn = sn elif sn.startswith("[CLS]"): sn = sn + " [SEP]" elif sn.endswith("[SEP]"): sn = "[CLS] " + sn else: sn = "[CLS] " + sn + " [SEP]" encoded_sn = self.tokenizer.encode_plus( sn.split(), add_special_tokens=False, padding=False, return_attention_mask=False, is_split_into_words=True, truncation=False, ) if len(encoded_sn) > self.args.seq_len / 2: print(f"Sentence {sn} is too long. Continue.") sn_word_ids = encoded_sn.word_ids() sn_input_ids = encoded_sn["input_ids"] sn_sp_repr = sn_word_ids mask = [0] * len(sn_word_ids) indices_pos_enc = list(range(0, len(sn_word_ids))) + list( range(0, self.args.seq_len - len(sn_word_ids)) ) words_for_mapping = sn.split() + (self.args.seq_len - len(sn.split())) * ["[PAD]"] data["mask"].append(mask) data["sn_sp_repr"].append(sn_sp_repr) data["sn_input_ids"].append(sn_input_ids) data["indices_pos_enc"].append(indices_pos_enc) data["words_for_mapping"].append(" ".join(words_for_mapping)) data["sn_repr_len"].append(len(sn_word_ids)) # padding data["mask"] = _collate_batch_helper( examples=data["mask"], pad_token_id=1, max_length=self.args.seq_len, ) data["sn_sp_repr"] = _collate_batch_helper( examples=data["sn_sp_repr"], pad_token_id=self.args.seq_len - 1, max_length=self.args.seq_len, ) data["sn_input_ids"] = _collate_batch_helper( examples=data["sn_input_ids"], pad_token_id=self.tokenizer.pad_token_id, max_length=self.args.seq_len, ) split = "inference" dataset = Dataset2.from_dict(data) dataset_dict = DatasetDict() dataset_dict[split] = dataset data_loader = text_dataset_loader( data=dataset_dict, data_args=self.args, split=split, deterministic=True, ) return data_loader def _load_scandl_module( self, path_to_scandl_module: str, ): logger.log("### Loading ScanDL Diffusion Module ...") scandl_module, diffusion = create_model_and_diffusion( **args_to_dict(self.args, load_defaults_config(config_path=self.path_to_config).keys()) ) # TODO Name scandl module, not model scandl_module.load_state_dict( dist_util.load_state_dict( os.path.join(self.path_to_scandl_module, "ema_0.9999_080000.pt"), map_location="cpu" ) ) pytorch_total_params = sum(p.numel() for p in scandl_module.parameters()) logger.log(f"### Total number of parameters: {pytorch_total_params}") scandl_module.eval().requires_grad_(False).to(dist_util.dev()) return scandl_module, diffusion def _get_sn_sp_repr_emb(self): sn_sp_repr_embedding = nn.Embedding( num_embeddings=self.args.hidden_t_dim, embedding_dim=self.args.hidden_dim, _weight=self.scandl_module.sn_sp_repr_embedding.weight.clone().cpu(), ) return sn_sp_repr_embedding def _get_args(self): args = self._get_parser().parse_args() # load the training arguments with open(os.path.join(self.path_to_scandl_module, "training_args.json")) as f: training_args = json.load(f) training_args["batch_size"] = args.batch_size args.__dict__.update(training_args) if args.clamp_first == "yes": args.clamp_first_bool = True else: args.clamp_first_bool = False # TODO self.args.clamp_first_bool as argument # set mask_padding to False args.mask_padding = False return args def _load_tokenizer(self): tokenizer = BertTokenizerFast.from_pretrained(self.args.config_name) self.args.vocab_size = tokenizer.vocab_size return tokenizer def _get_parser(self) -> argparse.ArgumentParser: defaults = dict( model_path="", step=0, out_dir="", top_p=0, clamp_first="yes", test_set_sns="mixed", atten_vis=False, notes="-", tsne_vis=False, sp_vis=False, no_inst=0, atten_vis_sp=False, load_ids="-", load_test_data="-", setting="-", fold=0, ) decode_defaults = dict( split="valid", clamp_step=0, seed2=105, clip_denoised=False, ) defaults.update(load_defaults_config(config_path=self.path_to_config)) defaults.update(decode_defaults) parser = argparse.ArgumentParser() add_dict_to_argparser(parser, defaults) return parser class FixdurModule(nn.Module): def __init__( self, text_type: str, # sentence, paragraph bsz: Optional[int] = 2, ): super(FixdurModule, self).__init__() base_path = os.path.dirname(__file__) if text_type == "paragraph": self.path_to_config = os.path.join(base_path, "config_emtec.json") self.path_to_fixdur_module = PAR_FIXDUR_MODULE elif text_type == "sentence": self.path_to_config = os.path.join(base_path, "config.json") self.path_to_fixdur_module = SENT_FIXDUR_MODULE else: raise NotImplementedError(f"Text type {text_type} not implemented.") self.bsz = bsz self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.config = load_defaults_config(config_path=self.path_to_config) self.args = self._get_args(config=self.config) self.hyperparameters = self._get_hyperparams( path_to_fixdur_module=self.path_to_fixdur_module ) # load GPT-2 model and tokenizer, and BERT embeddings self.gpt2_model, self.tokenizer, self.bert_embeddings = self._load_gpt_and_bert( config=self.config ) # load the Fixdur module and the MinMax Scaler self.fixdur_module = self._load_fixdur_module() self.scaler = self._load_scaler() def _get_args(self, config: Dict[str, Any]) -> Dict[str, Any]: args = { "max_length": config["seq_len"], "normalize": True, "output_attentions": False, "bsz": self.bsz, "corpus": config["corpus"], "sp_pad_token": config["seq_len"] - 1, } return args def _get_hyperparams(self, path_to_fixdur_module: str) -> Dict[str, Any]: with open(os.path.join(path_to_fixdur_module, "hyperparameters.json")) as f: return json.load(f) def _load_gpt_and_bert(self, config: Dict[str, Any]): """ Load GPT-2 and GPT-2 tokenizer to get the contextualized embeddings. Load BERT model (for embeddings of CLS and PAD tokens) """ # GPT-2 gpt_config_name = config["gpt_config_name"] tokenizer = GPT2TokenizerFast.from_pretrained(gpt_config_name, add_prefix_space=True) gpt2_model = GPT2Model.from_pretrained(gpt_config_name) tokenizer.pad_token = tokenizer.eos_token # freeze parameters for param in gpt2_model.parameters(): param.requires_grad = False # BERT bert_config_name = config["config_name"] bert_embeddings = BertModel.from_pretrained(bert_config_name).embeddings.word_embeddings # freeze parameters for param in bert_embeddings.parameters(): param.requires_grad = False return gpt2_model, tokenizer, bert_embeddings def _load_fixdur_module(self): fixdur_module_config = AutoConfig.from_pretrained("bert-base-cased") fixdur_module_config.num_attention_heads = self.hyperparameters["num_heads"] fixdur_module_config.num_hidden_layers = self.hyperparameters["num_layers"] fixdur_module = Seq2SeqModel( config=fixdur_module_config, output_dim=self.args["max_length"], num_linear=self.hyperparameters["num_linear"], dropout=self.hyperparameters["dropout"], ) fixdur_module.load_state_dict( torch.load( os.path.join(self.path_to_fixdur_module, "seq2seq_fixdur.pt"), map_location=self.device, ) ) fixdur_module.eval() fixdur_module.to(self.device) return fixdur_module def _load_scaler(self): scaler = joblib.load(os.path.join(self.path_to_fixdur_module, "min_max_scaler.pkl")) return scaler def _prepare_data( self, scandl_module_output: Dict[str, Union[List[List[str]], List[List[int]], List[str]]], ): data_dict = { "sp_embeddings": [], "attention_mask": [], "unique_idx": [], } for idx in range(len(scandl_module_output["predicted_sp_words"])): sn_words = scandl_module_output["original_sn"][idx] sp_ids = scandl_module_output["predicted_sp_ids"][idx] unique_id = scandl_module_output["unique_idx"][idx] # make the scanpath ids start at 0 for re-ordering of the embeddings sp_ids = [i - 1 for i in sp_ids] sp_words = scandl_module_output["predicted_sp_words"][idx] # get the sentence encoding sn_enc = self.tokenizer( sn_words, add_special_tokens=False, return_tensors="pt", is_split_into_words=True, ) sn_word_ids = torch.Tensor(sn_enc.word_ids()) # get the embeddings with torch.no_grad(): last_hidden = self.gpt2_model(sn_enc.input_ids).last_hidden_state # aggregate the embeddings to word level sn_embeddings = aggregate_input_embeddings( embeddings=last_hidden, word_ids=sn_word_ids, aggregate="mean", ) # convert sp_ids to tensor sp_ids = torch.Tensor(sp_ids).long() # re-order the embeddings as scanpath try: sp_embeddings = sn_embeddings[:, sp_ids, :] except: breakpoint() # pad the embeddings to max input length and get the attentino mask sp_embeddings_padded, attention_mask = padding_and_mask_seq2seq( sp_embeddings=sp_embeddings, bert_embeddings=self.bert_embeddings, max_length=self.args["max_length"], inference=True, ) data_dict["sp_embeddings"].append(sp_embeddings_padded) data_dict["attention_mask"].append(attention_mask) data_dict["unique_idx"].append(unique_id) return data_dict def forward( self, scandl_module_output: Dict[str, Union[List[List[str]], List[List[int]], List[str]]], ) -> Dict[str, Union[List[List[str]], List[List[int]], List[str], List[List[float]]]]: output_dict = { "predicted_sp_words": [], "predicted_sp_ids": [], "original_sn": [], "predicted_fix_durs": [], "unique_idx": [], } data_df = pd.DataFrame(scandl_module_output) data_dict = self._prepare_data( scandl_module_output=scandl_module_output, ) dataset = FixdurDataset(data=data_dict) data_loader = DataLoader( dataset, batch_size=self.bsz, shuffle=False, ) print("\t\t### FixDur Module generates fixation durations ...") for batch_idx, batch in tqdm(enumerate(data_loader)): sp_embeddings = batch["sp_embeddings"].squeeze(1).to(self.device) attention_mask = batch["attention_mask"].squeeze(1).to(self.device) unique_indices = batch["unique_idx"] out = self.fixdur_module( sp_embeddings=sp_embeddings, attention_mask=attention_mask, output_attentions=self.args["output_attentions"], ) # scale the output back to the original range out_transformed = self.scaler.inverse_transform(out.detach().cpu().numpy()) out_transformed_rounded = np.round(out_transformed, 2) # iterate over the individual predictions for out_idx, out_instance in enumerate(out_transformed_rounded): predicted_fix_durs = out_instance unique_idx = unique_indices[out_idx].item() # find predicted_sp_words, predicted_sp_ids, and original_sn in data_df conditioned on unique_idx predicted_sp_words = data_df.loc[ data_df["unique_idx"] == unique_idx, "predicted_sp_words" ].values[0] predicted_sp_ids = data_df.loc[ data_df["unique_idx"] == unique_idx, "predicted_sp_ids" ].values[0] original_sn = data_df.loc[ data_df["unique_idx"] == unique_idx, "original_sn" ].values[0] sp_len = len(predicted_sp_ids) # cut off the predicted_fix_durs to the length of the scanpath # the predicted fixation durations still contain predictions for the CLS and SEP token as well pred_fix_durs = predicted_fix_durs[: sp_len + 2].tolist()[1:-1] pred_fix_durs = [round(d, 2) for d in pred_fix_durs] # add to output_dict output_dict["predicted_sp_words"].append(predicted_sp_words) output_dict["predicted_sp_ids"].append(predicted_sp_ids) output_dict["original_sn"].append(original_sn) output_dict["predicted_fix_durs"].append(pred_fix_durs) output_dict["unique_idx"].append(unique_idx) print(f"fixdur original sn: {original_sn}") return output_dict