From 2042779efa3aa2ad48bcefad34076cc57c0b0620 Mon Sep 17 00:00:00 2001 From: Lorenzo Stella Date: Fri, 5 Apr 2024 17:15:33 +0200 Subject: [PATCH] Simplify tokenizer creation (#44) *Description of changes:* Minor simplification to how the tokenizer is constructed from the config By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. --- src/chronos/chronos.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/chronos/chronos.py b/src/chronos/chronos.py index 175b5db..cc1d3fb 100644 --- a/src/chronos/chronos.py +++ b/src/chronos/chronos.py @@ -5,6 +5,7 @@ import warnings from dataclasses import dataclass from typing import Any, Dict, List, Literal, Optional, Tuple, Union +import chronos import torch import torch.nn as nn from transformers import ( @@ -45,9 +46,8 @@ class ChronosConfig: ), f"Special token id's must be smaller than {self.n_special_tokens=}" def create_tokenizer(self) -> "ChronosTokenizer": - if self.tokenizer_class == "MeanScaleUniformBins": - return MeanScaleUniformBins(**self.tokenizer_kwargs, config=self) - raise ValueError + class_ = getattr(chronos, self.tokenizer_class) + return class_(**self.tokenizer_kwargs, config=self) class ChronosTokenizer: