<kaggle_start><data_title>kenh14-small<data_name>kenh14small
<code>import yaml
import os
import abc
import time
import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
import pandas as pd
import pickle
import random
import re
from tqdm import tqdm
from tensorflow.keras import layers
from tensorflow.keras import backend as K
from collections import namedtuple

# detect and init the TPU
tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
# instantiate a distribution strategy
tpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)


def load_yaml(filename):
    """Load yaml file
    Args:
        filename
    Returns:
        dict
    """
    try:
        with open(filename, "r") as f:
            config = yaml.load(f, yaml.SafeLoader)
        return config
    except FileNotFoundError:
        raise
    except Exception as e:
        raise IOError(f"load {filename} error!")


def flat_config(config):
    """Flat config load tu yaml file to flat dict
    Args:
        config (dict): config load tu yaml file
    Returns:
        dict
    """
    f_config = {}
    category = config.keys()
    for cate in category:
        for key, val in config[cate].items():
            f_config[key] = val
    return f_config


def check_type(config):
    int_parameters = [
        "word_size",
        "his_size",
        "title_size",
        "body_size",
        "npratio",
        "word_emb_dim",
        "attention_hidden_dim",
        "epochs",
        "batch_size",
        "show_step",
        "save_epoch",
        "head_num",
        "head_dim",
        "user_num",
        "filter_num",
        "window_size",
        "gru_unit",
        "user_emb_dim",
        "vert_emb_dim",
        "subvert_emb_dim",
    ]
    for param in int_parameters:
        if param in config and not isinstance(config[param], int):
            raise TypeError("Parameters {0} must be int".format(param))
    float_parameters = ["learning_rate", "dropout"]
    for param in float_parameters:
        if param in config and not isinstance(config[param], float):
            raise TypeError("Parameters {0} must be float".format(param))
        str_parameters = [
            "wordEmb_file",
            "wordDict_file",
            "userDict_file",
            "vertDict_file",
            "subvertDict_file",
            "method",
            "loss",
            "optimizer",
            "cnn_activation",
            "dense_activation" "type",
        ]
    for param in str_parameters:
        if param in config and not isinstance(config[param], str):
            raise TypeError("Parameters {0} must be str".format(param))
    list_parameters = ["layer_sizes", "activation"]
    for param in list_parameters:
        if param in config and not isinstance(config[param], list):
            raise TypeError("Parameters {0} must be list".format(param))
    bool_parameters = ["support_quick_scoring"]
    for param in bool_parameters:
        if param in config and not isinstance(config[param], bool):
            raise TypeError("Parameters {0} must be bool".format(param))


def check_nn_config(f_config):
    """Check neural net config
    Args:
        f_config (dict): file config duoc flat tu yaml file
    Raises:
        ValueError: Neu params bi sai -> Raise error
    """
    if f_config["model_type"] in ["nrms", "NRMS"]:
        required_parameters = [
            "title_size",
            "his_size",
            "wordEmb_file",
            "wordDict_file",
            "userDict_file",
            "npratio",
            "data_format",
            "word_emb_dim",
            "head_num",
            "head_dim",
            "attention_hidden_dim",
            "loss",
            "data_format",
            "dropout",
        ]
    else:
        required_parameters = []
    # check
    for param in required_parameters:
        if param not in f_config:
            raise ValueError("Parameter {0} must be set!".format(param))
    if f_config["model_type"] in ["nrms", "NRMS"]:
        if f_config["data_format"] != "news":
            raise ValueError(
                "Voi NRMS model, dataformat phai la news, dua cai {0} vao lam gi".format(
                    f_config["data_format"]
                )
            )
    check_type(f_config)


def get_hparams(**kwargs):
    return namedtuple("GenericDict", kwargs.keys())(**kwargs)


def create_hparams(flags):
    """Create model's params
    Args:
        flags (dict): Dict co requirement
    Returns:
        object: namedtuple
    """
    return get_hparams(
        # data
        data_format=flags.get("data_format", None),
        iterator_type=flags.get("iterator_type", None),
        support_quick_scoring=flags.get("support_quick_scoring", False),
        wordEmb_file=flags.get("wordEmb_file", None),
        wordDict_file=flags.get("wordDict_file", None),
        userDict_file=flags.get("userDict_file", None),
        vertDict_file=flags.get("vertDict_file", None),
        subvertDict_file=flags.get("subvertDict_file", None),
        # models
        title_size=flags.get("title_size", None),
        body_size=flags.get("body_size", None),
        word_emb_dim=flags.get("word_emb_dim", None),
        word_size=flags.get("word_size", None),
        user_num=flags.get("user_num", None),
        vert_num=flags.get("vert_num", None),
        subvert_num=flags.get("subvert_num", None),
        his_size=flags.get("his_size", None),
        npratio=flags.get("npratio"),
        dropout=flags.get("dropout", 0.0),
        attention_hidden_dim=flags.get("attention_hidden_dim", 200),
        # nrms
        head_num=flags.get("head_num", 4),
        head_dim=flags.get("head_dim", 100),
        # train
        learning_rate=flags.get("learning_rate", 0.001),
        loss=flags.get("loss", None),
        optimizer=flags.get("optimizer", "adam"),
        epochs=flags.get("epochs", 10),
        batch_size=flags.get("batch_size", 1),
        # show info
        show_step=flags.get("show_step", 1),
        metrics=flags.get("metrics", None),
    )


def prepare_hparams(yaml_file=None, **kwargs):
    """Prepare hyperparams and make sure it's ok
    Args:
        yaml_file: path to yaml file
    Returns:
        TF Hyperparams object (tf.contrib.training.HParams)
    """
    if yaml_file is not None:
        config = load_yaml(yaml_file)
        config = flat_config(config)
    else:
        config = {}
    config.update(kwargs)
    check_nn_config(config)
    return create_hparams(config)


yaml_file = "/kaggle/input/kenh14small/nrms.yaml"
wordEmb_file = "/kaggle/input/mindlike-final/embedding.npy"
wordDict_file = "/kaggle/input/mindlike-final/word_dict.pkl"
userDict_file = "/kaggle/input/mindlike-final/uid2index.pkl"
epochs = 1
seed = 42
batch_size = 1024
hparams = prepare_hparams(
    yaml_file,
    his_size=25,
    wordEmb_file=wordEmb_file,
    wordDict_file=wordDict_file,
    userDict_file=userDict_file,
    batch_size=batch_size,
    epochs=epochs,
    show_step=10,
)
print(hparams)
from sklearn.metrics import (
    roc_auc_score,
    log_loss,
    mean_squared_error,
    accuracy_score,
    f1_score,
)


def mrr_score(y_true, y_score):
    """Computing mrr score metric.
    Args:
        y_true (np.ndarray): Ground-truth labels.
        y_score (np.ndarray): Predicted labels.
    Returns:
        numpy.ndarray: mrr scores.
    """
    order = np.argsort(y_score)[::-1]
    y_true = np.take(y_true, order)
    rr_score = y_true / (np.arange(len(y_true)) + 1)
    return np.sum(rr_score) / np.sum(y_true)


def ndcg_score(y_true, y_score, k=10):
    """Computing ndcg score metric at k.
    Args:
        y_true (np.ndarray): Ground-truth labels.
        y_score (np.ndarray): Predicted labels.
    Returns:
        numpy.ndarray: ndcg scores.
    """
    best = dcg_score(y_true, y_true, k)
    actual = dcg_score(y_true, y_score, k)
    return actual / best


def hit_score(y_true, y_score, k=10):
    """Computing hit score metric at k.
    Args:
        y_true (np.ndarray): ground-truth labels.
        y_score (np.ndarray): predicted labels.
    Returns:
        np.ndarray: hit score.
    """
    ground_truth = np.where(y_true == 1)[0]
    argsort = np.argsort(y_score)[::-1][:k]
    for idx in argsort:
        if idx in ground_truth:
            return 1
    return 0


def dcg_score(y_true, y_score, k=10):
    """Computing dcg score metric at k.
    Args:
        y_true (np.ndarray): Ground-truth labels.
        y_score (np.ndarray): Predicted labels.
    Returns:
        np.ndarray: dcg scores.
    """
    k = min(np.shape(y_true)[-1], k)
    order = np.argsort(y_score)[::-1]
    y_true = np.take(y_true, order[:k])
    gains = 2**y_true - 1
    discounts = np.log2(np.arange(len(y_true)) + 2)
    return np.sum(gains / discounts)


def cal_metric(labels, preds, metrics):
    print(labels)
    print(preds)
    """Calculate metrics.
    Available options are: `auc`, `rmse`, `logloss`, `acc` (accurary), `f1`, `mean_mrr`,
    `ndcg` (format like: ndcg@2;4;6;8), `hit` (format like: hit@2;4;6;8), `group_auc`.
    Args:
        labels (array-like): Labels.
        preds (array-like): Predictions.
        metrics (list): List of metric names.
    Return:
        dict: Metrics.
    Examples:
        >>> cal_metric(labels, preds, ["ndcg@2;4;6", "group_auc"])
        {'ndcg@2': 0.4026, 'ndcg@4': 0.4953, 'ndcg@6': 0.5346, 'group_auc': 0.8096}
    """
    res = {}
    for metric in metrics:
        if metric == "auc":
            auc = roc_auc_score(np.asarray(labels), np.asarray(preds))
            res["auc"] = round(auc, 4)
        elif metric == "rmse":
            rmse = mean_squared_error(np.asarray(labels), np.asarray(preds))
            res["rmse"] = np.sqrt(round(rmse, 4))
        elif metric == "logloss":
            # avoid logloss nan
            preds = [max(min(p, 1.0 - 10e-12), 10e-12) for p in preds]
            logloss = log_loss(np.asarray(labels), np.asarray(preds))
            res["logloss"] = round(logloss, 4)
        elif metric == "acc":
            pred = np.asarray(preds)
            pred[pred >= 0.5] = 1
            pred[pred < 0.5] = 0
            acc = accuracy_score(np.asarray(labels), pred)
            res["acc"] = round(acc, 4)
        elif metric == "f1":
            pred = np.asarray(preds)
            pred[pred >= 0.5] = 1
            pred[pred < 0.5] = 0
            f1 = f1_score(np.asarray(labels), pred)
            res["f1"] = round(f1, 4)
        elif metric == "mean_mrr":
            mean_mrr = np.mean(
                [
                    mrr_score(each_labels, each_preds)
                    for each_labels, each_preds in zip(labels, preds)
                ]
            )
            res["mean_mrr"] = round(mean_mrr, 4)
        elif metric.startswith("ndcg"):  # format like:  ndcg@2;4;6;8
            ndcg_list = [1, 2]
            ks = metric.split("@")
            if len(ks) > 1:
                ndcg_list = [int(token) for token in ks[1].split(";")]
            for k in ndcg_list:
                ndcg_temp = np.mean(
                    [
                        ndcg_score(each_labels, each_preds, k)
                        for each_labels, each_preds in zip(labels, preds)
                    ]
                )
                res["ndcg@{0}".format(k)] = round(ndcg_temp, 4)
        elif metric.startswith("hit"):  # format like:  hit@2;4;6;8
            hit_list = [1, 2]
            ks = metric.split("@")
            if len(ks) > 1:
                hit_list = [int(token) for token in ks[1].split(";")]
            for k in hit_list:
                hit_temp = np.mean(
                    [
                        hit_score(each_labels, each_preds, k)
                        for each_labels, each_preds in zip(labels, preds)
                    ]
                )
                res["hit@{0}".format(k)] = round(hit_temp, 4)
        elif metric == "group_auc":
            group_auc = np.mean(
                [
                    roc_auc_score(each_labels, each_preds)
                    for each_labels, each_preds in zip(labels, preds)
                ]
            )
            res["group_auc"] = round(group_auc, 4)
        else:
            raise ValueError("Metric {0} not defined".format(metric))
    return res


def word_tokenize(sent):
    """Split sentence into word list using regex.
    Args:
        sent (str): Input sentence
    Return:
        list: word list
    """
    pat = re.compile(r"[\w]+|[.,!?;|]")
    if isinstance(sent, str):
        return pat.findall(sent.lower())
    else:
        return []


def newsample(news, ratio):
    """Sample ratio samples from news list.
    If length of news is less than ratio, pad zeros.
    Args:
        news (list): input news list
        ratio (int): sample number
    Returns:
        list: output of sample list.
    """
    if ratio > len(news):
        return news + [0] * (ratio - len(news))
    else:
        return random.sample(news, ratio)


def abstractmethod(funcobj):
    """A decorator indicating abstract methods.
    Requires that the metaclass is ABCMeta or derived from it.  A
    class that has a metaclass derived from ABCMeta cannot be
    instantiated unless all of its abstract methods are overridden.
    The abstract methods can be called using any of the normal
    'super' call mechanisms.  abstractmethod() may be used to declare
    abstract methods for properties and descriptors.
    Usage:
        class C(metaclass=ABCMeta):
            @abstractmethod
            def my_abstract_method(self, ...):
                ...
    """
    funcobj.__isabstractmethod__ = True
    return funcobj


class BaseIterator(object):
    """Abstract base iterator class"""

    @abstractmethod
    def parser_one_line(self, line):
        """Abstract method. Parse one string line into feature values.
        Args:
            line (str): A string indicating one instance.
        """
        pass

    @abstractmethod
    def load_data_from_file(self, infile):
        """Abstract method. Read and parse data from a file.
        Args:
            infile (str): Text input file. Each line in this file is an instance.
        """
        pass

    @abstractmethod
    def _convert_data(self, labels, features):
        pass

    @abstractmethod
    def gen_feed_dict(self, data_dict):
        """Abstract method. Construct a dictionary that maps graph elements to values.
        Args:
            data_dict (dict): A dictionary that maps string name to numpy arrays.
        """
        pass


def load_dict(file_path):
    with open(file_path, "rb") as f:
        return pickle.load(f)


