Change default to smaller model

NatLibFi · Aug 28, 2024 · 7379061 · 7379061
1 parent 3e02a72
commit 7379061
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/annif/backend/xtransformer.py b/annif/backend/xtransformer.py
@@ -93,7 +93,7 @@ class XTransformerBackend(mixins.TfidfVectorizerMixin, backend.AnnifBackend):
         "imbalanced_depth": 100,
         "max_match_clusters": 32768,
         "do_fine_tune": True,
-        "model_shortcut": "bert-base-multilingual-uncased",
+        "model_shortcut": "distilbert-base-multilingual-uncased",
         "beam_size": 20,
         "limit": 100,
         "post_processor": "sigmoid",

diff --git a/tests/test_backend_xtransformer.py b/tests/test_backend_xtransformer.py
@@ -41,8 +41,8 @@ def test_xtransformer_default_params(project):
         "imbalanced_depth": 100,
         "max_match_clusters": 32768,
         "do_fine_tune": True,
-        # 'model_shortcut': 'distilbert-base-multilingual-cased',
-        "model_shortcut": "bert-base-multilingual-uncased",
+        'model_shortcut': 'distilbert-base-multilingual-uncased',
+        # "model_shortcut": "bert-base-multilingual-uncased",
         "beam_size": 20,
         "limit": 100,
         "post_processor": "sigmoid",
@@ -144,8 +144,8 @@ def test_xtransformer_train(datadir, document_corpus, project, mocked_xtransform
                 "max_leaf_size": 100,
                 "imbalanced_ratio": 0.0,
                 "imbalanced_depth": 100,
-                "model_shortcut": "bert-base-multilingual-uncased",
-                # 'model_shortcut': 'distilbert-base-multilingual-cased',
+                # "model_shortcut": "bert-base-multilingual-uncased",
+                'model_shortcut': 'distilbert-base-multilingual-uncased',
                 "post_processor": "sigmoid",
                 "negative_sampling": "tfn",
                 "ensemble_method": "transformer-only",