import itertools


def generate_class_dict(train_df, test_df):
    """
    Given the train and test dataframe, containing "concepts" columns, where every entry is a list of strings representing
    the concepts or classes we are trying to predict, return a dictionary mapping a concept to a index.
    :param train_df: Train dataframe, must contain the "concepts" column.
    :param test_df: Test dataframe, must contain the "concepts" column.
    :return:
    """
    class_dict = dict()
    # make a set of concepts by merging the sets obtained by concepts from train and test dataframes
    concepts = set(itertools.chain(*train_df["concepts"].values)) | set(itertools.chain(*test_df["concepts"].values))
    # add to dict and return
    for concept in sorted(concepts):
        class_dict[concept] = len(class_dict)
    return class_dict