refactor: abstract supported models into a class
Inherited by SequenceClassifier and Processor to make code more Pythonic
@@ -34,12 +34,20 @@ MODEL_CLASS.update(
)
-def _list_supported_models():
+class SupportedModels:
+ """Generic Class to be inherited by classes that want to support PyTorch Transformers
+
+ Returns:
+ list -- list of supported PyTorch Transformers
+ """
+
+ @staticmethod
+ def list_supported_models():
return list(MODEL_CLASS)
-class Processor:
- def __init__(self, model_name, tokenize=None, to_lower=False, cache_dir="."):
+class Processor(SupportedModels):
+ def __init__(self, model_name="bert-base-cased", tokenize=None, to_lower=False, cache_dir="."):
self.tokenizer = TOKENIZER_CLASS[model_name].from_pretrained(
model_name, do_lower_case=to_lower, cache_dir=cache_dir
)
@@ -52,10 +60,6 @@ class Processor:
else:
raise ValueError("Model not supported.")
- @staticmethod
- def list_supported_models():
- return _list_supported_models()
-
def preprocess(self, text, labels, max_len):
"""preprocess data or batches"""
if max_len > MAX_SEQ_LEN:
@@ -85,7 +89,7 @@ class Processor:
return td
-class SequenceClassifier:
+class SequenceClassifier(SupportedModels):
def __init__(self, model_name="bert-base-cased", num_labels=2, cache_dir=".", seed=0):
self.model_name = model_name
self.model = MODEL_CLASS[model_name].from_pretrained(
@@ -93,10 +97,6 @@ class SequenceClassifier:
)
self.seed = seed
- @staticmethod
- def list_supported_models():
- return _list_supported_models()
-
def fit(
self,
train_dataset,
@@ -145,3 +145,4 @@ class SequenceClassifier:
def predict():
pass
+