class MINDIterator(BaseIterator):
    """Train data loader for NRMS model.
    The model require a special type of data format, where each instance contains a label, impresion id, user id,
    the candidate news articles and user's clicked news article. Articles are represented by title words,
    body words, verts and subverts.
    Iterator will not load the whole data into memory. Instead, it loads data into memory
    per mini-batch, so that large files can be used as input data.
    Attributes:
        col_spliter (str): column spliter in one line.
        ID_spliter (str): ID spliter in one line.
        batch_size (int): the samples num in one batch.
        title_size (int): max word num in news title.
        his_size (int): max clicked news num in user click history.
        npratio (int): negaive and positive ratio used in negative sampling. -1 means no need of negtive sampling.
    """

    def __init__(
        self,
        hparams,
        npratio=-1,
        col_spliter="\t",
        ID_spliter="%",
    ):
        """Initialize an iterator. Create necessary placeholders for the model.
        Args:
            hparams (object): Global hyper-parameters. Some key setttings such as head_num and head_dim are there.
            npratio (int): negaive and positive ratio used in negative sampling. -1 means no need of negtive sampling.
            col_spliter (str): column spliter in one line.
            ID_spliter (str): ID spliter in one line.
        """
        self.col_spliter = col_spliter
        self.ID_spliter = ID_spliter
        self.batch_size = hparams.batch_size
        self.title_size = hparams.title_size
        self.body_size = hparams.body_size
        self.his_size = hparams.his_size
        self.npratio = npratio
        self.word_dict = self.load_dict(hparams.wordDict_file)
        # self.cat_dict = self.load_dict(hparams.catDict_file)
        # self.subcat_dict = self.load_dict(hparams.subcatDict_file)
        # self.entity_dict = self.load_dict(hparams.entityDict_file)
        # self.relation_dict = self.load_dict(hparams.relationDict_file)
        self.uid2index = self.load_dict(hparams.userDict_file)

    def load_dict(self, file_path):
        """load pickle file
        Args:
            file path (str): file path
        Returns:
            object: pickle loaded object
        """
        with open(file_path, "rb") as f:
            return pickle.load(f)

    def init_news(self, news_file):
        """init news information given news file, such as news_title_index and nid2index.
        Args:
            news_file: path of news file
        """
        self.nid2index = {}
        news_title = [""]
        # news_cat = [""]
        # news_subcat = [""]
        news_abstract = [""]
        # news_entitites = [""]
        # news_relation = [""]
        with tf.io.gfile.GFile(news_file, "r") as rd:
            for line in rd:
                nid, cat, title, ab, content = line.strip("\n").split(self.col_spliter)
                if nid in self.nid2index:
                    continue
                self.nid2index[nid] = len(self.nid2index) + 1
                title = word_tokenize(title)
                if ab is None:
                    abstract = ""
                else:
                    abstract = word_tokenize(ab)
                news_title.append(title)
                # news_cat.append(cat)
                # news_subcat.append(subcat)
                news_abstract.append(abstract)
                # news_entities.append(entity)
                # news_relation.append(relation)
        self.news_title_index = np.zeros(
            (len(news_title), self.title_size), dtype="int32"
        )
        self.news_abstract_index = np.zeros(
            (len(news_abstract), self.body_size), dtype="int32"
        )
        # self.news_cat_index = np.zeros((len(news_cat), 1), dtype="int32")
        # self.news_subcat_index = np.zeros((len(news_subcat, 1), dtype="int32")
        for news_index in range(len(news_title)):
            title = news_title[news_index]
            abstract = news_abstract[news_index]
            for word_index in range(min(self.title_size, len(title))):
                if title[word_index] in self.word_dict:
                    self.news_title_index[news_index, word_index] = self.word_dict[
                        title[word_index].lower()
                    ]
            for word_index in range(min(self.body_size, len(abstract))):
                if abstract[word_index] in self.word_dict:
                    self.news_abstract_index[news_index, word_index] = self.word_dict[
                        abstract[word_index].lower()
                    ]

    def init_behaviors(self, behaviors_file):
        """init behavior logs given behaviors file.
        Args:
        behaviors_file: path of behaviors file
        """
        self.histories = []
        self.imprs = []
        self.labels = []
        self.impr_indexes = []
        self.uindexes = []
        with tf.io.gfile.GFile(behaviors_file, "r") as rd:
            impr_index = 0
            for line in rd:
                uid, history, impr, _ = line.strip("\n").split(self.col_spliter)[-4:]
                history = [self.nid2index[i] for i in history.split()]
                history = [0] * (self.his_size - len(history)) + history[
                    : self.his_size
                ]
                impr_news = [self.nid2index[i.split("-")[0]] for i in impr.split()]
                label = [int(i.split("-")[1]) for i in impr.split()]
                uindex = self.uid2index[uid] if uid in self.uid2index else 0
                self.histories.append(history)
                self.imprs.append(impr_news)
                self.labels.append(label)
                self.impr_indexes.append(impr_index)
                self.uindexes.append(uindex)
                impr_index += 1

    def parser_one_line(self, line):
        """Parse one behavior sample into feature values.
        if npratio is larger than 0, return negtive sampled result.
        Args:
            line (int): sample index.
        Yields:
            list: Parsed results including label, impression id , user id,
            candidate_title_index, clicked_title_index.
        """
        if self.npratio > 0:
            impr_label = self.labels[line]
            impr = self.imprs[line]
            poss = []
            negs = []
            for news, click in zip(impr, impr_label):
                if click == 1:
                    poss.append(news)
                else:
                    negs.append(news)
            for p in poss:
                candidate_title_index = []
                candidate_abstract_index = []
                impr_index = []
                user_index = []
                label = [1] + [0] * self.npratio
                n = newsample(negs, self.npratio)
                candidate_title_index = self.news_title_index[[p] + n]
                candidate_abstract_index = self.news_abstract_index[[p] + n]
                click_title_index = self.news_title_index[self.histories[line]]
                click_abstract_index = self.news_abstract_index[self.histories[line]]
                impr_index.append(self.impr_indexes[line])
                user_index.append(self.uindexes[line])
                yield (
                    label,
                    impr_index,
                    user_index,
                    candidate_title_index,
                    candidate_abstract_index,
                    click_title_index,
                    click_abstract_index,
                )
        else:
            impr_label = self.labels[line]
            impr = self.imprs[line]
            for news, label in zip(impr, impr_label):
                candidate_title_index = []
                impr_index = []
                user_index = []
                label = [label]
                candidate_title_index.append(self.news_title_index[news])
                candidate_abstract_index = self.news_abstract_index[[p] + n]
                click_title_index = self.news_title_index[self.histories[line]]
                click_abstract_index = self.news_abstract_index[self.histories]
                impr_index.append(self.impr_indexes[line])
                user_index.append(self.uindexes[line])
                yield (
                    label,
                    impr_index,
                    user_index,
                    candidate_title_index,
                    candidate_abstract_index,
                    click_title_index,
                    click_abstract_index,
                )

    def load_data_from_file(self, news_file, behavior_file):
        """Read and parse data from news file and behavior file.
        Args:
            news_file (str): A file contains several informations of news.
            beahavior_file (str): A file contains information of user impressions.
        Yields:
            object: An iterator that yields parsed results, in the format of dict.
        """
        if not hasattr(self, "news_title_index"):
            self.init_news(news_file)
        if not hasattr(self, "impr_indexes"):
            self.init_behaviors(behavior_file)
        label_list = []
        imp_indexes = []
        user_indexes = []
        candidate_title_indexes = []
        candidate_abstract_indexes = []
        click_title_indexes = []
        click_abstract_indexes = []
        cnt = 0
        indexes = np.arange(len(self.labels))
        if self.npratio > 0:
            np.random.shuffle(indexes)
        for index in indexes:
            for (
                label,
                imp_index,
                user_index,
                candidate_title_index,
                candidate_abstract_index,
                click_title_index,
                click_abstract_index,
            ) in self.parser_one_line(index):
                candidate_title_indexes.append(candidate_title_index)
                candidate_abstract_indexes.append(candidate_abstract_index)
                click_title_indexes.append(click_title_index)
                click_abstract_indexes.append(click_abstract_index)
                imp_indexes.append(imp_index)
                user_indexes.append(user_index)
                label_list.append(label)
                cnt += 1
                if cnt >= self.batch_size:
                    yield self._convert_data(
                        label_list,
                        imp_indexes,
                        user_indexes,
                        candidate_title_indexes,
                        candidate_abstract_indexes,
                        click_title_indexes,
                        click_abstract_indexes,
                    )
                    label_list = []
                    imp_indexes = []
                    user_indexes = []
                    candidate_title_indexes = []
                    candidate_abstract_indexes = []
                    click_title_indexes = []
                    click_abstract_indexes = []
                    cnt = 0
        if cnt > 0:
            yield self._convert_data(
                label_list,
                imp_indexes,
                user_indexes,
                candidate_title_indexes,
                candidate_abstract_indexes,
                click_title_indexes,
                click_abstract_indexes,
            )

    def _convert_data(
        self,
        label_list,
        imp_indexes,
        user_indexes,
        candidate_title_indexes,
        candidate_abstract_indexes,
        click_title_indexes,
        click_abstract_indexes,
    ):
        """Convert data into numpy arrays that are good for further model operation.
        Args:
            label_list (list): a list of ground-truth labels.
            imp_indexes (list): a list of impression indexes.
            user_indexes (list): a list of user indexes.
            candidate_title_indexes (list): the candidate news titles' words indices.
            candidate_abstract_indexes (list): the candidate news abstract' words indices.
            click_title_indexes (list): words indices for user's clicked news titles.
            click_abstract_indexes (list): words indices for user's clicked news abstract.
        Returns:
            dict: A dictionary, containing multiple numpy arrays that are convenient for further operation.
        """
        labels = np.asarray(label_list, dtype=np.float32)
        imp_indexes = np.asarray(imp_indexes, dtype=np.int32)
        user_indexes = np.asarray(user_indexes, dtype=np.int32)
        candidate_title_index_batch = np.asarray(
            candidate_title_indexes, dtype=np.int64
        )
        candidate_abstract_index_batch = np.asarray(
            candidate_abstract_indexes, dtype=np.int64
        )
        click_title_index_batch = np.asarray(click_title_indexes, dtype=np.int64)
        click_abstract_index_batch = np.asarray(click_abstract_indexes, dtype=np.int64)
        return {
            "impression_index_batch": imp_indexes,
            "user_index_batch": user_indexes,
            "clicked_title_batch": click_title_index_batch,
            "clicked_abstract_batch": click_abstract_index_batch,
            "candidate_title_batch": candidate_title_index_batch,
            "candidate_abstract_batch": candidate_abstract_index_batch,
            "labels": labels,
        }

    def load_user_from_file(self, news_file, behavior_file):
        """Read and parse user data from news file and behavior file.
        Args:
            news_file (str): A file contains several informations of news.
            beahaviros_file (str): A file contains information of user impressions.
        Yields:
            object: An iterator that yields parsed user feature, in the format of dict.
        """
        if not hasattr(self, "news_title_index"):
            self.init_news(news_file)
        if not hasattr(self, "impr_indexes"):
            self.init_behaviors(behavior_file)
        user_indexes = []
        impr_indexes = []
        click_title_indexes = []
        click_abstract_indexes = []
        cnt = 0
        for index in range(len(self.impr_indexes)):
            click_title_indexes.append(self.news_title_index[self.histories[index]])
            click_abstract_indexes.append(
                self.news_abstract_index[self.histories[index]]
            )
            user_indexes.append(self.uindexes[index])
            impr_indexes.append(self.impr_indexes[index])
            cnt += 1
            if cnt >= self.batch_size:
                yield self._convert_user_data(
                    user_indexes,
                    impr_indexes,
                    click_title_indexes,
                    click_abstract_indexes,
                )
                user_indexes = []
                impr_indexes = []
                click_title_indexes = []
                click_abstract_indexes = []
                cnt = 0
        if cnt > 0:
            yield self._convert_user_data(
                user_indexes, impr_indexes, click_title_indexes, click_abstract_indexes
            )

    def _convert_user_data(
        self, user_indexes, impr_indexes, click_title_indexes, click_abstract_indexes
    ):
        """Convert data into numpy arrays that are good for further model operation.
        Args:
            user_indexes (list): a list of user indexes.
            click_title_indexes (list): words indices for user's clicked news titles.
        Returns:
            dict: A dictionary, containing multiple numpy arrays that are convenient for further operation.
        """
        user_indexes = np.asarray(user_indexes, dtype=np.int32)
        impr_indexes = np.asarray(impr_indexes, dtype=np.int32)
        click_title_index_batch = np.asarray(click_title_indexes, dtype=np.int64)
        click_abstract_index_batch = np.asarray(click_abstract_indexes, dtype=np.int64)
        return {
            "user_index_batch": user_indexes,
            "impr_index_batch": impr_indexes,
            "clicked_title_batch": click_title_index_batch,
            "clicked_abstract_batch": click_abstract_index_batch,
        }

    def load_news_from_file(self, news_file):
        """Read and parse user data from news file.
        Args:
            news_file (str): A file contains several informations of news.
        Yields:
            object: An iterator that yields parsed news feature, in the format of dict.
        """
        if not hasattr(self, "news_title_index"):
            self.init_news(news_file)
        news_indexes = []
        candidate_title_indexes = []
        candidate_abstract_indexes = []
        cnt = 0
        for index in range(len(self.news_title_index)):
            news_indexes.append(index)
            candidate_title_indexes.append(self.news_title_index[index])
            candidate_abstract_indexes.append(self.news_abstract_index[index])
            cnt += 1
            if cnt >= self.batch_size:
                yield self._convert_news_data(
                    news_indexes, candidate_title_indexes, candidate_abstract_indexes
                )
                news_indexes = []
                candidate_title_indexes = []
                candidate_abstract_indexes = []
                cnt = 0
        if cnt > 0:
            yield self._convert_news_data(
                news_indexes, candidate_title_indexes, candidate_abstract_indexes
            )

    def _convert_news_data(
        self, news_indexes, candidate_title_indexes, candidate_abstract_indexes
    ):
        """Convert data into numpy arrays that are good for further model operation.
        Args:
            news_indexes (list): a list of news indexes.
            candidate_title_indexes (list): the candidate news titles' words indices.
        Returns:
            dict: A dictionary, containing multiple numpy arrays that are convenient for further operation.
        """
        news_indexes_batch = np.asarray(news_indexes, dtype=np.int32)
        candidate_title_index_batch = np.asarray(
            candidate_title_indexes, dtype=np.int32
        )
        candidate_abstract_index_batch = np.asarray(
            candidate_abstract_indexes, dtype=np.int32
        )
        return {
            "news_index_batch": news_indexes_batch,
            "candidate_title_batch": candidate_title_index_batch,
            "candidate_abstract_batch": candidate_abstract_index_batch,
        }

    def load_impression_from_file(self, behaivors_file):
        """Read and parse impression data from behaivors file.
        Args:
            behaivors_file (str): A file contains several informations of behaviros.
        Yields:
            object: An iterator that yields parsed impression data, in the format of dict.
        """
        if not hasattr(self, "histories"):
            self.init_behaviors(behaivors_file)
        indexes = np.arange(len(self.labels))
        for index in indexes:
            impr_label = np.array(self.labels[index], dtype="int32")
            impr_news = np.array(self.imprs[index], dtype="int32")
            yield (
                self.impr_indexes[index],
                impr_news,
                self.uindexes[index],
                impr_label,
            )


# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
class AttLayer2(layers.Layer):
    """Soft alignment attention implement.
    Attributes:
        dim (int): attention hidden dim
    """

    def __init__(self, dim=200, seed=0, **kwargs):
        """Initialization steps for AttLayer2.
        Args:
            dim (int): attention hidden dim
        """
        self.dim = dim
        self.seed = seed
        super(AttLayer2, self).__init__(**kwargs)

    def build(self, input_shape):
        """Initialization for variables in AttLayer2
        There are there variables in AttLayer2, i.e. W, b and q.
        Args:
            input_shape (object): shape of input tensor.
        """
        assert len(input_shape) == 3
        dim = self.dim
        self.W = self.add_weight(
            name="W",
            shape=(int(input_shape[-1]), dim),
            initializer=keras.initializers.glorot_uniform(seed=self.seed),
            trainable=True,
        )
        self.b = self.add_weight(
            name="b",
            shape=(dim,),
            initializer=keras.initializers.Zeros(),
            trainable=True,
        )
        self.q = self.add_weight(
            name="q",
            shape=(dim, 1),
            initializer=keras.initializers.glorot_uniform(seed=self.seed),
            trainable=True,
        )
        super(AttLayer2, self).build(input_shape)  # be sure you call this somewhere!

    def call(self, inputs, mask=None, **kwargs):
        """Core implemention of soft attention
        Args:
            inputs (object): input tensor.
        Returns:
            object: weighted sum of input tensors.
        """
        attention = K.tanh(K.dot(inputs, self.W) + self.b)
        attention = K.dot(attention, self.q)
        attention = K.squeeze(attention, axis=2)
        if mask == None:
            attention = K.exp(attention)
        else:
            attention = K.exp(attention) * K.cast(mask, dtype="float32")
        attention_weight = attention / (
            K.sum(attention, axis=-1, keepdims=True) + K.epsilon()
        )
        attention_weight = K.expand_dims(attention_weight)
        weighted_input = inputs * attention_weight
        return K.sum(weighted_input, axis=1)

    def compute_mask(self, input, input_mask=None):
        """Compte output mask value
        Args:
            input (object): input tensor.
            input_mask: input mask
        Returns:
            object: output mask.
        """
        return None

    def compute_output_shape(self, input_shape):
        """Compute shape of output tensor
        Args:
            input_shape (tuple): shape of input tensor.
        Returns:
            tuple: shape of output tensor.
        """
        return input_shape[0], input_shape[-1]


