diff options
| author | ZJaume <jzaragoza@prompsit.com> | 2022-07-05 14:48:53 +0300 |
|---|---|---|
| committer | ZJaume <jzaragoza@prompsit.com> | 2022-07-05 14:48:53 +0300 |
| commit | cb5e911d21e87c4dfbd5b082ded178bfb0c28683 (patch) | |
| tree | 7c356c8428a5fe6f3a6a1dbbfba00646e75f255c | |
| parent | c591cc7fd78b9c4c17c3814a034bf55d58d7d490 (diff) | |
Change XLMR subclass name to pass HF signature check
| -rw-r--r-- | bicleaner_ai/models.py | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/bicleaner_ai/models.py b/bicleaner_ai/models.py index df7ebda..cdb522c 100644 --- a/bicleaner_ai/models.py +++ b/bicleaner_ai/models.py @@ -510,7 +510,7 @@ class BCXLMRoberta(BaseModel): def load_model(self, model_file): settings = self.settings - tf_model = BCXLMRobertaForSequenceClassification.from_pretrained( + tf_model = TFXLMRBicleaner.from_pretrained( model_file, num_labels=settings["n_classes"], head_hidden_size=settings["n_hidden"], @@ -615,7 +615,7 @@ class BCXLMRoberta(BaseModel): return y_true, y_pred -class BCXLMRobertaForSequenceClassification(TFXLMRobertaForSequenceClassification): +class TFXLMRBicleaner(TFXLMRobertaForSequenceClassification): """Model for sentence-level classification tasks.""" def __init__(self, config, head_hidden_size, head_dropout, head_activation): |