| import enum | |
| class DatasetName(enum.Enum): | |
| """ | |
| Supported dataset names enumeration | |
| """ | |
| IMDB = "imdb" | |
| POLARITY = "polarity" | |
| class DatasetParams: | |
| """ | |
| Abstarct class for dataset | |
| """ | |
| hugging_face_name = "" | |
| content_col_name = "" | |
| label_col_name = "" | |
| local_path = "" | |
| def get_dataset_params_by_name(dataset_name: DatasetName) -> DatasetParams: | |
| if dataset_name == DatasetName.IMDB: | |
| return ImbdParams() | |
| if dataset_name == DatasetName.POLARITY: | |
| return PolarityParams() | |
| raise ValueError(f"Unsupported dataset: {dataset_name}") | |
| class ImbdParams(DatasetParams): | |
| """ | |
| IMDB dataset params class | |
| """ | |
| hugging_face_name = "stanfordnlp/imdb" | |
| content_col_name = "text" | |
| label_col_name = "label" | |
| local_path = "imdb" | |
| class PolarityParams(DatasetParams): | |
| """ | |
| POLARITY dataset params class | |
| """ | |
| hugging_face_name = "fancyzhx/amazon_polarity" | |
| content_col_name = "content" | |
| label_col_name = "label" | |
| local_path = "polarity" | |