class SelfAttention(layers.Layer):
    """Multi-head self attention implement.
    Args:
        multiheads (int): The number of heads.
        head_dim (object): Dimention of each head.
        mask_right (boolean): whether to mask right words.
    Returns:
        object: Weighted sum after attention.
    """

    def __init__(self, multiheads, head_dim, seed=0, mask_right=False, **kwargs):
        """Initialization steps for AttLayer2.
        Args:
            multiheads (int): The number of heads.
            head_dim (object): Dimention of each head.
            mask_right (boolean): whether to mask right words.
        """
        self.multiheads = multiheads
        self.head_dim = head_dim
        self.output_dim = multiheads * head_dim
        self.mask_right = mask_right
        self.seed = seed
        super(SelfAttention, self).__init__(**kwargs)

    def compute_output_shape(self, input_shape):
        """Compute shape of output tensor.
        Returns:
            tuple: output shape tuple.
        """
        return (input_shape[0][0], input_shape[0][1], self.output_dim)

    def build(self, input_shape):
        """Initialization for variables in SelfAttention.
        There are three variables in SelfAttention, i.e. WQ, WK ans WV.
        WQ is used for linear transformation of query.
        WK is used for linear transformation of key.
        WV is used for linear transformation of value.
        Args:
            input_shape (object): shape of input tensor.
        """
        self.WQ = self.add_weight(
            name="WQ",
            shape=(int(input_shape[0][-1]), self.output_dim),
            initializer=keras.initializers.glorot_uniform(seed=self.seed),
            trainable=True,
        )
        self.WK = self.add_weight(
            name="WK",
            shape=(int(input_shape[1][-1]), self.output_dim),
            initializer=keras.initializers.glorot_uniform(seed=self.seed),
            trainable=True,
        )
        self.WV = self.add_weight(
            name="WV",
            shape=(int(input_shape[2][-1]), self.output_dim),
            initializer=keras.initializers.glorot_uniform(seed=self.seed),
            trainable=True,
        )
        super(SelfAttention, self).build(input_shape)

    def Mask(self, inputs, seq_len, mode="add"):
        """Mask operation used in multi-head self attention
        Args:
            seq_len (object): sequence length of inputs.
            mode (str): mode of mask.
        Returns:
            object: tensors after masking.
        """
        if seq_len == None:
            return inputs
        else:
            mask = K.one_hot(indices=seq_len[:, 0], num_classes=K.shape(inputs)[1])
            mask = 1 - K.cumsum(mask, axis=1)
            for _ in range(len(inputs.shape) - 2):
                mask = K.expand_dims(mask, 2)
            if mode == "mul":
                return inputs * mask
            elif mode == "add":
                return inputs - (1 - mask) * 1e12

    def call(self, QKVs):
        """Core logic of multi-head self attention.
        Args:
            QKVs (list): inputs of multi-head self attention i.e. qeury, key and value.
        Returns:
            object: ouput tensors.
        """
        if len(QKVs) == 3:
            Q_seq, K_seq, V_seq = QKVs
            Q_len, V_len = None, None
        elif len(QKVs) == 5:
            Q_seq, K_seq, V_seq, Q_len, V_len = QKVs
        Q_seq = K.dot(Q_seq, self.WQ)
        Q_seq = K.reshape(
            Q_seq, shape=(-1, K.shape(Q_seq)[1], self.multiheads, self.head_dim)
        )
        Q_seq = K.permute_dimensions(Q_seq, pattern=(0, 2, 1, 3))
        K_seq = K.dot(K_seq, self.WK)
        K_seq = K.reshape(
            K_seq, shape=(-1, K.shape(K_seq)[1], self.multiheads, self.head_dim)
        )
        K_seq = K.permute_dimensions(K_seq, pattern=(0, 2, 1, 3))
        V_seq = K.dot(V_seq, self.WV)
        V_seq = K.reshape(
            V_seq, shape=(-1, K.shape(V_seq)[1], self.multiheads, self.head_dim)
        )
        V_seq = K.permute_dimensions(V_seq, pattern=(0, 2, 1, 3))
        # tf.einsum('m b i k, m b j k -> m b i j', Q_seq , K_seq) # shape [10, 20, 50]
        A = tf.einsum("m b i k, m b j k -> m b i j", Q_seq, K_seq) / K.sqrt(
            K.cast(self.head_dim, dtype="float32")
        )
        A = K.permute_dimensions(
            A, pattern=(0, 3, 2, 1)
        )  # A.shape=[batch_size,K_sequence_length,Q_sequence_length,self.multiheads]
        A = self.Mask(A, V_len, "add")
        A = K.permute_dimensions(A, pattern=(0, 3, 2, 1))
        if self.mask_right:
            ones = K.ones_like(A[:1, :1])
            lower_triangular = K.tf.matrix_band_part(ones, num_lower=-1, num_upper=0)
            mask = (ones - lower_triangular) * 1e12
            A = A - mask
        A = K.softmax(A)
        # tf.einsum('m b i k, m b k j -> m b i j', A , V_seq)
        O_seq = tf.einsum("m b i k, m b k j -> m b i j", A, V_seq)
        O_seq = K.permute_dimensions(O_seq, pattern=(0, 2, 1, 3))
        O_seq = K.reshape(O_seq, shape=(-1, K.shape(O_seq)[1], self.output_dim))
        O_seq = self.Mask(O_seq, Q_len, "mul")
        return O_seq

    def get_config(self):
        """add multiheads, multiheads and mask_right into layer config.
        Returns:
            dict: config of SelfAttention layer.
        """
        config = super(SelfAttention, self).get_config()
        config.update(
            {
                "multiheads": self.multiheads,
                "head_dim": self.head_dim,
                "mask_right": self.mask_right,
            }
        )
        return config


__all__ = ["BaseModel"]


class BaseModel:
    """Basic class of models
    Attributes:
        hparams (object): A tf.contrib.training.HParams object, hold the entire set of hyperparameters.
        train_iterator (object): An iterator to load the data in training steps.
        test_iterator (object): An iterator to load the data in testing steps.
        graph (object): An optional graph.
        seed (int): Random seed.
    """

    def __init__(
        self,
        hparams,
        iterator_creator,
        seed=None,
    ):
        """Initializing the model. Create common logics which are needed by all deeprec models, such as loss function,
        parameter set.
        Args:
            hparams (object): A tf.contrib.training.HParams object, hold the entire set of hyperparameters.
            iterator_creator (object): An iterator to load the data.
            graph (object): An optional graph.
            seed (int): Random seed.
        """
        self.seed = seed
        tf.compat.v1.set_random_seed(seed)
        np.random.seed(seed)
        self.train_iterator = iterator_creator(
            hparams,
            hparams.npratio,
            col_spliter="\t",
        )
        self.test_iterator = iterator_creator(
            hparams,
            col_spliter="\t",
        )
        self.hparams = hparams
        self.support_quick_scoring = hparams.support_quick_scoring
        self.model, self.scorer = self._build_graph()
        self.loss = self._get_loss()
        self.train_optimizer = self._get_opt()
        self.model.compile(loss=self.loss, optimizer=self.train_optimizer)

    def _init_embedding(self, file_path):
        """Load pre-trained embeddings as a constant tensor.
        Args:
            file_path (str): the pre-trained glove embeddings file path.
        Returns:
            numpy.ndarray: A constant numpy array.
        """
        return np.load(file_path)

    @abc.abstractmethod
    def _build_graph(self):
        """Subclass will implement this."""
        pass

    @abc.abstractmethod
    def _get_input_label_from_iter(self, batch_data):
        """Subclass will implement this"""
        pass

    def _get_loss(self):
        """Make loss function, consists of data loss and regularization loss
        Returns:
            object: Loss function or loss function name
        """
        if self.hparams.loss == "cross_entropy_loss":
            data_loss = "categorical_crossentropy"
        elif self.hparams.loss == "log_loss":
            data_loss = "binary_crossentropy"
        else:
            raise ValueError("this loss not defined {0}".format(self.hparams.loss))
        return data_loss

    def _get_opt(self):
        """Get the optimizer according to configuration. Usually we will use Adam.
        Returns:
            object: An optimizer.
        """
        lr = self.hparams.learning_rate
        optimizer = self.hparams.optimizer
        if optimizer == "adam":
            train_opt = tf.keras.optimizers.Adam(lr=lr)
        return train_opt

    def _get_pred(self, logit, task):
        """Make final output as prediction score, according to different tasks.
        Args:
            logit (object): Base prediction value.
            task (str): A task (values: regression/classification)
        Returns:
            object: Transformed score
        """
        if task == "regression":
            pred = tf.identity(logit)
        elif task == "classification":
            pred = tf.sigmoid(logit)
        else:
            raise ValueError(
                "method must be regression or classification, but now is {0}".format(
                    task
                )
            )
        return pred

    def train(self, train_batch_data):
        """Go through the optimization step once with training data in feed_dict.
        Args:
            sess (object): The model session object.
            feed_dict (dict): Feed values to train the model. This is a dictionary that maps graph elements to values.
        Returns:
            list: A list of values, including update operation, total loss, data loss, and merged summary.
        """
        train_input, train_label = self._get_input_label_from_iter(train_batch_data)
        rslt = self.model.train_on_batch(train_input, train_label)
        return rslt

    def eval(self, eval_batch_data):
        """Evaluate the data in feed_dict with current model.
        Args:
            sess (object): The model session object.
            feed_dict (dict): Feed values for evaluation. This is a dictionary that maps graph elements to values.
        Returns:
            list: A list of evaluated results, including total loss value, data loss value, predicted scores, and ground-truth labels.
        """
        eval_input, eval_label = self._get_input_label_from_iter(eval_batch_data)
        imp_index = eval_batch_data["impression_index_batch"]
        pred_rslt = self.scorer.predict_on_batch(eval_input)
        return pred_rslt, eval_label, imp_index

    def fit(
        self,
        train_news_file,
        train_behaviors_file,
        valid_news_file,
        valid_behaviors_file,
        test_news_file=None,
        test_behaviors_file=None,
    ):
        """Fit the model with train_file. Evaluate the model on valid_file per epoch to observe the training status.
        If test_news_file is not None, evaluate it too.
        Args:
            train_file (str): training data set.
            valid_file (str): validation set.
            test_news_file (str): test set.
        Returns:
            object: An instance of self.
        """
        self.history = {}
        self.history["train_loss"] = []
        self.history["val_loss"] = []
        self.history["Val_auc"] = []
        for epoch in range(1, self.hparams.epochs + 1):
            step = 0
            # self.hparams.current_epoch = epoch
            epoch_loss = 0
            train_start = time.time()
            tqdm_util = tqdm(
                self.train_iterator.load_data_from_file(
                    train_news_file, train_behaviors_file
                )
            )
            for batch_data_input in tqdm_util:
                step_result = self.train(batch_data_input)
                step_data_loss = step_result
                epoch_loss += step_data_loss
                step += 1
                if step % self.hparams.show_step == 0:
                    tqdm_util.set_description(
                        "step {0:d} , total_loss: {1:.4f}, data_loss: {2:.4f}".format(
                            step, epoch_loss / step, step_data_loss
                        )
                    )
            self.history["train_loss"].append(epoch_loss / step)
            self.history["val_loss"].append(step_data_loss)
            train_end = time.time()
            train_time = train_end - train_start
            eval_start = time.time()
            train_info = ",".join(
                [
                    str(item[0]) + ":" + str(item[1])
                    for item in [("logloss loss", epoch_loss / step)]
                ]
            )
            eval_res = self.run_eval(valid_news_file, valid_behaviors_file)
            eval_info = ", ".join(
                [
                    str(item[0]) + ":" + str(item[1])
                    for item in sorted(eval_res.items(), key=lambda x: x[0])
                ]
            )
            if test_news_file is not None:
                test_res = self.run_eval(test_news_file, test_behaviors_file)
                test_info = ", ".join(
                    [
                        str(item[0]) + ":" + str(item[1])
                        for item in sorted(test_res.items(), key=lambda x: x[0])
                    ]
                )
            eval_end = time.time()
            eval_time = eval_end - eval_start
            if test_news_file is not None:
                print(
                    "at epoch {0:d}".format(epoch)
                    + "\ntrain info: "
                    + train_info
                    + "\neval info: "
                    + eval_info
                    + "\ntest info: "
                    + test_info
                )
            else:
                print(
                    "at epoch {0:d}".format(epoch)
                    + "\ntrain info: "
                    + train_info
                    + "\neval info: "
                    + eval_info
                )
            print(
                "at epoch {0:d} , train time: {1:.1f} eval time: {2:.1f}".format(
                    epoch, train_time, eval_time
                )
            )
        return self

    def group_labels(self, labels, preds, group_keys):
        """Devide labels and preds into several group according to values in group keys.
        Args:
            labels (list): ground truth label list.
            preds (list): prediction score list.
            group_keys (list): group key list.
        Returns:
            list, list, list:
            - Keys after group.
            - Labels after group.
            - Preds after group.
        """
        all_keys = list(set(group_keys))
        all_keys.sort()
        group_labels = {k: [] for k in all_keys}
        group_preds = {k: [] for k in all_keys}
        for l, p, k in zip(labels, preds, group_keys):
            group_labels[k].append(l)
            group_preds[k].append(p)
        all_labels = []
        all_preds = []
        for k in all_keys:
            all_labels.append(group_labels[k])
            all_preds.append(group_preds[k])
        return all_keys, all_labels, all_preds

    def run_eval(self, news_filename, behaviors_file):
        """Evaluate the given file and returns some evaluation metrics.
        Args:
            filename (str): A file name that will be evaluated.
        Returns:
            dict: A dictionary that contains evaluation metrics.
        """
        if self.support_quick_scoring:
            _, group_labels, group_preds = self.run_fast_eval(
                news_filename, behaviors_file
            )
        else:
            _, group_labels, group_preds = self.run_slow_eval(
                news_filename, behaviors_file
            )
        res = cal_metric(group_labels, group_preds, self.hparams.metrics)
        return res

    def user(self, batch_user_input):
        user_input = self._get_user_feature_from_iter(batch_user_input)
        user_vec = self.userencoder.predict_on_batch(user_input)
        user_index = batch_user_input["impr_index_batch"]
        return user_index, user_vec

    def news(self, batch_news_input):
        news_input = self._get_news_feature_from_iter(batch_news_input)
        news_vec = self.newsencoder.predict_on_batch(news_input)
        news_index = batch_news_input["news_index_batch"]
        return news_index, news_vec

    def run_user(self, news_filename, behaviors_file):
        if not hasattr(self, "userencoder"):
            raise ValueError("model must have attribute userencoder")
        user_indexes = []
        user_vecs = []
        for batch_data_input in tqdm(
            self.test_iterator.load_user_from_file(news_filename, behaviors_file)
        ):
            # print("\nbatch_data_input:", batch_data_input)
            user_index, user_vec = self.user(batch_data_input)
            user_indexes.extend(np.reshape(user_index, -1))
            user_vecs.extend(user_vec)
        return dict(zip(user_indexes, user_vecs))

    def run_news(self, news_filename):
        if not hasattr(self, "newsencoder"):
            raise ValueError("model must have attribute newsencoder")
        news_indexes = []
        news_vecs = []
        for batch_data_input in tqdm(
            self.test_iterator.load_news_from_file(news_filename)
        ):
            news_index, news_vec = self.news(batch_data_input)
            news_indexes.extend(np.reshape(news_index, -1))
            news_vecs.extend(news_vec)
        return dict(zip(news_indexes, news_vecs))

    def run_slow_eval(self, news_filename, behaviors_file):
        preds = []
        labels = []
        imp_indexes = []
        for batch_data_input in tqdm(
            self.test_iterator.load_data_from_file(news_filename, behaviors_file)
        ):
            step_pred, step_labels, step_imp_index = self.eval(batch_data_input)
            preds.extend(np.reshape(step_pred, -1))
            labels.extend(np.reshape(step_labels, -1))
            imp_indexes.extend(np.reshape(step_imp_index, -1))
        group_impr_indexes, group_labels, group_preds = self.group_labels(
            labels, preds, imp_indexes
        )
        return group_impr_indexes, group_labels, group_preds

    def run_fast_eval(self, news_filename, behaviors_file):
        news_vecs = self.run_news(news_filename)
        user_vecs = self.run_user(news_filename, behaviors_file)
        self.news_vecs = news_vecs
        self.user_vecs = user_vecs
        group_impr_indexes = []
        group_labels = []
        group_preds = []
        for (
            impr_index,
            news_index,
            user_index,
            label,
        ) in tqdm(self.test_iterator.load_impression_from_file(behaviors_file)):
            pred = np.dot(
                np.stack([news_vecs[i] for i in news_index], axis=0),
                user_vecs[impr_index],
            )
            group_impr_indexes.append(impr_index)
            group_labels.append(label)
            group_preds.append(pred)
        return group_impr_indexes, group_labels, group_preds


# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
__all__ = ["NRMSModel"]


