# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/simentities.data.dataset.ipynb (unless otherwise specified).

__all__ = ['Dataset']

# Cell
import pandas as pd

from pathlib import Path

# Cell
class Dataset:
    """Load csv file and takes required data

    Takes data dedicated to the task. Series of authors and affiliations

    """
    def __init__(self, path: Path, compression: str = None):
        self.df = pd.read_csv(path, compression=compression)
        self.names = self._process_names(col=self.df["authors"])
        self.affiliations = self._process_affiliations(self.df["affiliations"])

    def _process_names(self, col: pd.Series) -> pd.Series:
        name_col = col.dropna().apply(eval)
        name_col = name_col.explode(ignore_index=True)
        return name_col

    def _process_affiliations(self, col: pd.Series) -> pd.Series:
        aff_col = col.dropna().str.split(".")
        aff_col = aff_col.explode(ignore_index=True)
        return aff_col