class NRMSModel(BaseModel):
    """NRMS model(Neural News Recommendation with Multi-Head Self-Attention)
    Chuhan Wu, Fangzhao Wu, Suyu Ge, Tao Qi, Yongfeng Huang,and Xing Xie, "Neural News
    Recommendation with Multi-Head Self-Attention" in Proceedings of the 2019 Conference
    on Empirical Methods in Natural Language Processing and the 9th International Joint Conference
    on Natural Language Processing (EMNLP-IJCNLP)
    Attributes:
        word2vec_embedding (numpy.ndarray): Pretrained word embedding matrix.
        hparam (object): Global hyper-parameters.
    """

    def __init__(
        self,
        hparams,
        iterator_creator,
        seed=None,
    ):
        """Initialization steps for NRMS.
        Compared with the BaseModel, NRMS need word embedding.
        After creating word embedding matrix, BaseModel's __init__ method will be called.
        Args:
            hparams (object): Global hyper-parameters. Some key setttings such as head_num and head_dim are there.
            iterator_creator_train (object): NRMS data loader class for train data.
            iterator_creator_test (object): NRMS data loader class for test and validation data
        """
        self.word2vec_embedding = self._init_embedding(hparams.wordEmb_file)
        super().__init__(
            hparams,
            iterator_creator,
            seed=seed,
        )

    def _get_input_label_from_iter(self, batch_data):
        """get input and labels for trainning from iterator
        Args:
            batch data: input batch data from iterator
        Returns:
            list: input feature fed into model (clicked_title_batch & candidate_title_batch)
            numpy.ndarray: labels
        """
        input_feat = [
            batch_data["clicked_title_batch"],
            batch_data["clicked_abstract_batch"],
            batch_data["candidate_title_batch"],
            batch_data["candidate_abstract_batch"],
        ]
        input_label = batch_data["labels"]
        return input_feat, input_label

    def _get_user_feature_from_iter(self, batch_data):
        """get input of user encoder
        Args:
            batch_data: input batch data from user iterator
        Returns:
            numpy.ndarray: input user feature (clicked title batch)
        """
        # print("\nclicked_title_batch:", batch_data["clicked_title_batch"].shape)
        # print("\nclicked_abstract_batch:", batch_data["clicked_abstract_batch"].shape)
        input_feature = [
            batch_data["clicked_title_batch"],
            batch_data["clicked_abstract_batch"],
        ]
        input_feature = np.concatenate(input_feature, axis=-1)
        # print("\ninput_feature:", input_feature.shape)
        return input_feature

    def _get_news_feature_from_iter(self, batch_data):
        """get input of news encoder
        Args:
            batch_data: input batch data from news iterator
        Returns:
            numpy.ndarray: input news feature (candidate title batch)
        """
        input_feature = [
            batch_data["candidate_title_batch"],
            batch_data["candidate_abstract_batch"],
        ]
        input_feature = np.concatenate(input_feature, axis=-1)
        return input_feature

    def _build_graph(self):
        """Build NRMS model and scorer.
        Returns:
            object: a model used to train.
            object: a model used to evaluate and inference.
        """
        hparams = self.hparams
        model, scorer = self._build_nrms()
        return model, scorer

    def _build_userencoder(self, newsencoder):
        """The main function to create user encoder of NRMS.
        Args:
            newsencoder (object): the news encoder of NRMS.
        Return:
            object: the user encoder of NRMS.
        """
        hparams = self.hparams
        his_input_title_abstract = keras.Input(
            shape=(hparams.his_size, hparams.title_size + hparams.body_size),
            dtype="int32",
        )
        click_new_presents = layers.TimeDistributed(newsencoder)(
            his_input_title_abstract
        )
        y = SelfAttention(hparams.head_num, hparams.head_dim, seed=self.seed)(
            [click_new_presents] * 3
        )
        user_present = AttLayer2(hparams.attention_hidden_dim, seed=self.seed)(y)
        model = keras.Model(his_input_title_abstract, user_present, name="user_encoder")
        print("\n", model.summary())
        return model

    def _build_newsencoder(self, embedding_layer):
        """The main function to create news encoder of NRMS.
        Args:
            embedding_layer (object): a word embedding layer.
        Return:
            object: the news encoder of NRMS
        """
        hparams = self.hparams
        input_title_abstract = tf.keras.Input(
            shape=(hparams.title_size + hparams.body_size,), dtype="int32"
        )
        # print("\ninput_title_abstract:", input_title_abstract)
        sequences_input_title = layers.Lambda(lambda x: x[:, : hparams.title_size])(
            input_title_abstract
        )
        # print("\nsequences_input_title:", sequences_input_title)
        sequences_input_abstract = layers.Lambda(
            lambda x: x[:, hparams.title_size : hparams.title_size + hparams.body_size]
        )(input_title_abstract)
        title_rep = self._build_titleencoder(embedding_layer)(sequences_input_title)
        abstract_rep = self._build_abstractencoder(embedding_layer)(
            sequences_input_abstract
        )
        # print("title_rep:", title_rep)
        # print("abstract_rep:", abstract_rep)
        concate_rep = layers.Concatenate(axis=-2)([title_rep, abstract_rep])
        y = SelfAttention(hparams.head_num, hparams.head_dim, seed=self.seed)(
            [concate_rep, concate_rep, concate_rep]
        )
        pred_title = AttLayer2(hparams.attention_hidden_dim, seed=self.seed)(y)
        model = keras.Model(input_title_abstract, pred_title, name="news_encoder")
        print("\n", model.summary())
        return model

    def _build_titleencoder(self, embedding_layer):
        hparams = self.hparams
        sequences_input_title = keras.Input(shape=(hparams.title_size,), dtype="int32")
        embedded_sequences_title = embedding_layer(sequences_input_title)
        y = layers.Dropout(hparams.dropout)(embedded_sequences_title)
        y = SelfAttention(hparams.head_num, hparams.head_dim, seed=self.seed)([y, y, y])
        y = layers.Dropout(hparams.dropout)(y)
        pred_title = AttLayer2(hparams.attention_hidden_dim, seed=self.seed)(y)
        pred_title = layers.Reshape((1, 400))(pred_title)
        model = keras.Model(sequences_input_title, pred_title, name="title_encoder")
        print("\n", model.summary())
        return model

    def _build_abstractencoder(self, embedding_layer):
        hparams = self.hparams
        sequences_input_abstract = keras.Input(
            shape=(hparams.body_size,), dtype="int32"
        )
        embedded_sequences_abstract = embedding_layer(sequences_input_abstract)
        y = layers.Dropout(hparams.dropout)(embedded_sequences_abstract)
        y = SelfAttention(hparams.head_num, hparams.head_dim, seed=self.seed)([y, y, y])
        y = layers.Dropout(hparams.dropout)(y)
        pred_abstract = AttLayer2(hparams.attention_hidden_dim, seed=self.seed)(y)
        pred_abstract = layers.Reshape((1, 400))(pred_abstract)
        model = keras.Model(
            sequences_input_abstract, pred_abstract, name="abstract_encoder"
        )
        print("\n", model.summary())
        return model

    def _build_nrms(self):
        """The main function to create NRMS's logic. The core of NRMS
        is a user encoder and a news encoder.
        Returns:
            object: a model used to train.
            object: a model used to evaluate and inference.
        """
        hparams = self.hparams
        his_input_title = keras.Input(
            shape=(hparams.his_size, hparams.title_size), dtype="int32"
        )
        his_input_abstract = keras.Input(
            shape=(hparams.his_size, hparams.body_size), dtype="int32"
        )
        pred_input_title = keras.Input(
            shape=(hparams.npratio + 1, hparams.title_size), dtype="int32"
        )
        pred_input_abstract = keras.Input(
            shape=(hparams.npratio + 1, hparams.body_size), dtype="int32"
        )
        pred_input_title_one = keras.Input(
            shape=(
                1,
                hparams.title_size,
            ),
            dtype="int32",
        )
        pred_input_abstract_one = keras.Input(
            shape=(
                1,
                hparams.body_size,
            ),
            dtype="int32",
        )
        his_title_abstract = layers.Concatenate(axis=-1)(
            [his_input_title, his_input_abstract]
        )
        pred_title_abstract = layers.Concatenate(axis=-1)(
            [pred_input_title, pred_input_abstract]
        )
        pred_title_abstract_one = layers.Concatenate(axis=-1)(
            [
                pred_input_title_one,
                pred_input_abstract_one,
            ]
        )
        pred_title_abstract_one = layers.Reshape((-1,))(pred_title_abstract_one)
        embedding_layer = layers.Embedding(
            self.word2vec_embedding.shape[0],
            hparams.word_emb_dim,
            weights=[self.word2vec_embedding],
            trainable=True,
        )
        self.newsencoder = self._build_newsencoder(embedding_layer)
        self.userencoder = self._build_userencoder(self.newsencoder)
        user_present = self.userencoder(his_title_abstract)
        # print("\npred_title_abstract:", pred_title_abstract)
        news_present = layers.TimeDistributed(self.newsencoder)(pred_title_abstract)
        news_present_one = self.newsencoder(pred_title_abstract_one)
        preds = layers.Dot(axes=-1)([news_present, user_present])
        preds = layers.Activation(activation="softmax")(preds)
        pred_one = layers.Dot(axes=-1)([news_present_one, user_present])
        pred_one = layers.Activation(activation="sigmoid")(pred_one)
        model = keras.Model(
            [
                his_input_title,
                his_input_abstract,
                pred_input_title,
                pred_input_abstract,
            ],
            preds,
        )
        scorer = keras.Model(
            [
                his_input_title,
                his_input_abstract,
                pred_input_title_one,
                pred_input_abstract_one,
            ],
            pred_one,
        )
        print("\n", model.summary())
        return model, scorer


data_path = "/kaggle/input/mindlike-final/"
train_news_file = os.path.join(data_path, r"news_train.tsv")
train_behaviors_file = os.path.join(data_path, r"behaviors_train.tsv")
valid_news_file = os.path.join(data_path, r"news_val.tsv")
valid_behaviors_file = os.path.join(data_path, r"behaviors_val.tsv")
# test_news_file = os.path.join(data_path, r'news_test.tsv')
# est_behaviors_file = os.path.join(data_path, r'behaviors_test.tsv')
iterator = MINDIterator
with tpu_strategy.scope():
    model = NRMSModel(hparams, iterator)
# print(model.run_eval(train_news_file, train_behaviors_file))
# print("\n", model.run_eval(valid_news_file, valid_behaviors_file))
# %%time
model.fit(train_news_file, train_behaviors_file, valid_news_file, valid_behaviors_file)
model.model.save_weights("/kaggle/working/models/nrms_ckpt")


# adsasd
# model.model.save("/kaggle/working/model.ckpt")
#!cd /kaggle/working
#!zip -r file1.zip /kaggle/working/model.ckpt
def init_behaviors2(iterator, behaviors_file):
    """init behavior logs given behaviors file.
    Args:
    behaviors_file: path of behaviors file
    """
    iterator.histories = []
    iterator.imprs = []
    iterator.labels = []
    iterator.impr_indexes = []
    iterator.uindexes = []
    with tf.io.gfile.GFile(behaviors_file, "r") as rd:
        impr_index = 0
        for line in rd:
            uid, time, history, impr = line.strip("\n").split(iterator.col_spliter)[-4:]
            history = [iterator.nid2index[i] for i in history.split()]
            history = [0] * (iterator.his_size - len(history)) + history[
                : iterator.his_size
            ]
            impr_news = [iterator.nid2index[i.split("-")[0]] for i in impr.split()]
            label = [0 for i in impr.split()]
            uindex = iterator.uid2index[uid] if uid in iterator.uid2index else 0
            iterator.histories.append(history)
            iterator.imprs.append(impr_news)
            iterator.labels.append(label)
            iterator.impr_indexes.append(impr_index)
            iterator.uindexes.append(uindex)
            impr_index += 1


def load_impression_from_file2(iterator, behaivors_file):
    """Read and parse impression data from behaivors file.
    Args:
        behaivors_file (str): A file contains several informations of behaviros.
    Yields:
        object: An iterator that yields parsed impression data, in the format of dict.
    """
    init_behaviors2(iterator, behaivors_file)
    indexes = np.arange(len(iterator.labels))
    for index in indexes:
        impr_label = np.array(iterator.labels[index], dtype="int32")
        impr_news = np.array(iterator.imprs[index], dtype="int32")
        yield (
            iterator.impr_indexes[index],
            impr_news,
            iterator.uindexes[index],
            impr_label,
        )


def load_user_from_file2(iterator, news_file, behavior_file):
    """Read and parse user data from news file and behavior file.
    Args:
        news_file (str): A file contains several informations of news.
        beahaviros_file (str): A file contains information of user impressions.
    Yields:
        object: An iterator that yields parsed user feature, in the format of dict.
    """
    iterator.init_news(news_file)
    init_behaviors2(iterator, behavior_file)
    user_indexes = []
    impr_indexes = []
    click_title_indexes = []
    click_abstract_indexes = []
    cnt = 0
    for index in range(len(iterator.impr_indexes)):
        click_title_indexes.append(iterator.news_title_index[iterator.histories[index]])
        click_abstract_indexes.append(
            iterator.news_abstract_index[iterator.histories[index]]
        )
        user_indexes.append(iterator.uindexes[index])
        impr_indexes.append(iterator.impr_indexes[index])
        cnt += 1
        if cnt >= iterator.batch_size:
            yield iterator._convert_user_data(
                user_indexes, impr_indexes, click_title_indexes, click_abstract_indexes
            )
            user_indexes = []
            impr_indexes = []
            click_title_indexes = []
            click_abstract_indexes = []
            cnt = 0
    if cnt > 0:
        yield iterator._convert_user_data(
            user_indexes, impr_indexes, click_title_indexes, click_abstract_indexes
        )


def load_news_from_file2(iterator, news_file):
    print(news_file)
    """Read and parse user data from news file.
        Args:
            news_file (str): A file contains several informations of news.
        Yields:
            object: An iterator that yields parsed news feature, in the format of dict.
        """
    # if not hasattr(iterator, "news_title_index"):
    iterator.init_news(news_file)
    news_indexes = []
    candidate_title_indexes = []
    candidate_abstract_indexes = []
    cnt = 0
    for index in range(len(iterator.news_title_index)):
        news_indexes.append(index)
        candidate_title_indexes.append(iterator.news_title_index[index])
        candidate_abstract_indexes.append(iterator.news_abstract_index[index])
        cnt += 1
        if cnt >= iterator.batch_size:
            yield iterator._convert_news_data(
                news_indexes, candidate_title_indexes, candidate_abstract_indexes
            )
            news_indexes = []
            candidate_title_indexes = []
            candidate_abstract_indexes = []
            cnt = 0
    if cnt > 0:
        yield iterator._convert_news_data(
            news_indexes, candidate_title_indexes, candidate_abstract_indexes
        )


def run_news2(model, news_filename):
    # print(news_filename)
    if not hasattr(model, "newsencoder"):
        raise ValueError("model must have attribute newsencoder")
    news_indexes = []
    news_vecs = []
    for batch_data_input in tqdm(
        load_news_from_file2(model.test_iterator, news_filename)
    ):
        news_index, news_vec = model.news(batch_data_input)
        news_indexes.extend(np.reshape(news_index, -1))
        news_vecs.extend(news_vec)
    return dict(zip(news_indexes, news_vecs))


def run_user2(model, news_filename, behaviors_file):
    if not hasattr(model, "userencoder"):
        raise ValueError("model must have attribute userencoder")
    user_indexes = []
    user_vecs = []
    for batch_data_input in tqdm(
        load_user_from_file2(model.test_iterator, news_filename, behaviors_file)
    ):
        user_index, user_vec = model.user(batch_data_input)
        user_indexes.extend(np.reshape(user_index, -1))
        user_vecs.extend(user_vec)
    return dict(zip(user_indexes, user_vecs))


def run_fast_eval2(model, news_filename, behaviors_file):
    news_vecs = run_news2(model, news_filename)
    # print(news_filename)
    user_vecs = run_user2(model, news_filename, behaviors_file)
    model.news_vecs = news_vecs
    model.user_vecs = user_vecs
    group_impr_indexes = []
    group_labels = []
    group_preds = []
    for (
        impr_index,
        news_index,
        user_index,
        label,
    ) in tqdm(load_impression_from_file2(model.test_iterator, behaviors_file)):
        pred = np.dot(
            np.stack([news_vecs[i] for i in news_index], axis=0),
            user_vecs[impr_index],
        )
        group_impr_indexes.append(impr_index)
        group_labels.append(label)
        group_preds.append(pred)
    return group_impr_indexes, group_labels, group_preds


group_impr_indexes, group_labels, group_preds = run_fast_eval2(
    model, test_news_file, test_behaviors_file
)
output_path = "/kaggle/working/"
with open(os.path.join(output_path, "prediction.txt"), "w") as f:
    for impr_index, preds in tqdm(zip(group_impr_indexes, group_preds)):
        impr_index += 1
        pred_rank = (np.argsort(np.argsort(preds)[::-1]) + 1).tolist()
        pred_rank = "[" + ",".join([str(i) for i in pred_rank]) + "]"
        f.write(" ".join([str(impr_index), pred_rank]) + "\n")
import zipfile

f = zipfile.ZipFile(
    os.path.join(output_path, "prediction.zip"), "w", zipfile.ZIP_DEFLATED
)
f.write(os.path.join(output_path, "prediction.txt"), arcname="prediction.txt")
f.close()
import matplotlib.pyplot as plt

fig = plt.figure()
plt.plot(np.arange(0, epochs), model.history["train_loss"], label="validation auc")
plt.plot(np.arange(0, epochs), model.history["val_loss"], label="validation auc")
plt.xticks(range(0, epochs))
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
<|endoftext|>
[{"Id": 94040247, "ScriptId": 26386486, "ParentScriptVersionId": NaN, "ScriptLanguageId": 9, "AuthorUserId": 8375020, "CreationDate": "04/26/2022 09:21:58", "VersionNumber": 4.0, "Title": "nrms_tf2_k14", "EvaluationDate": "04/26/2022", "IsChange": false, "TotalLines": 2241.0, "LinesInsertedFromPrevious": 0.0, "LinesChangedFromPrevious": 0.0, "LinesUnchangedFromPrevious": 2241.0, "LinesInsertedFromFork": NaN, "LinesDeletedFromFork": NaN, "LinesChangedFromFork": NaN, "LinesUnchangedFromFork": NaN, "TotalVotes": 0}]
[{"Id": 127926474, "KernelVersionId": 94040247, "SourceDatasetVersionId": 3491230}]
[{"Id": 3491230, "DatasetId": 2101503, "DatasourceVersionId": 3543661, "CreatorUserId": 8375020, "LicenseName": "Unknown", "CreationDate": "04/19/2022 09:41:36", "VersionNumber": 1.0, "Title": "kenh14-small", "Slug": "kenh14small", "Subtitle": NaN, "Description": NaN, "VersionNotes": "Initial release", "TotalCompressedBytes": 0.0, "TotalUncompressedBytes": 0.0}]
[{"Id": 2101503, "CreatorUserId": 8375020, "OwnerUserId": 8375020.0, "OwnerOrganizationId": NaN, "CurrentDatasetVersionId": 3491230.0, "CurrentDatasourceVersionId": 3543661.0, "ForumId": 2126900, "Type": 2, "CreationDate": "04/19/2022 09:41:36", "LastActivityDate": "04/19/2022", "TotalViews": 491, "TotalDownloads": 5, "TotalVotes": 0, "TotalKernels": 1}]
import yaml
import os
import abc
import time
import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
import pandas as pd
import pickle
import random
import re
from tqdm import tqdm
from tensorflow.keras import layers
from tensorflow.keras import backend as K
from collections import namedtuple

# detect and init the TPU
tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
# instantiate a distribution strategy
tpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)


def load_yaml(filename):
    """Load yaml file
    Args:
        filename
    Returns:
        dict
    """
    try:
        with open(filename, "r") as f:
            config = yaml.load(f, yaml.SafeLoader)
        return config
    except FileNotFoundError:
        raise
    except Exception as e:
        raise IOError(f"load {filename} error!")


def flat_config(config):
    """Flat config load tu yaml file to flat dict
    Args:
        config (dict): config load tu yaml file
    Returns:
        dict
    """
    f_config = {}
    category = config.keys()
    for cate in category:
        for key, val in config[cate].items():
            f_config[key] = val
    return f_config


def check_type(config):
    int_parameters = [
        "word_size",
        "his_size",
        "title_size",
        "body_size",
        "npratio",
        "word_emb_dim",
        "attention_hidden_dim",
        "epochs",
        "batch_size",
        "show_step",
        "save_epoch",
        "head_num",
        "head_dim",
        "user_num",
        "filter_num",
        "window_size",
        "gru_unit",
        "user_emb_dim",
        "vert_emb_dim",
        "subvert_emb_dim",
    ]
    for param in int_parameters:
        if param in config and not isinstance(config[param], int):
            raise TypeError("Parameters {0} must be int".format(param))
    float_parameters = ["learning_rate", "dropout"]
    for param in float_parameters:
        if param in config and not isinstance(config[param], float):
            raise TypeError("Parameters {0} must be float".format(param))
        str_parameters = [
            "wordEmb_file",
            "wordDict_file",
            "userDict_file",
            "vertDict_file",
            "subvertDict_file",
            "method",
            "loss",
            "optimizer",
            "cnn_activation",
            "dense_activation" "type",
        ]
    for param in str_parameters:
        if param in config and not isinstance(config[param], str):
            raise TypeError("Parameters {0} must be str".format(param))
    list_parameters = ["layer_sizes", "activation"]
    for param in list_parameters:
        if param in config and not isinstance(config[param], list):
            raise TypeError("Parameters {0} must be list".format(param))
    bool_parameters = ["support_quick_scoring"]
    for param in bool_parameters:
        if param in config and not isinstance(config[param], bool):
            raise TypeError("Parameters {0} must be bool".format(param))


def check_nn_config(f_config):
    """Check neural net config
    Args:
        f_config (dict): file config duoc flat tu yaml file
    Raises:
        ValueError: Neu params bi sai -> Raise error
    """
    if f_config["model_type"] in ["nrms", "NRMS"]:
        required_parameters = [
            "title_size",
            "his_size",
            "wordEmb_file",
            "wordDict_file",
            "userDict_file",
            "npratio",
            "data_format",
            "word_emb_dim",
            "head_num",
            "head_dim",
            "attention_hidden_dim",
            "loss",
            "data_format",
            "dropout",
        ]
    else:
        required_parameters = []
    # check
    for param in required_parameters:
        if param not in f_config:
            raise ValueError("Parameter {0} must be set!".format(param))
    if f_config["model_type"] in ["nrms", "NRMS"]:
        if f_config["data_format"] != "news":
            raise ValueError(
                "Voi NRMS model, dataformat phai la news, dua cai {0} vao lam gi".format(
                    f_config["data_format"]
                )
            )
    check_type(f_config)


def get_hparams(**kwargs):
    return namedtuple("GenericDict", kwargs.keys())(**kwargs)


def create_hparams(flags):
    """Create model's params
    Args:
        flags (dict): Dict co requirement
    Returns:
        object: namedtuple
    """
    return get_hparams(
        # data
        data_format=flags.get("data_format", None),
        iterator_type=flags.get("iterator_type", None),
        support_quick_scoring=flags.get("support_quick_scoring", False),
        wordEmb_file=flags.get("wordEmb_file", None),
        wordDict_file=flags.get("wordDict_file", None),
        userDict_file=flags.get("userDict_file", None),
        vertDict_file=flags.get("vertDict_file", None),
        subvertDict_file=flags.get("subvertDict_file", None),
        # models
        title_size=flags.get("title_size", None),
        body_size=flags.get("body_size", None),
        word_emb_dim=flags.get("word_emb_dim", None),
        word_size=flags.get("word_size", None),
        user_num=flags.get("user_num", None),
        vert_num=flags.get("vert_num", None),
        subvert_num=flags.get("subvert_num", None),
        his_size=flags.get("his_size", None),
        npratio=flags.get("npratio"),
        dropout=flags.get("dropout", 0.0),
        attention_hidden_dim=flags.get("attention_hidden_dim", 200),
        # nrms
        head_num=flags.get("head_num", 4),
        head_dim=flags.get("head_dim", 100),
        # train
        learning_rate=flags.get("learning_rate", 0.001),
        loss=flags.get("loss", None),
        optimizer=flags.get("optimizer", "adam"),
        epochs=flags.get("epochs", 10),
        batch_size=flags.get("batch_size", 1),
        # show info
        show_step=flags.get("show_step", 1),
        metrics=flags.get("metrics", None),
    )


def prepare_hparams(yaml_file=None, **kwargs):
    """Prepare hyperparams and make sure it's ok
    Args:
        yaml_file: path to yaml file
    Returns:
        TF Hyperparams object (tf.contrib.training.HParams)
    """
    if yaml_file is not None:
        config = load_yaml(yaml_file)
        config = flat_config(config)
    else:
        config = {}
    config.update(kwargs)
    check_nn_config(config)
    return create_hparams(config)


yaml_file = "/kaggle/input/kenh14small/nrms.yaml"
wordEmb_file = "/kaggle/input/mindlike-final/embedding.npy"
wordDict_file = "/kaggle/input/mindlike-final/word_dict.pkl"
userDict_file = "/kaggle/input/mindlike-final/uid2index.pkl"
epochs = 1
seed = 42
batch_size = 1024
hparams = prepare_hparams(
    yaml_file,
    his_size=25,
    wordEmb_file=wordEmb_file,
    wordDict_file=wordDict_file,
    userDict_file=userDict_file,
    batch_size=batch_size,
    epochs=epochs,
    show_step=10,
)
print(hparams)
from sklearn.metrics import (
    roc_auc_score,
    log_loss,
    mean_squared_error,
    accuracy_score,
    f1_score,
)


def mrr_score(y_true, y_score):
    """Computing mrr score metric.
    Args:
        y_true (np.ndarray): Ground-truth labels.
        y_score (np.ndarray): Predicted labels.
    Returns:
        numpy.ndarray: mrr scores.
    """
    order = np.argsort(y_score)[::-1]
    y_true = np.take(y_true, order)
    rr_score = y_true / (np.arange(len(y_true)) + 1)
    return np.sum(rr_score) / np.sum(y_true)


def ndcg_score(y_true, y_score, k=10):
    """Computing ndcg score metric at k.
    Args:
        y_true (np.ndarray): Ground-truth labels.
        y_score (np.ndarray): Predicted labels.
    Returns:
        numpy.ndarray: ndcg scores.
    """
    best = dcg_score(y_true, y_true, k)
    actual = dcg_score(y_true, y_score, k)
    return actual / best


def hit_score(y_true, y_score, k=10):
    """Computing hit score metric at k.
    Args:
        y_true (np.ndarray): ground-truth labels.
        y_score (np.ndarray): predicted labels.
    Returns:
        np.ndarray: hit score.
    """
    ground_truth = np.where(y_true == 1)[0]
    argsort = np.argsort(y_score)[::-1][:k]
    for idx in argsort:
        if idx in ground_truth:
            return 1
    return 0


def dcg_score(y_true, y_score, k=10):
    """Computing dcg score metric at k.
    Args:
        y_true (np.ndarray): Ground-truth labels.
        y_score (np.ndarray): Predicted labels.
    Returns:
        np.ndarray: dcg scores.
    """
    k = min(np.shape(y_true)[-1], k)
    order = np.argsort(y_score)[::-1]
    y_true = np.take(y_true, order[:k])
    gains = 2**y_true - 1
    discounts = np.log2(np.arange(len(y_true)) + 2)
    return np.sum(gains / discounts)


def cal_metric(labels, preds, metrics):
    print(labels)
    print(preds)
    """Calculate metrics.
    Available options are: `auc`, `rmse`, `logloss`, `acc` (accurary), `f1`, `mean_mrr`,
    `ndcg` (format like: ndcg@2;4;6;8), `hit` (format like: hit@2;4;6;8), `group_auc`.
    Args:
        labels (array-like): Labels.
        preds (array-like): Predictions.
        metrics (list): List of metric names.
    Return:
        dict: Metrics.
    Examples:
        >>> cal_metric(labels, preds, ["ndcg@2;4;6", "group_auc"])
        {'ndcg@2': 0.4026, 'ndcg@4': 0.4953, 'ndcg@6': 0.5346, 'group_auc': 0.8096}
    """
    res = {}
    for metric in metrics:
        if metric == "auc":
            auc = roc_auc_score(np.asarray(labels), np.asarray(preds))
            res["auc"] = round(auc, 4)
        elif metric == "rmse":
            rmse = mean_squared_error(np.asarray(labels), np.asarray(preds))
            res["rmse"] = np.sqrt(round(rmse, 4))
        elif metric == "logloss":
            # avoid logloss nan
            preds = [max(min(p, 1.0 - 10e-12), 10e-12) for p in preds]
            logloss = log_loss(np.asarray(labels), np.asarray(preds))
            res["logloss"] = round(logloss, 4)
        elif metric == "acc":
            pred = np.asarray(preds)
            pred[pred >= 0.5] = 1
            pred[pred < 0.5] = 0
            acc = accuracy_score(np.asarray(labels), pred)
            res["acc"] = round(acc, 4)
        elif metric == "f1":
            pred = np.asarray(preds)
            pred[pred >= 0.5] = 1
            pred[pred < 0.5] = 0
            f1 = f1_score(np.asarray(labels), pred)
            res["f1"] = round(f1, 4)
        elif metric == "mean_mrr":
            mean_mrr = np.mean(
                [
                    mrr_score(each_labels, each_preds)
                    for each_labels, each_preds in zip(labels, preds)
                ]
            )
            res["mean_mrr"] = round(mean_mrr, 4)
        elif metric.startswith("ndcg"):  # format like:  ndcg@2;4;6;8
            ndcg_list = [1, 2]
            ks = metric.split("@")
            if len(ks) > 1:
                ndcg_list = [int(token) for token in ks[1].split(";")]
            for k in ndcg_list:
                ndcg_temp = np.mean(
                    [
                        ndcg_score(each_labels, each_preds, k)
                        for each_labels, each_preds in zip(labels, preds)
                    ]
                )
                res["ndcg@{0}".format(k)] = round(ndcg_temp, 4)
        elif metric.startswith("hit"):  # format like:  hit@2;4;6;8
            hit_list = [1, 2]
            ks = metric.split("@")
            if len(ks) > 1:
                hit_list = [int(token) for token in ks[1].split(";")]
            for k in hit_list:
                hit_temp = np.mean(
                    [
                        hit_score(each_labels, each_preds, k)
                        for each_labels, each_preds in zip(labels, preds)
                    ]
                )
                res["hit@{0}".format(k)] = round(hit_temp, 4)
        elif metric == "group_auc":
            group_auc = np.mean(
                [
                    roc_auc_score(each_labels, each_preds)
                    for each_labels, each_preds in zip(labels, preds)
                ]
            )
            res["group_auc"] = round(group_auc, 4)
        else:
            raise ValueError("Metric {0} not defined".format(metric))
    return res


def word_tokenize(sent):
    """Split sentence into word list using regex.
    Args:
        sent (str): Input sentence
    Return:
        list: word list
    """
    pat = re.compile(r"[\w]+|[.,!?;|]")
    if isinstance(sent, str):
        return pat.findall(sent.lower())
    else:
        return []


def newsample(news, ratio):
    """Sample ratio samples from news list.
    If length of news is less than ratio, pad zeros.
    Args:
        news (list): input news list
        ratio (int): sample number
    Returns:
        list: output of sample list.
    """
    if ratio > len(news):
        return news + [0] * (ratio - len(news))
    else:
        return random.sample(news, ratio)


def abstractmethod(funcobj):
    """A decorator indicating abstract methods.
    Requires that the metaclass is ABCMeta or derived from it.  A
    class that has a metaclass derived from ABCMeta cannot be
    instantiated unless all of its abstract methods are overridden.
    The abstract methods can be called using any of the normal
    'super' call mechanisms.  abstractmethod() may be used to declare
    abstract methods for properties and descriptors.
    Usage:
        class C(metaclass=ABCMeta):
            @abstractmethod
            def my_abstract_method(self, ...):
                ...
    """
    funcobj.__isabstractmethod__ = True
    return funcobj


class BaseIterator(object):
    """Abstract base iterator class"""

    @abstractmethod
    def parser_one_line(self, line):
        """Abstract method. Parse one string line into feature values.
        Args:
            line (str): A string indicating one instance.
        """
        pass

    @abstractmethod
    def load_data_from_file(self, infile):
        """Abstract method. Read and parse data from a file.
        Args:
            infile (str): Text input file. Each line in this file is an instance.
        """
        pass

    @abstractmethod
    def _convert_data(self, labels, features):
        pass

    @abstractmethod
    def gen_feed_dict(self, data_dict):
        """Abstract method. Construct a dictionary that maps graph elements to values.
        Args:
            data_dict (dict): A dictionary that maps string name to numpy arrays.
        """
        pass


def load_dict(file_path):
    with open(file_path, "rb") as f:
        return pickle.load(f)


class MINDIterator(BaseIterator):
    """Train data loader for NRMS model.
    The model require a special type of data format, where each instance contains a label, impresion id, user id,
    the candidate news articles and user's clicked news article. Articles are represented by title words,
    body words, verts and subverts.
    Iterator will not load the whole data into memory. Instead, it loads data into memory
    per mini-batch, so that large files can be used as input data.
    Attributes:
        col_spliter (str): column spliter in one line.
        ID_spliter (str): ID spliter in one line.
        batch_size (int): the samples num in one batch.
        title_size (int): max word num in news title.
        his_size (int): max clicked news num in user click history.
        npratio (int): negaive and positive ratio used in negative sampling. -1 means no need of negtive sampling.
    """

    def __init__(
        self,
        hparams,
        npratio=-1,
        col_spliter="\t",
        ID_spliter="%",
    ):
        """Initialize an iterator. Create necessary placeholders for the model.
        Args:
            hparams (object): Global hyper-parameters. Some key setttings such as head_num and head_dim are there.
            npratio (int): negaive and positive ratio used in negative sampling. -1 means no need of negtive sampling.
            col_spliter (str): column spliter in one line.
            ID_spliter (str): ID spliter in one line.
        """
        self.col_spliter = col_spliter
        self.ID_spliter = ID_spliter
        self.batch_size = hparams.batch_size
        self.title_size = hparams.title_size
        self.body_size = hparams.body_size
        self.his_size = hparams.his_size
        self.npratio = npratio
        self.word_dict = self.load_dict(hparams.wordDict_file)
        # self.cat_dict = self.load_dict(hparams.catDict_file)
        # self.subcat_dict = self.load_dict(hparams.subcatDict_file)
        # self.entity_dict = self.load_dict(hparams.entityDict_file)
        # self.relation_dict = self.load_dict(hparams.relationDict_file)
        self.uid2index = self.load_dict(hparams.userDict_file)

    def load_dict(self, file_path):
        """load pickle file
        Args:
            file path (str): file path
        Returns:
            object: pickle loaded object
        """
        with open(file_path, "rb") as f:
            return pickle.load(f)

    def init_news(self, news_file):
        """init news information given news file, such as news_title_index and nid2index.
        Args:
            news_file: path of news file
        """
        self.nid2index = {}
        news_title = [""]
        # news_cat = [""]
        # news_subcat = [""]
        news_abstract = [""]
        # news_entitites = [""]
        # news_relation = [""]
        with tf.io.gfile.GFile(news_file, "r") as rd:
            for line in rd:
                nid, cat, title, ab, content = line.strip("\n").split(self.col_spliter)
                if nid in self.nid2index:
                    continue
                self.nid2index[nid] = len(self.nid2index) + 1
                title = word_tokenize(title)
                if ab is None:
                    abstract = ""
                else:
                    abstract = word_tokenize(ab)
                news_title.append(title)
                # news_cat.append(cat)
                # news_subcat.append(subcat)
                news_abstract.append(abstract)
                # news_entities.append(entity)
                # news_relation.append(relation)
        self.news_title_index = np.zeros(
            (len(news_title), self.title_size), dtype="int32"
        )
        self.news_abstract_index = np.zeros(
            (len(news_abstract), self.body_size), dtype="int32"
        )
        # self.news_cat_index = np.zeros((len(news_cat), 1), dtype="int32")
        # self.news_subcat_index = np.zeros((len(news_subcat, 1), dtype="int32")
        for news_index in range(len(news_title)):
            title = news_title[news_index]
            abstract = news_abstract[news_index]
            for word_index in range(min(self.title_size, len(title))):
                if title[word_index] in self.word_dict:
                    self.news_title_index[news_index, word_index] = self.word_dict[
                        title[word_index].lower()
                    ]
            for word_index in range(min(self.body_size, len(abstract))):
                if abstract[word_index] in self.word_dict:
                    self.news_abstract_index[news_index, word_index] = self.word_dict[
                        abstract[word_index].lower()
                    ]

    def init_behaviors(self, behaviors_file):
        """init behavior logs given behaviors file.
        Args:
        behaviors_file: path of behaviors file
        """
        self.histories = []
        self.imprs = []
        self.labels = []
        self.impr_indexes = []
        self.uindexes = []
        with tf.io.gfile.GFile(behaviors_file, "r") as rd:
            impr_index = 0
            for line in rd:
                uid, history, impr, _ = line.strip("\n").split(self.col_spliter)[-4:]
                history = [self.nid2index[i] for i in history.split()]
                history = [0] * (self.his_size - len(history)) + history[
                    : self.his_size
                ]
                impr_news = [self.nid2index[i.split("-")[0]] for i in impr.split()]
                label = [int(i.split("-")[1]) for i in impr.split()]
                uindex = self.uid2index[uid] if uid in self.uid2index else 0
                self.histories.append(history)
                self.imprs.append(impr_news)
                self.labels.append(label)
                self.impr_indexes.append(impr_index)
                self.uindexes.append(uindex)
                impr_index += 1

    def parser_one_line(self, line):
        """Parse one behavior sample into feature values.
        if npratio is larger than 0, return negtive sampled result.
        Args:
            line (int): sample index.
        Yields:
            list: Parsed results including label, impression id , user id,
            candidate_title_index, clicked_title_index.
        """
        if self.npratio > 0:
            impr_label = self.labels[line]
            impr = self.imprs[line]
            poss = []
            negs = []
            for news, click in zip(impr, impr_label):
                if click == 1:
                    poss.append(news)
                else:
                    negs.append(news)
            for p in poss:
                candidate_title_index = []
                candidate_abstract_index = []
                impr_index = []
                user_index = []
                label = [1] + [0] * self.npratio
                n = newsample(negs, self.npratio)
                candidate_title_index = self.news_title_index[[p] + n]
                candidate_abstract_index = self.news_abstract_index[[p] + n]
                click_title_index = self.news_title_index[self.histories[line]]
                click_abstract_index = self.news_abstract_index[self.histories[line]]
                impr_index.append(self.impr_indexes[line])
                user_index.append(self.uindexes[line])
                yield (
                    label,
                    impr_index,
                    user_index,
                    candidate_title_index,
                    candidate_abstract_index,
                    click_title_index,
                    click_abstract_index,
                )
        else:
            impr_label = self.labels[line]
            impr = self.imprs[line]
            for news, label in zip(impr, impr_label):
                candidate_title_index = []
                impr_index = []
                user_index = []
                label = [label]
                candidate_title_index.append(self.news_title_index[news])
                candidate_abstract_index = self.news_abstract_index[[p] + n]
                click_title_index = self.news_title_index[self.histories[line]]
                click_abstract_index = self.news_abstract_index[self.histories]
                impr_index.append(self.impr_indexes[line])
                user_index.append(self.uindexes[line])
                yield (
                    label,
                    impr_index,
                    user_index,
                    candidate_title_index,
                    candidate_abstract_index,
                    click_title_index,
                    click_abstract_index,
                )

    def load_data_from_file(self, news_file, behavior_file):
        """Read and parse data from news file and behavior file.
        Args:
            news_file (str): A file contains several informations of news.
            beahavior_file (str): A file contains information of user impressions.
        Yields:
            object: An iterator that yields parsed results, in the format of dict.
        """
        if not hasattr(self, "news_title_index"):
            self.init_news(news_file)
        if not hasattr(self, "impr_indexes"):
            self.init_behaviors(behavior_file)
        label_list = []
        imp_indexes = []
        user_indexes = []
        candidate_title_indexes = []
        candidate_abstract_indexes = []
        click_title_indexes = []
        click_abstract_indexes = []
        cnt = 0
        indexes = np.arange(len(self.labels))
        if self.npratio > 0:
            np.random.shuffle(indexes)
        for index in indexes:
            for (
                label,
                imp_index,
                user_index,
                candidate_title_index,
                candidate_abstract_index,
                click_title_index,
                click_abstract_index,
            ) in self.parser_one_line(index):
                candidate_title_indexes.append(candidate_title_index)
                candidate_abstract_indexes.append(candidate_abstract_index)
                click_title_indexes.append(click_title_index)
                click_abstract_indexes.append(click_abstract_index)
                imp_indexes.append(imp_index)
                user_indexes.append(user_index)
                label_list.append(label)
                cnt += 1
                if cnt >= self.batch_size:
                    yield self._convert_data(
                        label_list,
                        imp_indexes,
                        user_indexes,
                        candidate_title_indexes,
                        candidate_abstract_indexes,
                        click_title_indexes,
                        click_abstract_indexes,
                    )
                    label_list = []
                    imp_indexes = []
                    user_indexes = []
                    candidate_title_indexes = []
                    candidate_abstract_indexes = []
                    click_title_indexes = []
                    click_abstract_indexes = []
                    cnt = 0
        if cnt > 0:
            yield self._convert_data(
                label_list,
                imp_indexes,
                user_indexes,
                candidate_title_indexes,
                candidate_abstract_indexes,
                click_title_indexes,
                click_abstract_indexes,
            )

    def _convert_data(
        self,
        label_list,
        imp_indexes,
        user_indexes,
        candidate_title_indexes,
        candidate_abstract_indexes,
        click_title_indexes,
        click_abstract_indexes,
    ):
        """Convert data into numpy arrays that are good for further model operation.
        Args:
            label_list (list): a list of ground-truth labels.
            imp_indexes (list): a list of impression indexes.
            user_indexes (list): a list of user indexes.
            candidate_title_indexes (list): the candidate news titles' words indices.
            candidate_abstract_indexes (list): the candidate news abstract' words indices.
            click_title_indexes (list): words indices for user's clicked news titles.
            click_abstract_indexes (list): words indices for user's clicked news abstract.
        Returns:
            dict: A dictionary, containing multiple numpy arrays that are convenient for further operation.
        """
        labels = np.asarray(label_list, dtype=np.float32)
        imp_indexes = np.asarray(imp_indexes, dtype=np.int32)
        user_indexes = np.asarray(user_indexes, dtype=np.int32)
        candidate_title_index_batch = np.asarray(
            candidate_title_indexes, dtype=np.int64
        )
        candidate_abstract_index_batch = np.asarray(
            candidate_abstract_indexes, dtype=np.int64
        )
        click_title_index_batch = np.asarray(click_title_indexes, dtype=np.int64)
        click_abstract_index_batch = np.asarray(click_abstract_indexes, dtype=np.int64)
        return {
            "impression_index_batch": imp_indexes,
            "user_index_batch": user_indexes,
            "clicked_title_batch": click_title_index_batch,
            "clicked_abstract_batch": click_abstract_index_batch,
            "candidate_title_batch": candidate_title_index_batch,
            "candidate_abstract_batch": candidate_abstract_index_batch,
            "labels": labels,
        }

    def load_user_from_file(self, news_file, behavior_file):
        """Read and parse user data from news file and behavior file.
        Args:
            news_file (str): A file contains several informations of news.
            beahaviros_file (str): A file contains information of user impressions.
        Yields:
            object: An iterator that yields parsed user feature, in the format of dict.
        """
        if not hasattr(self, "news_title_index"):
            self.init_news(news_file)
        if not hasattr(self, "impr_indexes"):
            self.init_behaviors(behavior_file)
        user_indexes = []
        impr_indexes = []
        click_title_indexes = []
        click_abstract_indexes = []
        cnt = 0
        for index in range(len(self.impr_indexes)):
            click_title_indexes.append(self.news_title_index[self.histories[index]])
            click_abstract_indexes.append(
                self.news_abstract_index[self.histories[index]]
            )
            user_indexes.append(self.uindexes[index])
            impr_indexes.append(self.impr_indexes[index])
            cnt += 1
            if cnt >= self.batch_size:
                yield self._convert_user_data(
                    user_indexes,
                    impr_indexes,
                    click_title_indexes,
                    click_abstract_indexes,
                )
                user_indexes = []
                impr_indexes = []
                click_title_indexes = []
                click_abstract_indexes = []
                cnt = 0
        if cnt > 0:
            yield self._convert_user_data(
                user_indexes, impr_indexes, click_title_indexes, click_abstract_indexes
            )

    def _convert_user_data(
        self, user_indexes, impr_indexes, click_title_indexes, click_abstract_indexes
    ):
        """Convert data into numpy arrays that are good for further model operation.
        Args:
            user_indexes (list): a list of user indexes.
            click_title_indexes (list): words indices for user's clicked news titles.
        Returns:
            dict: A dictionary, containing multiple numpy arrays that are convenient for further operation.
        """
        user_indexes = np.asarray(user_indexes, dtype=np.int32)
        impr_indexes = np.asarray(impr_indexes, dtype=np.int32)
        click_title_index_batch = np.asarray(click_title_indexes, dtype=np.int64)
        click_abstract_index_batch = np.asarray(click_abstract_indexes, dtype=np.int64)
        return {
            "user_index_batch": user_indexes,
            "impr_index_batch": impr_indexes,
            "clicked_title_batch": click_title_index_batch,
            "clicked_abstract_batch": click_abstract_index_batch,
        }

    def load_news_from_file(self, news_file):
        """Read and parse user data from news file.
        Args:
            news_file (str): A file contains several informations of news.
        Yields:
            object: An iterator that yields parsed news feature, in the format of dict.
        """
        if not hasattr(self, "news_title_index"):
            self.init_news(news_file)
        news_indexes = []
        candidate_title_indexes = []
        candidate_abstract_indexes = []
        cnt = 0
        for index in range(len(self.news_title_index)):
            news_indexes.append(index)
            candidate_title_indexes.append(self.news_title_index[index])
            candidate_abstract_indexes.append(self.news_abstract_index[index])
            cnt += 1
            if cnt >= self.batch_size:
                yield self._convert_news_data(
                    news_indexes, candidate_title_indexes, candidate_abstract_indexes
                )
                news_indexes = []
                candidate_title_indexes = []
                candidate_abstract_indexes = []
                cnt = 0
        if cnt > 0:
            yield self._convert_news_data(
                news_indexes, candidate_title_indexes, candidate_abstract_indexes
            )

    def _convert_news_data(
        self, news_indexes, candidate_title_indexes, candidate_abstract_indexes
    ):
        """Convert data into numpy arrays that are good for further model operation.
        Args:
            news_indexes (list): a list of news indexes.
            candidate_title_indexes (list): the candidate news titles' words indices.
        Returns:
            dict: A dictionary, containing multiple numpy arrays that are convenient for further operation.
        """
        news_indexes_batch = np.asarray(news_indexes, dtype=np.int32)
        candidate_title_index_batch = np.asarray(
            candidate_title_indexes, dtype=np.int32
        )
        candidate_abstract_index_batch = np.asarray(
            candidate_abstract_indexes, dtype=np.int32
        )
        return {
            "news_index_batch": news_indexes_batch,
            "candidate_title_batch": candidate_title_index_batch,
            "candidate_abstract_batch": candidate_abstract_index_batch,
        }

    def load_impression_from_file(self, behaivors_file):
        """Read and parse impression data from behaivors file.
        Args:
            behaivors_file (str): A file contains several informations of behaviros.
        Yields:
            object: An iterator that yields parsed impression data, in the format of dict.
        """
        if not hasattr(self, "histories"):
            self.init_behaviors(behaivors_file)
        indexes = np.arange(len(self.labels))
        for index in indexes:
            impr_label = np.array(self.labels[index], dtype="int32")
            impr_news = np.array(self.imprs[index], dtype="int32")
            yield (
                self.impr_indexes[index],
                impr_news,
                self.uindexes[index],
                impr_label,
            )


# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
class AttLayer2(layers.Layer):
    """Soft alignment attention implement.
    Attributes:
        dim (int): attention hidden dim
    """

    def __init__(self, dim=200, seed=0, **kwargs):
        """Initialization steps for AttLayer2.
        Args:
            dim (int): attention hidden dim
        """
        self.dim = dim
        self.seed = seed
        super(AttLayer2, self).__init__(**kwargs)

    def build(self, input_shape):
        """Initialization for variables in AttLayer2
        There are there variables in AttLayer2, i.e. W, b and q.
        Args:
            input_shape (object): shape of input tensor.
        """
        assert len(input_shape) == 3
        dim = self.dim
        self.W = self.add_weight(
            name="W",
            shape=(int(input_shape[-1]), dim),
            initializer=keras.initializers.glorot_uniform(seed=self.seed),
            trainable=True,
        )
        self.b = self.add_weight(
            name="b",
            shape=(dim,),
            initializer=keras.initializers.Zeros(),
            trainable=True,
        )
        self.q = self.add_weight(
            name="q",
            shape=(dim, 1),
            initializer=keras.initializers.glorot_uniform(seed=self.seed),
            trainable=True,
        )
        super(AttLayer2, self).build(input_shape)  # be sure you call this somewhere!

    def call(self, inputs, mask=None, **kwargs):
        """Core implemention of soft attention
        Args:
            inputs (object): input tensor.
        Returns:
            object: weighted sum of input tensors.
        """
        attention = K.tanh(K.dot(inputs, self.W) + self.b)
        attention = K.dot(attention, self.q)
        attention = K.squeeze(attention, axis=2)
        if mask == None:
            attention = K.exp(attention)
        else:
            attention = K.exp(attention) * K.cast(mask, dtype="float32")
        attention_weight = attention / (
            K.sum(attention, axis=-1, keepdims=True) + K.epsilon()
        )
        attention_weight = K.expand_dims(attention_weight)
        weighted_input = inputs * attention_weight
        return K.sum(weighted_input, axis=1)

    def compute_mask(self, input, input_mask=None):
        """Compte output mask value
        Args:
            input (object): input tensor.
            input_mask: input mask
        Returns:
            object: output mask.
        """
        return None

    def compute_output_shape(self, input_shape):
        """Compute shape of output tensor
        Args:
            input_shape (tuple): shape of input tensor.
        Returns:
            tuple: shape of output tensor.
        """
        return input_shape[0], input_shape[-1]


class SelfAttention(layers.Layer):
    """Multi-head self attention implement.
    Args:
        multiheads (int): The number of heads.
        head_dim (object): Dimention of each head.
        mask_right (boolean): whether to mask right words.
    Returns:
        object: Weighted sum after attention.
    """

    def __init__(self, multiheads, head_dim, seed=0, mask_right=False, **kwargs):
        """Initialization steps for AttLayer2.
        Args:
            multiheads (int): The number of heads.
            head_dim (object): Dimention of each head.
            mask_right (boolean): whether to mask right words.
        """
        self.multiheads = multiheads
        self.head_dim = head_dim
        self.output_dim = multiheads * head_dim
        self.mask_right = mask_right
        self.seed = seed
        super(SelfAttention, self).__init__(**kwargs)

    def compute_output_shape(self, input_shape):
        """Compute shape of output tensor.
        Returns:
            tuple: output shape tuple.
        """
        return (input_shape[0][0], input_shape[0][1], self.output_dim)

    def build(self, input_shape):
        """Initialization for variables in SelfAttention.
        There are three variables in SelfAttention, i.e. WQ, WK ans WV.
        WQ is used for linear transformation of query.
        WK is used for linear transformation of key.
        WV is used for linear transformation of value.
        Args:
            input_shape (object): shape of input tensor.
        """
        self.WQ = self.add_weight(
            name="WQ",
            shape=(int(input_shape[0][-1]), self.output_dim),
            initializer=keras.initializers.glorot_uniform(seed=self.seed),
            trainable=True,
        )
        self.WK = self.add_weight(
            name="WK",
            shape=(int(input_shape[1][-1]), self.output_dim),
            initializer=keras.initializers.glorot_uniform(seed=self.seed),
            trainable=True,
        )
        self.WV = self.add_weight(
            name="WV",
            shape=(int(input_shape[2][-1]), self.output_dim),
            initializer=keras.initializers.glorot_uniform(seed=self.seed),
            trainable=True,
        )
        super(SelfAttention, self).build(input_shape)

    def Mask(self, inputs, seq_len, mode="add"):
        """Mask operation used in multi-head self attention
        Args:
            seq_len (object): sequence length of inputs.
            mode (str): mode of mask.
        Returns:
            object: tensors after masking.
        """
        if seq_len == None:
            return inputs
        else:
            mask = K.one_hot(indices=seq_len[:, 0], num_classes=K.shape(inputs)[1])
            mask = 1 - K.cumsum(mask, axis=1)
            for _ in range(len(inputs.shape) - 2):
                mask = K.expand_dims(mask, 2)
            if mode == "mul":
                return inputs * mask
            elif mode == "add":
                return inputs - (1 - mask) * 1e12

    def call(self, QKVs):
        """Core logic of multi-head self attention.
        Args:
            QKVs (list): inputs of multi-head self attention i.e. qeury, key and value.
        Returns:
            object: ouput tensors.
        """
        if len(QKVs) == 3:
            Q_seq, K_seq, V_seq = QKVs
            Q_len, V_len = None, None
        elif len(QKVs) == 5:
            Q_seq, K_seq, V_seq, Q_len, V_len = QKVs
        Q_seq = K.dot(Q_seq, self.WQ)
        Q_seq = K.reshape(
            Q_seq, shape=(-1, K.shape(Q_seq)[1], self.multiheads, self.head_dim)
        )
        Q_seq = K.permute_dimensions(Q_seq, pattern=(0, 2, 1, 3))
        K_seq = K.dot(K_seq, self.WK)
        K_seq = K.reshape(
            K_seq, shape=(-1, K.shape(K_seq)[1], self.multiheads, self.head_dim)
        )
        K_seq = K.permute_dimensions(K_seq, pattern=(0, 2, 1, 3))
        V_seq = K.dot(V_seq, self.WV)
        V_seq = K.reshape(
            V_seq, shape=(-1, K.shape(V_seq)[1], self.multiheads, self.head_dim)
        )
        V_seq = K.permute_dimensions(V_seq, pattern=(0, 2, 1, 3))
        # tf.einsum('m b i k, m b j k -> m b i j', Q_seq , K_seq) # shape [10, 20, 50]
        A = tf.einsum("m b i k, m b j k -> m b i j", Q_seq, K_seq) / K.sqrt(
            K.cast(self.head_dim, dtype="float32")
        )
        A = K.permute_dimensions(
            A, pattern=(0, 3, 2, 1)
        )  # A.shape=[batch_size,K_sequence_length,Q_sequence_length,self.multiheads]
        A = self.Mask(A, V_len, "add")
        A = K.permute_dimensions(A, pattern=(0, 3, 2, 1))
        if self.mask_right:
            ones = K.ones_like(A[:1, :1])
            lower_triangular = K.tf.matrix_band_part(ones, num_lower=-1, num_upper=0)
            mask = (ones - lower_triangular) * 1e12
            A = A - mask
        A = K.softmax(A)
        # tf.einsum('m b i k, m b k j -> m b i j', A , V_seq)
        O_seq = tf.einsum("m b i k, m b k j -> m b i j", A, V_seq)
        O_seq = K.permute_dimensions(O_seq, pattern=(0, 2, 1, 3))
        O_seq = K.reshape(O_seq, shape=(-1, K.shape(O_seq)[1], self.output_dim))
        O_seq = self.Mask(O_seq, Q_len, "mul")
        return O_seq

    def get_config(self):
        """add multiheads, multiheads and mask_right into layer config.
        Returns:
            dict: config of SelfAttention layer.
        """
        config = super(SelfAttention, self).get_config()
        config.update(
            {
                "multiheads": self.multiheads,
                "head_dim": self.head_dim,
                "mask_right": self.mask_right,
            }
        )
        return config


__all__ = ["BaseModel"]


class BaseModel:
    """Basic class of models
    Attributes:
        hparams (object): A tf.contrib.training.HParams object, hold the entire set of hyperparameters.
        train_iterator (object): An iterator to load the data in training steps.
        test_iterator (object): An iterator to load the data in testing steps.
        graph (object): An optional graph.
        seed (int): Random seed.
    """

    def __init__(
        self,
        hparams,
        iterator_creator,
        seed=None,
    ):
        """Initializing the model. Create common logics which are needed by all deeprec models, such as loss function,
        parameter set.
        Args:
            hparams (object): A tf.contrib.training.HParams object, hold the entire set of hyperparameters.
            iterator_creator (object): An iterator to load the data.
            graph (object): An optional graph.
            seed (int): Random seed.
        """
        self.seed = seed
        tf.compat.v1.set_random_seed(seed)
        np.random.seed(seed)
        self.train_iterator = iterator_creator(
            hparams,
            hparams.npratio,
            col_spliter="\t",
        )
        self.test_iterator = iterator_creator(
            hparams,
            col_spliter="\t",
        )
        self.hparams = hparams
        self.support_quick_scoring = hparams.support_quick_scoring
        self.model, self.scorer = self._build_graph()
        self.loss = self._get_loss()
        self.train_optimizer = self._get_opt()
        self.model.compile(loss=self.loss, optimizer=self.train_optimizer)

    def _init_embedding(self, file_path):
        """Load pre-trained embeddings as a constant tensor.
        Args:
            file_path (str): the pre-trained glove embeddings file path.
        Returns:
            numpy.ndarray: A constant numpy array.
        """
        return np.load(file_path)

    @abc.abstractmethod
    def _build_graph(self):
        """Subclass will implement this."""
        pass

    @abc.abstractmethod
    def _get_input_label_from_iter(self, batch_data):
        """Subclass will implement this"""
        pass

    def _get_loss(self):
        """Make loss function, consists of data loss and regularization loss
        Returns:
            object: Loss function or loss function name
        """
        if self.hparams.loss == "cross_entropy_loss":
            data_loss = "categorical_crossentropy"
        elif self.hparams.loss == "log_loss":
            data_loss = "binary_crossentropy"
        else:
            raise ValueError("this loss not defined {0}".format(self.hparams.loss))
        return data_loss

    def _get_opt(self):
        """Get the optimizer according to configuration. Usually we will use Adam.
        Returns:
            object: An optimizer.
        """
        lr = self.hparams.learning_rate
        optimizer = self.hparams.optimizer
        if optimizer == "adam":
            train_opt = tf.keras.optimizers.Adam(lr=lr)
        return train_opt

    def _get_pred(self, logit, task):
        """Make final output as prediction score, according to different tasks.
        Args:
            logit (object): Base prediction value.
            task (str): A task (values: regression/classification)
        Returns:
            object: Transformed score
        """
        if task == "regression":
            pred = tf.identity(logit)
        elif task == "classification":
            pred = tf.sigmoid(logit)
        else:
            raise ValueError(
                "method must be regression or classification, but now is {0}".format(
                    task
                )
            )
        return pred

    def train(self, train_batch_data):
        """Go through the optimization step once with training data in feed_dict.
        Args:
            sess (object): The model session object.
            feed_dict (dict): Feed values to train the model. This is a dictionary that maps graph elements to values.
        Returns:
            list: A list of values, including update operation, total loss, data loss, and merged summary.
        """
        train_input, train_label = self._get_input_label_from_iter(train_batch_data)
        rslt = self.model.train_on_batch(train_input, train_label)
        return rslt

    def eval(self, eval_batch_data):
        """Evaluate the data in feed_dict with current model.
        Args:
            sess (object): The model session object.
            feed_dict (dict): Feed values for evaluation. This is a dictionary that maps graph elements to values.
        Returns:
            list: A list of evaluated results, including total loss value, data loss value, predicted scores, and ground-truth labels.
        """
        eval_input, eval_label = self._get_input_label_from_iter(eval_batch_data)
        imp_index = eval_batch_data["impression_index_batch"]
        pred_rslt = self.scorer.predict_on_batch(eval_input)
        return pred_rslt, eval_label, imp_index

    def fit(
        self,
        train_news_file,
        train_behaviors_file,
        valid_news_file,
        valid_behaviors_file,
        test_news_file=None,
        test_behaviors_file=None,
    ):
        """Fit the model with train_file. Evaluate the model on valid_file per epoch to observe the training status.
        If test_news_file is not None, evaluate it too.
        Args:
            train_file (str): training data set.
            valid_file (str): validation set.
            test_news_file (str): test set.
        Returns:
            object: An instance of self.
        """
        self.history = {}
        self.history["train_loss"] = []
        self.history["val_loss"] = []
        self.history["Val_auc"] = []
        for epoch in range(1, self.hparams.epochs + 1):
            step = 0
            # self.hparams.current_epoch = epoch
            epoch_loss = 0
            train_start = time.time()
            tqdm_util = tqdm(
                self.train_iterator.load_data_from_file(
                    train_news_file, train_behaviors_file
                )
            )
            for batch_data_input in tqdm_util:
                step_result = self.train(batch_data_input)
                step_data_loss = step_result
                epoch_loss += step_data_loss
                step += 1
                if step % self.hparams.show_step == 0:
                    tqdm_util.set_description(
                        "step {0:d} , total_loss: {1:.4f}, data_loss: {2:.4f}".format(
                            step, epoch_loss / step, step_data_loss
                        )
                    )
            self.history["train_loss"].append(epoch_loss / step)
            self.history["val_loss"].append(step_data_loss)
            train_end = time.time()
            train_time = train_end - train_start
            eval_start = time.time()
            train_info = ",".join(
                [
                    str(item[0]) + ":" + str(item[1])
                    for item in [("logloss loss", epoch_loss / step)]
                ]
            )
            eval_res = self.run_eval(valid_news_file, valid_behaviors_file)
            eval_info = ", ".join(
                [
                    str(item[0]) + ":" + str(item[1])
                    for item in sorted(eval_res.items(), key=lambda x: x[0])
                ]
            )
            if test_news_file is not None:
                test_res = self.run_eval(test_news_file, test_behaviors_file)
                test_info = ", ".join(
                    [
                        str(item[0]) + ":" + str(item[1])
                        for item in sorted(test_res.items(), key=lambda x: x[0])
                    ]
                )
            eval_end = time.time()
            eval_time = eval_end - eval_start
            if test_news_file is not None:
                print(
                    "at epoch {0:d}".format(epoch)
                    + "\ntrain info: "
                    + train_info
                    + "\neval info: "
                    + eval_info
                    + "\ntest info: "
                    + test_info
                )
            else:
                print(
                    "at epoch {0:d}".format(epoch)
                    + "\ntrain info: "
                    + train_info
                    + "\neval info: "
                    + eval_info
                )
            print(
                "at epoch {0:d} , train time: {1:.1f} eval time: {2:.1f}".format(
                    epoch, train_time, eval_time
                )
            )
        return self

    def group_labels(self, labels, preds, group_keys):
        """Devide labels and preds into several group according to values in group keys.
        Args:
            labels (list): ground truth label list.
            preds (list): prediction score list.
            group_keys (list): group key list.
        Returns:
            list, list, list:
            - Keys after group.
            - Labels after group.
            - Preds after group.
        """
        all_keys = list(set(group_keys))
        all_keys.sort()
        group_labels = {k: [] for k in all_keys}
        group_preds = {k: [] for k in all_keys}
        for l, p, k in zip(labels, preds, group_keys):
            group_labels[k].append(l)
            group_preds[k].append(p)
        all_labels = []
        all_preds = []
        for k in all_keys:
            all_labels.append(group_labels[k])
            all_preds.append(group_preds[k])
        return all_keys, all_labels, all_preds

    def run_eval(self, news_filename, behaviors_file):
        """Evaluate the given file and returns some evaluation metrics.
        Args:
            filename (str): A file name that will be evaluated.
        Returns:
            dict: A dictionary that contains evaluation metrics.
        """
        if self.support_quick_scoring:
            _, group_labels, group_preds = self.run_fast_eval(
                news_filename, behaviors_file
            )
        else:
            _, group_labels, group_preds = self.run_slow_eval(
                news_filename, behaviors_file
            )
        res = cal_metric(group_labels, group_preds, self.hparams.metrics)
        return res

    def user(self, batch_user_input):
        user_input = self._get_user_feature_from_iter(batch_user_input)
        user_vec = self.userencoder.predict_on_batch(user_input)
        user_index = batch_user_input["impr_index_batch"]
        return user_index, user_vec

    def news(self, batch_news_input):
        news_input = self._get_news_feature_from_iter(batch_news_input)
        news_vec = self.newsencoder.predict_on_batch(news_input)
        news_index = batch_news_input["news_index_batch"]
        return news_index, news_vec

    def run_user(self, news_filename, behaviors_file):
        if not hasattr(self, "userencoder"):
            raise ValueError("model must have attribute userencoder")
        user_indexes = []
        user_vecs = []
        for batch_data_input in tqdm(
            self.test_iterator.load_user_from_file(news_filename, behaviors_file)
        ):
            # print("\nbatch_data_input:", batch_data_input)
            user_index, user_vec = self.user(batch_data_input)
            user_indexes.extend(np.reshape(user_index, -1))
            user_vecs.extend(user_vec)
        return dict(zip(user_indexes, user_vecs))

    def run_news(self, news_filename):
        if not hasattr(self, "newsencoder"):
            raise ValueError("model must have attribute newsencoder")
        news_indexes = []
        news_vecs = []
        for batch_data_input in tqdm(
            self.test_iterator.load_news_from_file(news_filename)
        ):
            news_index, news_vec = self.news(batch_data_input)
            news_indexes.extend(np.reshape(news_index, -1))
            news_vecs.extend(news_vec)
        return dict(zip(news_indexes, news_vecs))

    def run_slow_eval(self, news_filename, behaviors_file):
        preds = []
        labels = []
        imp_indexes = []
        for batch_data_input in tqdm(
            self.test_iterator.load_data_from_file(news_filename, behaviors_file)
        ):
            step_pred, step_labels, step_imp_index = self.eval(batch_data_input)
            preds.extend(np.reshape(step_pred, -1))
            labels.extend(np.reshape(step_labels, -1))
            imp_indexes.extend(np.reshape(step_imp_index, -1))
        group_impr_indexes, group_labels, group_preds = self.group_labels(
            labels, preds, imp_indexes
        )
        return group_impr_indexes, group_labels, group_preds

    def run_fast_eval(self, news_filename, behaviors_file):
        news_vecs = self.run_news(news_filename)
        user_vecs = self.run_user(news_filename, behaviors_file)
        self.news_vecs = news_vecs
        self.user_vecs = user_vecs
        group_impr_indexes = []
        group_labels = []
        group_preds = []
        for (
            impr_index,
            news_index,
            user_index,
            label,
        ) in tqdm(self.test_iterator.load_impression_from_file(behaviors_file)):
            pred = np.dot(
                np.stack([news_vecs[i] for i in news_index], axis=0),
                user_vecs[impr_index],
            )
            group_impr_indexes.append(impr_index)
            group_labels.append(label)
            group_preds.append(pred)
        return group_impr_indexes, group_labels, group_preds


# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
__all__ = ["NRMSModel"]


class NRMSModel(BaseModel):
    """NRMS model(Neural News Recommendation with Multi-Head Self-Attention)
    Chuhan Wu, Fangzhao Wu, Suyu Ge, Tao Qi, Yongfeng Huang,and Xing Xie, "Neural News
    Recommendation with Multi-Head Self-Attention" in Proceedings of the 2019 Conference
    on Empirical Methods in Natural Language Processing and the 9th International Joint Conference
    on Natural Language Processing (EMNLP-IJCNLP)
    Attributes:
        word2vec_embedding (numpy.ndarray): Pretrained word embedding matrix.
        hparam (object): Global hyper-parameters.
    """

    def __init__(
        self,
        hparams,
        iterator_creator,
        seed=None,
    ):
        """Initialization steps for NRMS.
        Compared with the BaseModel, NRMS need word embedding.
        After creating word embedding matrix, BaseModel's __init__ method will be called.
        Args:
            hparams (object): Global hyper-parameters. Some key setttings such as head_num and head_dim are there.
            iterator_creator_train (object): NRMS data loader class for train data.
            iterator_creator_test (object): NRMS data loader class for test and validation data
        """
        self.word2vec_embedding = self._init_embedding(hparams.wordEmb_file)
        super().__init__(
            hparams,
            iterator_creator,
            seed=seed,
        )

    def _get_input_label_from_iter(self, batch_data):
        """get input and labels for trainning from iterator
        Args:
            batch data: input batch data from iterator
        Returns:
            list: input feature fed into model (clicked_title_batch & candidate_title_batch)
            numpy.ndarray: labels
        """
        input_feat = [
            batch_data["clicked_title_batch"],
            batch_data["clicked_abstract_batch"],
            batch_data["candidate_title_batch"],
            batch_data["candidate_abstract_batch"],
        ]
        input_label = batch_data["labels"]
        return input_feat, input_label

    def _get_user_feature_from_iter(self, batch_data):
        """get input of user encoder
        Args:
            batch_data: input batch data from user iterator
        Returns:
            numpy.ndarray: input user feature (clicked title batch)
        """
        # print("\nclicked_title_batch:", batch_data["clicked_title_batch"].shape)
        # print("\nclicked_abstract_batch:", batch_data["clicked_abstract_batch"].shape)
        input_feature = [
            batch_data["clicked_title_batch"],
            batch_data["clicked_abstract_batch"],
        ]
        input_feature = np.concatenate(input_feature, axis=-1)
        # print("\ninput_feature:", input_feature.shape)
        return input_feature

    def _get_news_feature_from_iter(self, batch_data):
        """get input of news encoder
        Args:
            batch_data: input batch data from news iterator
        Returns:
            numpy.ndarray: input news feature (candidate title batch)
        """
        input_feature = [
            batch_data["candidate_title_batch"],
            batch_data["candidate_abstract_batch"],
        ]
        input_feature = np.concatenate(input_feature, axis=-1)
        return input_feature

    def _build_graph(self):
        """Build NRMS model and scorer.
        Returns:
            object: a model used to train.
            object: a model used to evaluate and inference.
        """
        hparams = self.hparams
        model, scorer = self._build_nrms()
        return model, scorer

    def _build_userencoder(self, newsencoder):
        """The main function to create user encoder of NRMS.
        Args:
            newsencoder (object): the news encoder of NRMS.
        Return:
            object: the user encoder of NRMS.
        """
        hparams = self.hparams
        his_input_title_abstract = keras.Input(
            shape=(hparams.his_size, hparams.title_size + hparams.body_size),
            dtype="int32",
        )
        click_new_presents = layers.TimeDistributed(newsencoder)(
            his_input_title_abstract
        )
        y = SelfAttention(hparams.head_num, hparams.head_dim, seed=self.seed)(
            [click_new_presents] * 3
        )
        user_present = AttLayer2(hparams.attention_hidden_dim, seed=self.seed)(y)
        model = keras.Model(his_input_title_abstract, user_present, name="user_encoder")
        print("\n", model.summary())
        return model

    def _build_newsencoder(self, embedding_layer):
        """The main function to create news encoder of NRMS.
        Args:
            embedding_layer (object): a word embedding layer.
        Return:
            object: the news encoder of NRMS
        """
        hparams = self.hparams
        input_title_abstract = tf.keras.Input(
            shape=(hparams.title_size + hparams.body_size,), dtype="int32"
        )
        # print("\ninput_title_abstract:", input_title_abstract)
        sequences_input_title = layers.Lambda(lambda x: x[:, : hparams.title_size])(
            input_title_abstract
        )
        # print("\nsequences_input_title:", sequences_input_title)
        sequences_input_abstract = layers.Lambda(
            lambda x: x[:, hparams.title_size : hparams.title_size + hparams.body_size]
        )(input_title_abstract)
        title_rep = self._build_titleencoder(embedding_layer)(sequences_input_title)
        abstract_rep = self._build_abstractencoder(embedding_layer)(
            sequences_input_abstract
        )
        # print("title_rep:", title_rep)
        # print("abstract_rep:", abstract_rep)
        concate_rep = layers.Concatenate(axis=-2)([title_rep, abstract_rep])
        y = SelfAttention(hparams.head_num, hparams.head_dim, seed=self.seed)(
            [concate_rep, concate_rep, concate_rep]
        )
        pred_title = AttLayer2(hparams.attention_hidden_dim, seed=self.seed)(y)
        model = keras.Model(input_title_abstract, pred_title, name="news_encoder")
        print("\n", model.summary())
        return model

    def _build_titleencoder(self, embedding_layer):
        hparams = self.hparams
        sequences_input_title = keras.Input(shape=(hparams.title_size,), dtype="int32")
        embedded_sequences_title = embedding_layer(sequences_input_title)
        y = layers.Dropout(hparams.dropout)(embedded_sequences_title)
        y = SelfAttention(hparams.head_num, hparams.head_dim, seed=self.seed)([y, y, y])
        y = layers.Dropout(hparams.dropout)(y)
        pred_title = AttLayer2(hparams.attention_hidden_dim, seed=self.seed)(y)
        pred_title = layers.Reshape((1, 400))(pred_title)
        model = keras.Model(sequences_input_title, pred_title, name="title_encoder")
        print("\n", model.summary())
        return model

    def _build_abstractencoder(self, embedding_layer):
        hparams = self.hparams
        sequences_input_abstract = keras.Input(
            shape=(hparams.body_size,), dtype="int32"
        )
        embedded_sequences_abstract = embedding_layer(sequences_input_abstract)
        y = layers.Dropout(hparams.dropout)(embedded_sequences_abstract)
        y = SelfAttention(hparams.head_num, hparams.head_dim, seed=self.seed)([y, y, y])
        y = layers.Dropout(hparams.dropout)(y)
        pred_abstract = AttLayer2(hparams.attention_hidden_dim, seed=self.seed)(y)
        pred_abstract = layers.Reshape((1, 400))(pred_abstract)
        model = keras.Model(
            sequences_input_abstract, pred_abstract, name="abstract_encoder"
        )
        print("\n", model.summary())
        return model

    def _build_nrms(self):
        """The main function to create NRMS's logic. The core of NRMS
        is a user encoder and a news encoder.
        Returns:
            object: a model used to train.
            object: a model used to evaluate and inference.
        """
        hparams = self.hparams
        his_input_title = keras.Input(
            shape=(hparams.his_size, hparams.title_size), dtype="int32"
        )
        his_input_abstract = keras.Input(
            shape=(hparams.his_size, hparams.body_size), dtype="int32"
        )
        pred_input_title = keras.Input(
            shape=(hparams.npratio + 1, hparams.title_size), dtype="int32"
        )
        pred_input_abstract = keras.Input(
            shape=(hparams.npratio + 1, hparams.body_size), dtype="int32"
        )
        pred_input_title_one = keras.Input(
            shape=(
                1,
                hparams.title_size,
            ),
            dtype="int32",
        )
        pred_input_abstract_one = keras.Input(
            shape=(
                1,
                hparams.body_size,
            ),
            dtype="int32",
        )
        his_title_abstract = layers.Concatenate(axis=-1)(
            [his_input_title, his_input_abstract]
        )
        pred_title_abstract = layers.Concatenate(axis=-1)(
            [pred_input_title, pred_input_abstract]
        )
        pred_title_abstract_one = layers.Concatenate(axis=-1)(
            [
                pred_input_title_one,
                pred_input_abstract_one,
            ]
        )
        pred_title_abstract_one = layers.Reshape((-1,))(pred_title_abstract_one)
        embedding_layer = layers.Embedding(
            self.word2vec_embedding.shape[0],
            hparams.word_emb_dim,
            weights=[self.word2vec_embedding],
            trainable=True,
        )
        self.newsencoder = self._build_newsencoder(embedding_layer)
        self.userencoder = self._build_userencoder(self.newsencoder)
        user_present = self.userencoder(his_title_abstract)
        # print("\npred_title_abstract:", pred_title_abstract)
        news_present = layers.TimeDistributed(self.newsencoder)(pred_title_abstract)
        news_present_one = self.newsencoder(pred_title_abstract_one)
        preds = layers.Dot(axes=-1)([news_present, user_present])
        preds = layers.Activation(activation="softmax")(preds)
        pred_one = layers.Dot(axes=-1)([news_present_one, user_present])
        pred_one = layers.Activation(activation="sigmoid")(pred_one)
        model = keras.Model(
            [
                his_input_title,
                his_input_abstract,
                pred_input_title,
                pred_input_abstract,
            ],
            preds,
        )
        scorer = keras.Model(
            [
                his_input_title,
                his_input_abstract,
                pred_input_title_one,
                pred_input_abstract_one,
            ],
            pred_one,
        )
        print("\n", model.summary())
        return model, scorer


data_path = "/kaggle/input/mindlike-final/"
train_news_file = os.path.join(data_path, r"news_train.tsv")
train_behaviors_file = os.path.join(data_path, r"behaviors_train.tsv")
valid_news_file = os.path.join(data_path, r"news_val.tsv")
valid_behaviors_file = os.path.join(data_path, r"behaviors_val.tsv")
# test_news_file = os.path.join(data_path, r'news_test.tsv')
# est_behaviors_file = os.path.join(data_path, r'behaviors_test.tsv')
iterator = MINDIterator
with tpu_strategy.scope():
    model = NRMSModel(hparams, iterator)
# print(model.run_eval(train_news_file, train_behaviors_file))
# print("\n", model.run_eval(valid_news_file, valid_behaviors_file))
# %%time
model.fit(train_news_file, train_behaviors_file, valid_news_file, valid_behaviors_file)
model.model.save_weights("/kaggle/working/models/nrms_ckpt")


# adsasd
# model.model.save("/kaggle/working/model.ckpt")
#!cd /kaggle/working
#!zip -r file1.zip /kaggle/working/model.ckpt
def init_behaviors2(iterator, behaviors_file):
    """init behavior logs given behaviors file.
    Args:
    behaviors_file: path of behaviors file
    """
    iterator.histories = []
    iterator.imprs = []
    iterator.labels = []
    iterator.impr_indexes = []
    iterator.uindexes = []
    with tf.io.gfile.GFile(behaviors_file, "r") as rd:
        impr_index = 0
        for line in rd:
            uid, time, history, impr = line.strip("\n").split(iterator.col_spliter)[-4:]
            history = [iterator.nid2index[i] for i in history.split()]
            history = [0] * (iterator.his_size - len(history)) + history[
                : iterator.his_size
            ]
            impr_news = [iterator.nid2index[i.split("-")[0]] for i in impr.split()]
            label = [0 for i in impr.split()]
            uindex = iterator.uid2index[uid] if uid in iterator.uid2index else 0
            iterator.histories.append(history)
            iterator.imprs.append(impr_news)
            iterator.labels.append(label)
            iterator.impr_indexes.append(impr_index)
            iterator.uindexes.append(uindex)
            impr_index += 1


def load_impression_from_file2(iterator, behaivors_file):
    """Read and parse impression data from behaivors file.
    Args:
        behaivors_file (str): A file contains several informations of behaviros.
    Yields:
        object: An iterator that yields parsed impression data, in the format of dict.
    """
    init_behaviors2(iterator, behaivors_file)
    indexes = np.arange(len(iterator.labels))
    for index in indexes:
        impr_label = np.array(iterator.labels[index], dtype="int32")
        impr_news = np.array(iterator.imprs[index], dtype="int32")
        yield (
            iterator.impr_indexes[index],
            impr_news,
            iterator.uindexes[index],
            impr_label,
        )


def load_user_from_file2(iterator, news_file, behavior_file):
    """Read and parse user data from news file and behavior file.
    Args:
        news_file (str): A file contains several informations of news.
        beahaviros_file (str): A file contains information of user impressions.
    Yields:
        object: An iterator that yields parsed user feature, in the format of dict.
    """
    iterator.init_news(news_file)
    init_behaviors2(iterator, behavior_file)
    user_indexes = []
    impr_indexes = []
    click_title_indexes = []
    click_abstract_indexes = []
    cnt = 0
    for index in range(len(iterator.impr_indexes)):
        click_title_indexes.append(iterator.news_title_index[iterator.histories[index]])
        click_abstract_indexes.append(
            iterator.news_abstract_index[iterator.histories[index]]
        )
        user_indexes.append(iterator.uindexes[index])
        impr_indexes.append(iterator.impr_indexes[index])
        cnt += 1
        if cnt >= iterator.batch_size:
            yield iterator._convert_user_data(
                user_indexes, impr_indexes, click_title_indexes, click_abstract_indexes
            )
            user_indexes = []
            impr_indexes = []
            click_title_indexes = []
            click_abstract_indexes = []
            cnt = 0
    if cnt > 0:
        yield iterator._convert_user_data(
            user_indexes, impr_indexes, click_title_indexes, click_abstract_indexes
        )


def load_news_from_file2(iterator, news_file):
    print(news_file)
    """Read and parse user data from news file.
        Args:
            news_file (str): A file contains several informations of news.
        Yields:
            object: An iterator that yields parsed news feature, in the format of dict.
        """
    # if not hasattr(iterator, "news_title_index"):
    iterator.init_news(news_file)
    news_indexes = []
    candidate_title_indexes = []
    candidate_abstract_indexes = []
    cnt = 0
    for index in range(len(iterator.news_title_index)):
        news_indexes.append(index)
        candidate_title_indexes.append(iterator.news_title_index[index])
        candidate_abstract_indexes.append(iterator.news_abstract_index[index])
        cnt += 1
        if cnt >= iterator.batch_size:
            yield iterator._convert_news_data(
                news_indexes, candidate_title_indexes, candidate_abstract_indexes
            )
            news_indexes = []
            candidate_title_indexes = []
            candidate_abstract_indexes = []
            cnt = 0
    if cnt > 0:
        yield iterator._convert_news_data(
            news_indexes, candidate_title_indexes, candidate_abstract_indexes
        )


def run_news2(model, news_filename):
    # print(news_filename)
    if not hasattr(model, "newsencoder"):
        raise ValueError("model must have attribute newsencoder")
    news_indexes = []
    news_vecs = []
    for batch_data_input in tqdm(
        load_news_from_file2(model.test_iterator, news_filename)
    ):
        news_index, news_vec = model.news(batch_data_input)
        news_indexes.extend(np.reshape(news_index, -1))
        news_vecs.extend(news_vec)
    return dict(zip(news_indexes, news_vecs))


def run_user2(model, news_filename, behaviors_file):
    if not hasattr(model, "userencoder"):
        raise ValueError("model must have attribute userencoder")
    user_indexes = []
    user_vecs = []
    for batch_data_input in tqdm(
        load_user_from_file2(model.test_iterator, news_filename, behaviors_file)
    ):
        user_index, user_vec = model.user(batch_data_input)
        user_indexes.extend(np.reshape(user_index, -1))
        user_vecs.extend(user_vec)
    return dict(zip(user_indexes, user_vecs))


def run_fast_eval2(model, news_filename, behaviors_file):
    news_vecs = run_news2(model, news_filename)
    # print(news_filename)
    user_vecs = run_user2(model, news_filename, behaviors_file)
    model.news_vecs = news_vecs
    model.user_vecs = user_vecs
    group_impr_indexes = []
    group_labels = []
    group_preds = []
    for (
        impr_index,
        news_index,
        user_index,
        label,
    ) in tqdm(load_impression_from_file2(model.test_iterator, behaviors_file)):
        pred = np.dot(
            np.stack([news_vecs[i] for i in news_index], axis=0),
            user_vecs[impr_index],
        )
        group_impr_indexes.append(impr_index)
        group_labels.append(label)
        group_preds.append(pred)
    return group_impr_indexes, group_labels, group_preds


group_impr_indexes, group_labels, group_preds = run_fast_eval2(
    model, test_news_file, test_behaviors_file
)
output_path = "/kaggle/working/"
with open(os.path.join(output_path, "prediction.txt"), "w") as f:
    for impr_index, preds in tqdm(zip(group_impr_indexes, group_preds)):
        impr_index += 1
        pred_rank = (np.argsort(np.argsort(preds)[::-1]) + 1).tolist()
        pred_rank = "[" + ",".join([str(i) for i in pred_rank]) + "]"
        f.write(" ".join([str(impr_index), pred_rank]) + "\n")
import zipfile

f = zipfile.ZipFile(
    os.path.join(output_path, "prediction.zip"), "w", zipfile.ZIP_DEFLATED
)
f.write(os.path.join(output_path, "prediction.txt"), arcname="prediction.txt")
f.close()
import matplotlib.pyplot as plt

fig = plt.figure()
plt.plot(np.arange(0, epochs), model.history["train_loss"], label="validation auc")
plt.plot(np.arange(0, epochs), model.history["val_loss"], label="validation auc")
plt.xticks(range(0, epochs))
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
