Repository: ScandEval/ScandEval (language: python)

Issue:
enable mistral
mistral needs flash attention and a decent dtype to be finetuned on a 40GB GPU


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 3651e5e1e..0a64e3c2f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,5 @@
 # Byte-compiled / optimized / DLL files
-__pycache__/
+**/__pycache__/
 *.py[cod]
 
 # C extensions
@@ -89,6 +89,9 @@ target/
 # pytest cache
 .pytest_cache/
 
+# Ruff cache
+.ruff_cache/
+
 # Checkpoints
 checkpoint-*
 
@@ -103,6 +106,7 @@ docs/scandeval/
 docs/index.html
 docs/scandeval.html
 docs/search.json
+docs/search.js
 
 # Experiments
 gai-experiments.xlsx
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 56e1bb95c..ff6c792e3 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -4,7 +4,7 @@ repos:
     hooks:
     - id: black
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.0.290
+    rev: v0.1.6
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
@@ -14,7 +14,7 @@ repos:
     hooks:
     -   id: nbstripout
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.4.1
+    rev: v1.7.1
     hooks:
     -   id: mypy
         args: [--install-types, --non-interactive, --ignore-missing-imports, --show-error-codes]
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 608cf6275..90eb3ddd2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,12 @@ and this project adheres to
     - `rrn`, an Icelandic summarization dataset based on news articles.
     - `mlsum`, a German summarization dataset based on news articles.
     - `wiki-lingua-nl`, a Dutch summarization dataset based on WikiHow articles.
+  These are all of the task `summarization`, meaning that they can also all be run
+  using `scandeval --dataset-task summarization --model-id <model_id>`.
+- A `--use-flash-attention` flag has been added, which enables Flash Attention 2.0,
+  which is required by some models, such as Mistral-based ones. If `flash-attn` has not
+  been installed then an informative error message will be raised. Thanks to @peter-sk
+  for this contribution!
 
 ### Changed
 - Now uses 8-bit AdamW whenever CUDA is available, as opposed to regular AdamW.
@@ -29,6 +35,8 @@ and this project adheres to
 ### Fixed
 - A bug was removed which caused some overlap between the dataset splits of the
   ScandiQA datasets.
+- Now allows loading in models in the data type that they were trained in, which
+  previously caused errors if they weren't trained in float32.
 
 
 ## [v8.0.0] - 2023-11-29
diff --git a/README.md b/README.md
index b530622e1..07efd3b08 100644
--- a/README.md
+++ b/README.md
@@ -1,16 +1,22 @@
 <div align='center'>
-<img src="https://raw.githubusercontent.com/saattrupdan/ScandEval/main/gfx/scandeval.png" width="517" height="217">
+<img src="https://raw.githubusercontent.com/ScandEval/ScandEval/main/gfx/scandeval.png" width="517" height="217">
 </div>
 
-### Evaluation of pretrained language models on mono- or multilingual Scandinavian language tasks.
+### Evaluation of pretrained language models on mono- or multilingual language tasks.
 
 ______________________________________________________________________
 [![PyPI Status](https://badge.fury.io/py/scandeval.svg)](https://pypi.org/project/scandeval/)
 [![Paper](https://img.shields.io/badge/arXiv-2304.00906-b31b1b.svg)](https://arxiv.org/abs/2304.00906)
-[![License](https://img.shields.io/github/license/saattrupdan/ScandEval)](https://github.com/saattrupdan/ScandEval/blob/main/LICENSE)
-[![LastCommit](https://img.shields.io/github/last-commit/saattrupdan/ScandEval)](https://github.com/saattrupdan/ScandEval/commits/main)
-[![Code Coverage](https://img.shields.io/badge/Coverage-64%25-yellow.svg)](https://github.com/saattrupdan/ScandEval/tree/main/tests)
-[![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-2.0-4baaaa.svg)](https://github.com/saattrupdan/ScandEval/blob/main/CODE_OF_CONDUCT.md)
+[![License](https://img.shields.io/github/license/ScandEval/ScandEval)](https://github.com/ScandEval/ScandEval/blob/main/LICENSE)
+[![LastCommit](https://img.shields.io/github/last-commit/ScandEval/ScandEval)](https://github.com/ScandEval/ScandEval/commits/main)
+[![Code Coverage](https://img.shields.io/badge/Coverage-64%25-yellow.svg)](https://github.com/ScandEval/ScandEval/tree/main/tests)
+[![Contributor Covenant](https://img.shields.io/badge/Contributor%20Covenant-2.0-4baaaa.svg)](https://github.com/ScandEval/ScandEval/blob/main/CODE_OF_CONDUCT.md)
+
+
+## Maintainers
+
+- Dan Saattrup Nielsen (@saattrupdan, dan.nielsen@alexandra.dk)
+- Kenneth Enevoldsen (@KennethEnevoldsen, kenneth.enevoldsen@cas.au.dk)
 
 
 ## Installation
@@ -107,7 +113,6 @@ World](https://satwcomic.com/) team. Go check them out!
 ## Project structure
 ```
 .
-├── .flake8
 ├── .github
 │   └── workflows
 │       └── ci.yaml
@@ -116,9 +121,11 @@ World](https://satwcomic.com/) team. Go check them out!
 ├── CHANGELOG.md
 ├── LICENSE
 ├── README.md
+├── docs
 ├── gfx
 │   └── scandeval.png
 ├── makefile
+├── poetry.lock
 ├── poetry.toml
 ├── pyproject.toml
 ├── src
@@ -133,28 +140,47 @@ World](https://satwcomic.com/) team. Go check them out!
 │   │   ├── dataset_configs.py
 │   │   ├── dataset_factory.py
 │   │   ├── dataset_tasks.py
+│   │   ├── enums.py
 │   │   ├── exceptions.py
-│   │   ├── hf_hub.py
+│   │   ├── finetuning.py
+│   │   ├── generation.py
 │   │   ├── languages.py
+│   │   ├── model_config.py
 │   │   ├── model_loading.py
+│   │   ├── model_setups
+│   │   │   ├── __init__.py
+│   │   │   ├── fresh.py
+│   │   │   ├── hf.py
+│   │   │   ├── local.py
+│   │   │   ├── openai.py
+│   │   │   └── utils.py
 │   │   ├── named_entity_recognition.py
+│   │   ├── openai_models.py
+│   │   ├── protocols.py
 │   │   ├── question_answering.py
 │   │   ├── question_answering_trainer.py
 │   │   ├── scores.py
 │   │   ├── sequence_classification.py
 │   │   ├── speed_benchmark.py
+│   │   ├── text_to_text.py
 │   │   ├── types.py
 │   │   └── utils.py
 │   └── scripts
 │       ├── create_angry_tweets.py
 │       ├── create_dane.py
 │       ├── create_mim_gold_ner.py
+│       ├── create_mlsum.py
+│       ├── create_no_sammendrag.py
+│       ├── create_nordjylland_news.py
 │       ├── create_norec.py
 │       ├── create_norne.py
+│       ├── create_rrn.py
 │       ├── create_scala.py
 │       ├── create_scandiqa.py
 │       ├── create_suc3.py
+│       ├── create_swedn.py
 │       ├── create_swerec.py
+│       ├── create_wiki_lingua_nl.py
 │       ├── create_wikiann_fo.py
 │       ├── fill_in_missing_model_metadata.py
 │       ├── fix_dot_env_file.py
@@ -172,11 +198,13 @@ World](https://satwcomic.com/) team. Go check them out!
     ├── test_dataset_configs.py
     ├── test_dataset_factory.py
     ├── test_dataset_tasks.py
+    ├── test_enums.py
     ├── test_exceptions.py
-    ├── test_hf_hub.py
     ├── test_languages.py
+    ├── test_model_config.py
     ├── test_model_loading.py
     ├── test_named_entity_recognition.py
+    ├── test_openai_models.py
     ├── test_question_answering.py
     ├── test_question_answering_trainer.py
     ├── test_scores.py
diff --git a/poetry.lock b/poetry.lock
index 906768212..08e5f3932 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -13,13 +13,13 @@ files = [
 
 [[package]]
 name = "accelerate"
-version = "0.24.1"
+version = "0.25.0"
 description = "Accelerate"
 optional = false
 python-versions = ">=3.8.0"
 files = [
-    {file = "accelerate-0.24.1-py3-none-any.whl", hash = "sha256:866dec394da60e8da964be212379d8cf6cc0d0e5e28a7c0d7e09507715d21c61"},
-    {file = "accelerate-0.24.1.tar.gz", hash = "sha256:85ab2aeb4d06194b75113339f81b7d650523414a82c9e91b2912a655f53dfa8e"},
+    {file = "accelerate-0.25.0-py3-none-any.whl", hash = "sha256:c7bb817eb974bba0ff3ea1ba0f24d55afb86d50e3d4fe98d6922dc69cf2ccff1"},
+    {file = "accelerate-0.25.0.tar.gz", hash = "sha256:ecf55b0ab278a1dac8539dde0d276977aff04683f07ede73eaf02478538576a1"},
 ]
 
 [package.dependencies]
@@ -28,6 +28,7 @@ numpy = ">=1.17"
 packaging = ">=20.0"
 psutil = "*"
 pyyaml = "*"
+safetensors = ">=0.3.1"
 torch = ">=1.10.0"
 
 [package.extras]
@@ -37,7 +38,7 @@ rich = ["rich"]
 sagemaker = ["sagemaker"]
 test-dev = ["bitsandbytes", "datasets", "deepspeed", "evaluate", "scikit-learn", "scipy", "timm", "tqdm", "transformers"]
 test-prod = ["parameterized", "pytest", "pytest-subtests", "pytest-xdist"]
-test-trackers = ["comet-ml", "tensorboard", "wandb"]
+test-trackers = ["comet-ml", "dvclive", "tensorboard", "wandb"]
 testing = ["bitsandbytes", "datasets", "deepspeed", "evaluate", "parameterized", "pytest", "pytest-subtests", "pytest-xdist", "scikit-learn", "scipy", "timm", "tqdm", "transformers"]
 
 [[package]]
@@ -806,53 +807,53 @@ testing = ["black[jupyter] (==23.7.0)", "clu", "clu (<=0.0.9)", "einops", "gymna
 
 [[package]]
 name = "fonttools"
-version = "4.45.1"
+version = "4.46.0"
 description = "Tools to manipulate font files"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "fonttools-4.45.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:45fa321c458ea29224067700954ec44493ae869b47e7c5485a350a149a19fb53"},
-    {file = "fonttools-4.45.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0dc7617d96b1e668eea9250e1c1fe62d0c78c3f69573ce7e3332cc40e6d84356"},
-    {file = "fonttools-4.45.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:03ed3bda541e86725f6b4e1b94213f13ed1ae51a5a1f167028534cedea38c010"},
-    {file = "fonttools-4.45.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c4f4a5870e3b56788fb196da8cf30d0dfd51a76dc3b907861d018165f76ae4c2"},
-    {file = "fonttools-4.45.1-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:a3c11d9687479f01eddef729aa737abcdea0a44fdaffb62a930a18892f186c9b"},
-    {file = "fonttools-4.45.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:316cec50581e844c3ab69d7c82455b54c7cf18236b2f09e722faf665fbfcac58"},
-    {file = "fonttools-4.45.1-cp310-cp310-win32.whl", hash = "sha256:e2277cba9f0b525e30de2a9ad3cb4219aa4bc697230c1645666b0deee9f914f0"},
-    {file = "fonttools-4.45.1-cp310-cp310-win_amd64.whl", hash = "sha256:1b9e9ad2bcded9a1431afaa57c8d3c39143ac1f050862d66bddd863c515464a2"},
-    {file = "fonttools-4.45.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ff6a698bdd435d24c379f6e8a54908cd9bb7dda23719084d56bf8c87709bf3bd"},
-    {file = "fonttools-4.45.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:2c980d60cd6ec1376206fe55013d166e5627ad0b149b5c81e74eaa913ab6134f"},
-    {file = "fonttools-4.45.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a12dee6523c02ca78aeedd0a5e12bfa9b7b29896350edd5241542897b072ae23"},
-    {file = "fonttools-4.45.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:37cd1ced6efb3dd6fe82e9f9bf92fd74ac58a5aefc284045f59ecd517a5fb9ab"},
-    {file = "fonttools-4.45.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:e3d24248221bd7151dfff0d88b1b5da02dccd7134bd576ce8888199827bbaa19"},
-    {file = "fonttools-4.45.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:ba6c23591427844dfb0a13658f1718489de75de6a46b64234584c0d17573162d"},
-    {file = "fonttools-4.45.1-cp311-cp311-win32.whl", hash = "sha256:cebcddbe9351b67166292b4f71ffdbfcce01ba4b07d4267824eb46b277aeb19a"},
-    {file = "fonttools-4.45.1-cp311-cp311-win_amd64.whl", hash = "sha256:f22eb69996a0bd49f76bdefb30be54ce8dbb89a0d1246874d610f05c2aa2e69e"},
-    {file = "fonttools-4.45.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:794de93e83297db7b4943f2431e206d8b1ea69cb3ae14638a49cc50332bf0db8"},
-    {file = "fonttools-4.45.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:4ba17822a6681d06849078daaf6e03eccc9f467efe7c4c60280e28a78e8e5df9"},
-    {file = "fonttools-4.45.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e50f794d09df0675da8d9dbd7c66bfcab2f74a708343aabcad41936d26556891"},
-    {file = "fonttools-4.45.1-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8b07b857d4f9de3199a8c3d1b1bf2078c0f37447891ca1a8d9234106b9a27aff"},
-    {file = "fonttools-4.45.1-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:777ba42b94a27bb7fb2b4082522fccfd345667c32a56011e1c3e105979af5b79"},
-    {file = "fonttools-4.45.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:21e96b99878348c74aa58059b8578d7586f9519cbcdadacf56486737038aa043"},
-    {file = "fonttools-4.45.1-cp312-cp312-win32.whl", hash = "sha256:5cbf02cda8465b69769d07385f5d11e7bba19954e7787792f46fe679ec755ebb"},
-    {file = "fonttools-4.45.1-cp312-cp312-win_amd64.whl", hash = "sha256:800e354e0c3afaeb8d9552769773d02f228e98c37b8cb03041157c3d0687cffc"},
-    {file = "fonttools-4.45.1-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:6eb2c54f7a07c92108daabcf02caf31df97825738db02a28270633946bcda4d0"},
-    {file = "fonttools-4.45.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:43a3d267334109ff849c37cf3629476b5feb392ef1d2e464a167b83de8cd599c"},
-    {file = "fonttools-4.45.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8e1aefc2bf3c43e0f33f995f828a7bbeff4adc9393a7760b11456dbcf14388f6"},
-    {file = "fonttools-4.45.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1f53a19dcdd5737440839b8394eeebb35da9ec8109f7926cb6456639b5b58e47"},
-    {file = "fonttools-4.45.1-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:5a17706b9cc24b27721613fe5773d93331ab7f0ecaca9955aead89c6b843d3a7"},
-    {file = "fonttools-4.45.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:fb36e5f40191274a95938b40c0a1fa7f895e36935aea8709e1d6deff0b2d0d4f"},
-    {file = "fonttools-4.45.1-cp38-cp38-win32.whl", hash = "sha256:46eabddec12066829b8a1efe45ae552ba2f1796981ecf538d5f68284c354c589"},
-    {file = "fonttools-4.45.1-cp38-cp38-win_amd64.whl", hash = "sha256:b6de2f0fcd3302fb82f94801002cb473959e998c14c24ec28234adb674aed345"},
-    {file = "fonttools-4.45.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:392d0e3cc23daee910193625f7cf1b387aff9dd5b6f1a5f4a925680acb6dcbc2"},
-    {file = "fonttools-4.45.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:4b9544b1346d99848ac0e9b05b5d45ee703d7562fc4c9c48cf4b781de9632e57"},
-    {file = "fonttools-4.45.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8717db3e4895e4820ade64ea379187738827ee60748223cb0438ef044ee208c6"},
-    {file = "fonttools-4.45.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e29d5f298d616a93a4c5963682dc6cc8cc09f6d89cad2c29019fc5fb3b4d9472"},
-    {file = "fonttools-4.45.1-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:cb472905da3049960e80fc1cf808231880d79727a8410e156bf3e5063a1c574f"},
-    {file = "fonttools-4.45.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:ba299f1fbaa2a1e33210aaaf6fa816d4059e4d3cfe2ae9871368d4ab548c1c6a"},
-    {file = "fonttools-4.45.1-cp39-cp39-win32.whl", hash = "sha256:105099968b58a5b4cef6f3eb409db8ea8578b302a9d05e23fecba1b8b0177b5f"},
-    {file = "fonttools-4.45.1-cp39-cp39-win_amd64.whl", hash = "sha256:847f3f49dd3423e5a678c098e2ba92c7f4955d4aab3044f6a507b0bb0ecb07e0"},
-    {file = "fonttools-4.45.1-py3-none-any.whl", hash = "sha256:3bdd7dfca8f6c9f4779384064027e8477ad6a037d6a327b09381f43e0247c6f3"},
-    {file = "fonttools-4.45.1.tar.gz", hash = "sha256:6e441286d55fe7ec7c4fb36812bf914924813776ff514b744b510680fc2733f2"},
+    {file = "fonttools-4.46.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:d4e69e2c7f93b695d2e6f18f709d501d945f65c1d237dafaabdd23cd935a5276"},
+    {file = "fonttools-4.46.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:25852f0c63df0af022f698464a4a80f7d1d5bd974bcd22f995f6b4ad198e32dd"},
+    {file = "fonttools-4.46.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:adab73618d0a328b203a0e242b3eba60a2b5662d9cb2bd16ed9c52af8a7d86af"},
+    {file = "fonttools-4.46.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2cf923a4a556ab4cc4c52f69a4a2db624cf5a2cf360394368b40c5152fe3321e"},
+    {file = "fonttools-4.46.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:87c214197712cc14fd2a4621efce2a9c501a77041232b789568149a8a3161517"},
+    {file = "fonttools-4.46.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:156ae342a1ed1fe38e180de471e98fbf5b2b6ae280fa3323138569c4ca215844"},
+    {file = "fonttools-4.46.0-cp310-cp310-win32.whl", hash = "sha256:c506e3d3a9e898caee4dc094f34b49c5566870d5a2d1ca2125f0a9f35ecc2205"},
+    {file = "fonttools-4.46.0-cp310-cp310-win_amd64.whl", hash = "sha256:f8bc3973ed58893c4107993e0a7ae34901cb572b5e798249cbef35d30801ffd4"},
+    {file = "fonttools-4.46.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:982f69855ac258260f51048d9e0c53c5f19881138cc7ca06deb38dc4b97404b6"},
+    {file = "fonttools-4.46.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:2c23c59d321d62588620f2255cf951270bf637d88070f38ed8b5e5558775b86c"},
+    {file = "fonttools-4.46.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0e94244ec24a940ecfbe5b31c975c8a575d5ed2d80f9a280ce3b21fa5dc9c34"},
+    {file = "fonttools-4.46.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1a9f9cdd7ef63d1b8ac90db335762451452426b3207abd79f60da510cea62da5"},
+    {file = "fonttools-4.46.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:ca9eceebe70035b057ce549e2054cad73e95cac3fe91a9d827253d1c14618204"},
+    {file = "fonttools-4.46.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:8be6adfa4e15977075278dd0a0bae74dec59be7b969b5ceed93fb86af52aa5be"},
+    {file = "fonttools-4.46.0-cp311-cp311-win32.whl", hash = "sha256:7b5636f5706d49f13b6d610fe54ee662336cdf56b5a6f6683c0b803e23d826d2"},
+    {file = "fonttools-4.46.0-cp311-cp311-win_amd64.whl", hash = "sha256:49ea0983e55fd7586a809787cd4644a7ae471e53ab8ddc016f9093b400e32646"},
+    {file = "fonttools-4.46.0-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:7b460720ce81773da1a3e7cc964c48e1e11942b280619582a897fa0117b56a62"},
+    {file = "fonttools-4.46.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:8bee9f4fc8c99824a424ae45c789ee8c67cb84f8e747afa7f83b7d3cef439c3b"},
+    {file = "fonttools-4.46.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d3d7b96aba96e05e8c911ce2dfc5acc6a178b8f44f6aa69371ab91aa587563da"},
+    {file = "fonttools-4.46.0-cp312-cp312-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e6aeb5c340416d11a3209d75c48d13e72deea9e1517837dd1522c1fd1f17c11"},
+    {file = "fonttools-4.46.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:c779f8701deedf41908f287aeb775b8a6f59875ad1002b98ac6034ae4ddc1b7b"},
+    {file = "fonttools-4.46.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:ce199227ce7921eaafdd4f96536f16b232d6b580ce74ce337de544bf06cb2752"},
+    {file = "fonttools-4.46.0-cp312-cp312-win32.whl", hash = "sha256:1c9937c4dd1061afd22643389445fabda858af5e805860ec3082a4bc07c7a720"},
+    {file = "fonttools-4.46.0-cp312-cp312-win_amd64.whl", hash = "sha256:a9fa52ef8fd14d7eb3d813e1451e7ace3e1eebfa9b7237d3f81fee8f3de6a114"},
+    {file = "fonttools-4.46.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:c94564b1f3b5dd87e73577610d85115b1936edcc596deaf84a31bbe70e17456b"},
+    {file = "fonttools-4.46.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:a4a50a1dfad7f7ba5ca3f99cc73bf5cdac67ceade8e4b355a877521f20ad1b63"},
+    {file = "fonttools-4.46.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:89c2c520f9492844ecd6316d20c6c7a157b5c0cb73a1411b3db28ee304f30122"},
+    {file = "fonttools-4.46.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e5b7905fd68eacb7cc56a13139da5c312c45baae6950dd00b02563c54508a041"},
+    {file = "fonttools-4.46.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:8485cc468288e213f31afdaf1fdda3c79010f542559fbba936a54f4644df2570"},
+    {file = "fonttools-4.46.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:87c3299da7da55394fb324349db0ede38114a46aafd0e7dfcabfecd28cdd94c3"},
+    {file = "fonttools-4.46.0-cp38-cp38-win32.whl", hash = "sha256:f5f1423a504ccc329efb5aa79738de83d38c072be5308788dde6bd419969d7f5"},
+    {file = "fonttools-4.46.0-cp38-cp38-win_amd64.whl", hash = "sha256:6d4a4ebcc76e30898ff3296ea786491c70e183f738319ae2629e0d44f17ece42"},
+    {file = "fonttools-4.46.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:c9a0e422ab79e5cb2b47913be6a4b5fd20c4c7ac34a24f3691a4e099e965e0b8"},
+    {file = "fonttools-4.46.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:13ac0cba2fc63fa4b232f2a7971f35f35c6eaf10bd1271fa96d4ce6253a8acfd"},
+    {file = "fonttools-4.46.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:795150d5edc595e1a2cfb3d65e8f4f3d027704fc2579f8990d381bef6b188eb6"},
+    {file = "fonttools-4.46.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d00fc63131dcac6b25f50a5a129758438317e54e3ce5587163f7058de4b0e933"},
+    {file = "fonttools-4.46.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:3033b55f401a622de2630b3982234d97219d89b058607b87927eccb0f922313c"},
+    {file = "fonttools-4.46.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:e26e7fb908ae4f622813e7cb32cd2db6c24e3122bb3b98f25e832a2fe0e7e228"},
+    {file = "fonttools-4.46.0-cp39-cp39-win32.whl", hash = "sha256:2d0eba685938c603f2f648dfc0aadbf8c6a4fe1c7ca608c2970a6ef39e00f254"},
+    {file = "fonttools-4.46.0-cp39-cp39-win_amd64.whl", hash = "sha256:5200b01f463d97cc2b7ff8a1e3584151f4413e98cb8419da5f17d1dbb84cc214"},
+    {file = "fonttools-4.46.0-py3-none-any.whl", hash = "sha256:5b627ed142398ea9202bd752c04311592558964d1a765fb2f78dc441a05633f4"},
+    {file = "fonttools-4.46.0.tar.gz", hash = "sha256:2ae45716c27a41807d58a9f3f59983bdc8c0a46cb259e4450ab7e196253a9853"},
 ]
 
 [package.extras]
@@ -1062,23 +1063,6 @@ files = [
     {file = "iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3"},
 ]
 
-[[package]]
-name = "isort"
-version = "5.12.0"
-description = "A Python utility / library to sort Python imports."
-optional = false
-python-versions = ">=3.8.0"
-files = [
-    {file = "isort-5.12.0-py3-none-any.whl", hash = "sha256:f84c2818376e66cf843d497486ea8fed8700b340f308f076c6fb1229dff318b6"},
-    {file = "isort-5.12.0.tar.gz", hash = "sha256:8bef7dde241278824a6d83f44a544709b065191b95b6e50894bdc722fcba0504"},
-]
-
-[package.extras]
-colors = ["colorama (>=0.4.3)"]
-pipfile-deprecated-finder = ["pip-shims (>=0.5.2)", "pipreqs", "requirementslib"]
-plugins = ["setuptools"]
-requirements-deprecated-finder = ["pip-api", "pipreqs"]
-
 [[package]]
 name = "jax"
 version = "0.4.20"
@@ -1885,6 +1869,53 @@ files = [
 [package.dependencies]
 dill = ">=0.3.7"
 
+[[package]]
+name = "mypy"
+version = "1.7.1"
+description = "Optional static typing for Python"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "mypy-1.7.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:12cce78e329838d70a204293e7b29af9faa3ab14899aec397798a4b41be7f340"},
+    {file = "mypy-1.7.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1484b8fa2c10adf4474f016e09d7a159602f3239075c7bf9f1627f5acf40ad49"},
+    {file = "mypy-1.7.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:31902408f4bf54108bbfb2e35369877c01c95adc6192958684473658c322c8a5"},
+    {file = "mypy-1.7.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:f2c2521a8e4d6d769e3234350ba7b65ff5d527137cdcde13ff4d99114b0c8e7d"},
+    {file = "mypy-1.7.1-cp310-cp310-win_amd64.whl", hash = "sha256:fcd2572dd4519e8a6642b733cd3a8cfc1ef94bafd0c1ceed9c94fe736cb65b6a"},
+    {file = "mypy-1.7.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:4b901927f16224d0d143b925ce9a4e6b3a758010673eeded9b748f250cf4e8f7"},
+    {file = "mypy-1.7.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2f7f6985d05a4e3ce8255396df363046c28bea790e40617654e91ed580ca7c51"},
+    {file = "mypy-1.7.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:944bdc21ebd620eafefc090cdf83158393ec2b1391578359776c00de00e8907a"},
+    {file = "mypy-1.7.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:9c7ac372232c928fff0645d85f273a726970c014749b924ce5710d7d89763a28"},
+    {file = "mypy-1.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:f6efc9bd72258f89a3816e3a98c09d36f079c223aa345c659622f056b760ab42"},
+    {file = "mypy-1.7.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:6dbdec441c60699288adf051f51a5d512b0d818526d1dcfff5a41f8cd8b4aaf1"},
+    {file = "mypy-1.7.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4fc3d14ee80cd22367caaaf6e014494415bf440980a3045bf5045b525680ac33"},
+    {file = "mypy-1.7.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2c6e4464ed5f01dc44dc9821caf67b60a4e5c3b04278286a85c067010653a0eb"},
+    {file = "mypy-1.7.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:d9b338c19fa2412f76e17525c1b4f2c687a55b156320acb588df79f2e6fa9fea"},
+    {file = "mypy-1.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:204e0d6de5fd2317394a4eff62065614c4892d5a4d1a7ee55b765d7a3d9e3f82"},
+    {file = "mypy-1.7.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:84860e06ba363d9c0eeabd45ac0fde4b903ad7aa4f93cd8b648385a888e23200"},
+    {file = "mypy-1.7.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:8c5091ebd294f7628eb25ea554852a52058ac81472c921150e3a61cdd68f75a7"},
+    {file = "mypy-1.7.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:40716d1f821b89838589e5b3106ebbc23636ffdef5abc31f7cd0266db936067e"},
+    {file = "mypy-1.7.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:5cf3f0c5ac72139797953bd50bc6c95ac13075e62dbfcc923571180bebb662e9"},
+    {file = "mypy-1.7.1-cp38-cp38-win_amd64.whl", hash = "sha256:78e25b2fd6cbb55ddfb8058417df193f0129cad5f4ee75d1502248e588d9e0d7"},
+    {file = "mypy-1.7.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:75c4d2a6effd015786c87774e04331b6da863fc3fc4e8adfc3b40aa55ab516fe"},
+    {file = "mypy-1.7.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:2643d145af5292ee956aa0a83c2ce1038a3bdb26e033dadeb2f7066fb0c9abce"},
+    {file = "mypy-1.7.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:75aa828610b67462ffe3057d4d8a4112105ed211596b750b53cbfe182f44777a"},
+    {file = "mypy-1.7.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:ee5d62d28b854eb61889cde4e1dbc10fbaa5560cb39780c3995f6737f7e82120"},
+    {file = "mypy-1.7.1-cp39-cp39-win_amd64.whl", hash = "sha256:72cf32ce7dd3562373f78bd751f73c96cfb441de147cc2448a92c1a308bd0ca6"},
+    {file = "mypy-1.7.1-py3-none-any.whl", hash = "sha256:f7c5d642db47376a0cc130f0de6d055056e010debdaf0707cd2b0fc7e7ef30ea"},
+    {file = "mypy-1.7.1.tar.gz", hash = "sha256:fcb6d9afb1b6208b4c712af0dafdc650f518836065df0d4fb1d800f5d6773db2"},
+]
+
+[package.dependencies]
+mypy-extensions = ">=1.0.0"
+tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""}
+typing-extensions = ">=4.1.0"
+
+[package.extras]
+dmypy = ["psutil (>=4.0)"]
+install-types = ["pip"]
+mypyc = ["setuptools (>=50)"]
+reports = ["lxml"]
+
 [[package]]
 name = "mypy-extensions"
 version = "1.0.0"
@@ -2216,13 +2247,13 @@ test = ["dm-haiku (>=0.0.3)", "dm-tree (>=0.1.7)", "flax (==0.5.3)"]
 
 [[package]]
 name = "orbax-checkpoint"
-version = "0.4.3"
+version = "0.4.4"
 description = "Orbax Checkpoint"
 optional = false
 python-versions = ">=3.9"
 files = [
-    {file = "orbax_checkpoint-0.4.3-py3-none-any.whl", hash = "sha256:2cc19961eef4169a328989291030b9132d8184b9b138892c9149395c0fbf280a"},
-    {file = "orbax_checkpoint-0.4.3.tar.gz", hash = "sha256:dabd7c3690bd522972eccceff301245e85bc3365c383123ccf8a64ccef2d5245"},
+    {file = "orbax_checkpoint-0.4.4-py3-none-any.whl", hash = "sha256:e356288d7f62b30519b20ae3c0584743d6598234e4996b4c15bbbd32c13c1f04"},
+    {file = "orbax_checkpoint-0.4.4.tar.gz", hash = "sha256:85ab96268b3f39e83809254cb3d55aa5a47c2279d7d3e725bd5f7c2da10a4de9"},
 ]
 
 [package.dependencies]
@@ -3202,15 +3233,41 @@ nltk = "*"
 numpy = "*"
 six = ">=1.14.0"
 
+[[package]]
+name = "ruff"
+version = "0.1.6"
+description = "An extremely fast Python linter and code formatter, written in Rust."
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "ruff-0.1.6-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:88b8cdf6abf98130991cbc9f6438f35f6e8d41a02622cc5ee130a02a0ed28703"},
+    {file = "ruff-0.1.6-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:5c549ed437680b6105a1299d2cd30e4964211606eeb48a0ff7a93ef70b902248"},
+    {file = "ruff-0.1.6-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1cf5f701062e294f2167e66d11b092bba7af6a057668ed618a9253e1e90cfd76"},
+    {file = "ruff-0.1.6-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:05991ee20d4ac4bb78385360c684e4b417edd971030ab12a4fbd075ff535050e"},
+    {file = "ruff-0.1.6-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:87455a0c1f739b3c069e2f4c43b66479a54dea0276dd5d4d67b091265f6fd1dc"},
+    {file = "ruff-0.1.6-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:683aa5bdda5a48cb8266fcde8eea2a6af4e5700a392c56ea5fb5f0d4bfdc0240"},
+    {file = "ruff-0.1.6-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:137852105586dcbf80c1717facb6781555c4e99f520c9c827bd414fac67ddfb6"},
+    {file = "ruff-0.1.6-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:bd98138a98d48a1c36c394fd6b84cd943ac92a08278aa8ac8c0fdefcf7138f35"},
+    {file = "ruff-0.1.6-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3a0cd909d25f227ac5c36d4e7e681577275fb74ba3b11d288aff7ec47e3ae745"},
+    {file = "ruff-0.1.6-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:e8fd1c62a47aa88a02707b5dd20c5ff20d035d634aa74826b42a1da77861b5ff"},
+    {file = "ruff-0.1.6-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:fd89b45d374935829134a082617954120d7a1470a9f0ec0e7f3ead983edc48cc"},
+    {file = "ruff-0.1.6-py3-none-musllinux_1_2_i686.whl", hash = "sha256:491262006e92f825b145cd1e52948073c56560243b55fb3b4ecb142f6f0e9543"},
+    {file = "ruff-0.1.6-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:ea284789861b8b5ca9d5443591a92a397ac183d4351882ab52f6296b4fdd5462"},
+    {file = "ruff-0.1.6-py3-none-win32.whl", hash = "sha256:1610e14750826dfc207ccbcdd7331b6bd285607d4181df9c1c6ae26646d6848a"},
+    {file = "ruff-0.1.6-py3-none-win_amd64.whl", hash = "sha256:4558b3e178145491e9bc3b2ee3c4b42f19d19384eaa5c59d10acf6e8f8b57e33"},
+    {file = "ruff-0.1.6-py3-none-win_arm64.whl", hash = "sha256:03910e81df0d8db0e30050725a5802441c2022ea3ae4fe0609b76081731accbc"},
+    {file = "ruff-0.1.6.tar.gz", hash = "sha256:1b09f29b16c6ead5ea6b097ef2764b42372aebe363722f1605ecbcd2b9207184"},
+]
+
 [[package]]
 name = "sacrebleu"
-version = "2.3.2"
+version = "2.3.3"
 description = "Hassle-free computation of shareable, comparable, and reproducible BLEU, chrF, and TER scores"
 optional = false
 python-versions = ">=3.6"
 files = [
-    {file = "sacrebleu-2.3.2-py3-none-any.whl", hash = "sha256:ee8891d790cafbc8d3fe2a2ab2672b705fd856d287160040206785ded93a6e01"},
-    {file = "sacrebleu-2.3.2.tar.gz", hash = "sha256:29ec2eaecb9bdd0cd9ddb7a811182ed178545cc26a677f81d4d3291bd20578f4"},
+    {file = "sacrebleu-2.3.3-py3-none-any.whl", hash = "sha256:215422e29c615823281d7aa1c85c2955b9a810eee7c59909c605e8e62ee63478"},
+    {file = "sacrebleu-2.3.3.tar.gz", hash = "sha256:eff9cfd2c91041f57f6681c646c9b230ae332e442640762a556605d6ef67465b"},
 ]
 
 [package.dependencies]
@@ -3222,6 +3279,7 @@ regex = "*"
 tabulate = ">=0.8.9"
 
 [package.extras]
+dev = ["lxml-stubs", "mypy", "pytest", "types-tabulate", "wheel"]
 ja = ["ipadic (>=1.0,<2.0)", "mecab-python3 (>=1.0.5,<=1.0.6)"]
 ko = ["mecab-ko (>=1.0.0,<=1.0.1)", "mecab-ko-dic (>=1.0,<2.0)"]
 
@@ -3605,13 +3663,13 @@ numpy = ">=1.16.0"
 
 [[package]]
 name = "termcolor"
-version = "2.3.0"
+version = "2.4.0"
 description = "ANSI color formatting for output in terminal"
 optional = false
-python-versions = ">=3.7"
+python-versions = ">=3.8"
 files = [
-    {file = "termcolor-2.3.0-py3-none-any.whl", hash = "sha256:3afb05607b89aed0ffe25202399ee0867ad4d3cb4180d98aaf8eefa6a5f7d475"},
-    {file = "termcolor-2.3.0.tar.gz", hash = "sha256:b5b08f68937f138fe92f6c089b99f1e2da0ae56c52b78bf7075fd95420fd9a5a"},
+    {file = "termcolor-2.4.0-py3-none-any.whl", hash = "sha256:9297c0df9c99445c2412e832e882a7884038a25617c60cea2ad69488d4040d63"},
+    {file = "termcolor-2.4.0.tar.gz", hash = "sha256:aab9e56047c8ac41ed798fa36d892a37aca6b3e9159f3e0c24bc64a9b3ac7b7a"},
 ]
 
 [package.extras]
@@ -3630,40 +3688,47 @@ files = [
 
 [[package]]
 name = "tiktoken"
-version = "0.5.1"
+version = "0.5.2"
 description = "tiktoken is a fast BPE tokeniser for use with OpenAI's models"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "tiktoken-0.5.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:2b0bae3fd56de1c0a5874fb6577667a3c75bf231a6cef599338820210c16e40a"},
-    {file = "tiktoken-0.5.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e529578d017045e2f0ed12d2e00e7e99f780f477234da4aae799ec4afca89f37"},
-    {file = "tiktoken-0.5.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:edd2ffbb789712d83fee19ab009949f998a35c51ad9f9beb39109357416344ff"},
-    {file = "tiktoken-0.5.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e4c73d47bdc1a3f1f66ffa019af0386c48effdc6e8797e5e76875f6388ff72e9"},
-    {file = "tiktoken-0.5.1-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:46b8554b9f351561b1989157c6bb54462056f3d44e43aa4e671367c5d62535fc"},
-    {file = "tiktoken-0.5.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:92ed3bbf71a175a6a4e5fbfcdb2c422bdd72d9b20407e00f435cf22a68b4ea9b"},
-    {file = "tiktoken-0.5.1-cp310-cp310-win_amd64.whl", hash = "sha256:714efb2f4a082635d9f5afe0bf7e62989b72b65ac52f004eb7ac939f506c03a4"},
-    {file = "tiktoken-0.5.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a10488d1d1a5f9c9d2b2052fdb4cf807bba545818cb1ef724a7f5d44d9f7c3d4"},
-    {file = "tiktoken-0.5.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:8079ac065572fe0e7c696dbd63e1fdc12ce4cdca9933935d038689d4732451df"},
-    {file = "tiktoken-0.5.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7ef730db4097f5b13df8d960f7fdda2744fe21d203ea2bb80c120bb58661b155"},
-    {file = "tiktoken-0.5.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:426e7def5f3f23645dada816be119fa61e587dfb4755de250e136b47a045c365"},
-    {file = "tiktoken-0.5.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:323cec0031358bc09aa965c2c5c1f9f59baf76e5b17e62dcc06d1bb9bc3a3c7c"},
-    {file = "tiktoken-0.5.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:5abd9436f02e2c8eda5cce2ff8015ce91f33e782a7423de2a1859f772928f714"},
-    {file = "tiktoken-0.5.1-cp311-cp311-win_amd64.whl", hash = "sha256:1fe99953b63aabc0c9536fbc91c3c9000d78e4755edc28cc2e10825372046a2d"},
-    {file = "tiktoken-0.5.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:dcdc630461927718b317e6f8be7707bd0fc768cee1fdc78ddaa1e93f4dc6b2b1"},
-    {file = "tiktoken-0.5.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:1f2b3b253e22322b7f53a111e1f6d7ecfa199b4f08f3efdeb0480f4033b5cdc6"},
-    {file = "tiktoken-0.5.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:43ce0199f315776dec3ea7bf86f35df86d24b6fcde1babd3e53c38f17352442f"},
-    {file = "tiktoken-0.5.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a84657c083d458593c0235926b5c993eec0b586a2508d6a2020556e5347c2f0d"},
-    {file = "tiktoken-0.5.1-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:c008375c0f3d97c36e81725308699116cd5804fdac0f9b7afc732056329d2790"},
-    {file = "tiktoken-0.5.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:779c4dea5edd1d3178734d144d32231e0b814976bec1ec09636d1003ffe4725f"},
-    {file = "tiktoken-0.5.1-cp38-cp38-win_amd64.whl", hash = "sha256:b5dcfcf9bfb798e86fbce76d40a1d5d9e3f92131aecfa3d1e5c9ea1a20f1ef1a"},
-    {file = "tiktoken-0.5.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9b180a22db0bbcc447f691ffc3cf7a580e9e0587d87379e35e58b826ebf5bc7b"},
-    {file = "tiktoken-0.5.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:2b756a65d98b7cf760617a6b68762a23ab8b6ef79922be5afdb00f5e8a9f4e76"},
-    {file = "tiktoken-0.5.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ba9873c253ca1f670e662192a0afcb72b41e0ba3e730f16c665099e12f4dac2d"},
-    {file = "tiktoken-0.5.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:74c90d2be0b4c1a2b3f7dde95cd976757817d4df080d6af0ee8d461568c2e2ad"},
-    {file = "tiktoken-0.5.1-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:709a5220891f2b56caad8327fab86281787704931ed484d9548f65598dea9ce4"},
-    {file = "tiktoken-0.5.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:5d5a187ff9c786fae6aadd49f47f019ff19e99071dc5b0fe91bfecc94d37c686"},
-    {file = "tiktoken-0.5.1-cp39-cp39-win_amd64.whl", hash = "sha256:e21840043dbe2e280e99ad41951c00eff8ee3b63daf57cd4c1508a3fd8583ea2"},
-    {file = "tiktoken-0.5.1.tar.gz", hash = "sha256:27e773564232004f4f810fd1f85236673ec3a56ed7f1206fc9ed8670ebedb97a"},
+    {file = "tiktoken-0.5.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8c4e654282ef05ec1bd06ead22141a9a1687991cef2c6a81bdd1284301abc71d"},
+    {file = "tiktoken-0.5.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:7b3134aa24319f42c27718c6967f3c1916a38a715a0fa73d33717ba121231307"},
+    {file = "tiktoken-0.5.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6092e6e77730929c8c6a51bb0d7cfdf1b72b63c4d033d6258d1f2ee81052e9e5"},
+    {file = "tiktoken-0.5.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:72ad8ae2a747622efae75837abba59be6c15a8f31b4ac3c6156bc56ec7a8e631"},
+    {file = "tiktoken-0.5.2-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:51cba7c8711afa0b885445f0637f0fcc366740798c40b981f08c5f984e02c9d1"},
+    {file = "tiktoken-0.5.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:3d8c7d2c9313f8e92e987d585ee2ba0f7c40a0de84f4805b093b634f792124f5"},
+    {file = "tiktoken-0.5.2-cp310-cp310-win_amd64.whl", hash = "sha256:692eca18c5fd8d1e0dde767f895c17686faaa102f37640e884eecb6854e7cca7"},
+    {file = "tiktoken-0.5.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:138d173abbf1ec75863ad68ca289d4da30caa3245f3c8d4bfb274c4d629a2f77"},
+    {file = "tiktoken-0.5.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:7388fdd684690973fdc450b47dfd24d7f0cbe658f58a576169baef5ae4658607"},
+    {file = "tiktoken-0.5.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a114391790113bcff670c70c24e166a841f7ea8f47ee2fe0e71e08b49d0bf2d4"},
+    {file = "tiktoken-0.5.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ca96f001e69f6859dd52926d950cfcc610480e920e576183497ab954e645e6ac"},
+    {file = "tiktoken-0.5.2-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:15fed1dd88e30dfadcdd8e53a8927f04e1f6f81ad08a5ca824858a593ab476c7"},
+    {file = "tiktoken-0.5.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:93f8e692db5756f7ea8cb0cfca34638316dcf0841fb8469de8ed7f6a015ba0b0"},
+    {file = "tiktoken-0.5.2-cp311-cp311-win_amd64.whl", hash = "sha256:bcae1c4c92df2ffc4fe9f475bf8148dbb0ee2404743168bbeb9dcc4b79dc1fdd"},
+    {file = "tiktoken-0.5.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:b76a1e17d4eb4357d00f0622d9a48ffbb23401dcf36f9716d9bd9c8e79d421aa"},
+    {file = "tiktoken-0.5.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:01d8b171bb5df4035580bc26d4f5339a6fd58d06f069091899d4a798ea279d3e"},
+    {file = "tiktoken-0.5.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42adf7d4fb1ed8de6e0ff2e794a6a15005f056a0d83d22d1d6755a39bffd9e7f"},
+    {file = "tiktoken-0.5.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4c3f894dbe0adb44609f3d532b8ea10820d61fdcb288b325a458dfc60fefb7db"},
+    {file = "tiktoken-0.5.2-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:58ccfddb4e62f0df974e8f7e34a667981d9bb553a811256e617731bf1d007d19"},
+    {file = "tiktoken-0.5.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:58902a8bad2de4268c2a701f1c844d22bfa3cbcc485b10e8e3e28a050179330b"},
+    {file = "tiktoken-0.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:5e39257826d0647fcac403d8fa0a474b30d02ec8ffc012cfaf13083e9b5e82c5"},
+    {file = "tiktoken-0.5.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:8bde3b0fbf09a23072d39c1ede0e0821f759b4fa254a5f00078909158e90ae1f"},
+    {file = "tiktoken-0.5.2-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:2ddee082dcf1231ccf3a591d234935e6acf3e82ee28521fe99af9630bc8d2a60"},
+    {file = "tiktoken-0.5.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:35c057a6a4e777b5966a7540481a75a31429fc1cb4c9da87b71c8b75b5143037"},
+    {file = "tiktoken-0.5.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4c4a049b87e28f1dc60509f8eb7790bc8d11f9a70d99b9dd18dfdd81a084ffe6"},
+    {file = "tiktoken-0.5.2-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:5bf5ce759089f4f6521ea6ed89d8f988f7b396e9f4afb503b945f5c949c6bec2"},
+    {file = "tiktoken-0.5.2-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:0c964f554af1a96884e01188f480dad3fc224c4bbcf7af75d4b74c4b74ae0125"},
+    {file = "tiktoken-0.5.2-cp38-cp38-win_amd64.whl", hash = "sha256:368dd5726d2e8788e47ea04f32e20f72a2012a8a67af5b0b003d1e059f1d30a3"},
+    {file = "tiktoken-0.5.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:a2deef9115b8cd55536c0a02c0203512f8deb2447f41585e6d929a0b878a0dd2"},
+    {file = "tiktoken-0.5.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:2ed7d380195affbf886e2f8b92b14edfe13f4768ff5fc8de315adba5b773815e"},
+    {file = "tiktoken-0.5.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c76fce01309c8140ffe15eb34ded2bb94789614b7d1d09e206838fc173776a18"},
+    {file = "tiktoken-0.5.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:60a5654d6a2e2d152637dd9a880b4482267dfc8a86ccf3ab1cec31a8c76bfae8"},
+    {file = "tiktoken-0.5.2-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:41d4d3228e051b779245a8ddd21d4336f8975563e92375662f42d05a19bdff41"},
+    {file = "tiktoken-0.5.2-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:a5c1cdec2c92fcde8c17a50814b525ae6a88e8e5b02030dc120b76e11db93f13"},
+    {file = "tiktoken-0.5.2-cp39-cp39-win_amd64.whl", hash = "sha256:84ddb36faedb448a50b246e13d1b6ee3437f60b7169b723a4b2abad75e914f3e"},
+    {file = "tiktoken-0.5.2.tar.gz", hash = "sha256:f54c581f134a8ea96ce2023ab221d4d4d81ab614efa0b2fbce926387deb56c80"},
 ]
 
 [package.dependencies]
@@ -4016,13 +4081,13 @@ zstd = ["zstandard (>=0.18.0)"]
 
 [[package]]
 name = "virtualenv"
-version = "20.24.7"
+version = "20.25.0"
 description = "Virtual Python Environment builder"
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "virtualenv-20.24.7-py3-none-any.whl", hash = "sha256:a18b3fd0314ca59a2e9f4b556819ed07183b3e9a3702ecfe213f593d44f7b3fd"},
-    {file = "virtualenv-20.24.7.tar.gz", hash = "sha256:69050ffb42419c91f6c1284a7b24e0475d793447e35929b488bf6a0aade39353"},
+    {file = "virtualenv-20.25.0-py3-none-any.whl", hash = "sha256:4238949c5ffe6876362d9c0180fc6c3a824a7b12b80604eeb8085f2ed7460de3"},
+    {file = "virtualenv-20.25.0.tar.gz", hash = "sha256:bf51c0d9c7dd63ea8e44086fa1e4fb1093a31e963b86959257378aef020e1f1b"},
 ]
 
 [package.dependencies]
@@ -4272,4 +4337,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.10,<4.0"
-content-hash = "d67426d3c71d54ebf0f24b6031055847ca179250e0337ecac75edb71596272d5"
+content-hash = "9ab4ecec5d5833942bae07376695eff087cc2c190d6002309185d9441fad0a3f"
diff --git a/pyproject.toml b/pyproject.toml
index 5073111cd..16ccab44a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,18 +1,22 @@
 [tool.poetry]
 name = "ScandEval"
 version = "8.0.0"
-description = "Evaluation of pretrained language models on mono- or multilingual Scandinavian language tasks."
-authors = ["Dan Saattrup Nielsen <saattrupdan@gmail.com>"]
+description = "Evaluation of pretrained language models on mono- or multilingual language tasks."
+authors = ["Dan Saattrup Nielsen <dan.nielsen@alexandra.dk>"]
+maintainers = [
+    "Dan Saattrup Nielsen <dan.nielsen@alexandra.dk>",
+    "Kenneth Enevoldsen <kenneth.enevoldsen@cas.au.dk>",
+]
 readme = "README.md"
 license = "MIT"
 homepage = "https://scandeval.github.io"
-repository = "https://github.com/saattrupdan/ScandEval"
+repository = "https://github.com/ScandEval/ScandEval"
 
 [tool.poetry.dependencies]
 python = ">=3.10,<4.0"
 tqdm = "^4.62.0"
 huggingface-hub = ">=0.7.0,<1.0.0"
-datasets = "^2.7.0"
+datasets = "^2.15.0"
 click = "^8.1.3"
 termcolor = "^2.0.0"
 numpy = "^1.23.0"
@@ -34,14 +38,14 @@ bitsandbytes = {markers="sys_platform != 'darwin' or platform_machine != 'arm64'
 requests = "^2.31.0"
 pydantic = "^2.3.0"
 torch = "^2.1.1"
-transformers = "^4.34.0"
+transformers = "^4.35.0"
 accelerate = ">=0.23.0,<1.0.0"
-evaluate = ">=0.4.0,<1.0.0"
 nltk = "^3.8.1"
 rouge-score = ">=0.1.2,<1.0.0"
 bert-score = ">=0.3.13,<1.0.0"
 sacrebleu = "^2.3.2"
 cer = "^1.2.0"
+evaluate = ">=0.4.1,<1.0.0"
 
 [tool.poetry.group.dev.dependencies]
 pdoc = "^14.0.0"
@@ -49,9 +53,10 @@ pytest = "^7.4.0"
 pre-commit = "^3.3.3"
 black = "^23.7.0"
 lxml = "^4.9.3"
-isort = "^5.12.0"
 pytest-cov = "^4.1.0"
 readme-coverage-badger = ">=0.1.2,<1.0.0"
+ruff = ">=0.1.6,<1.0.0"
+mypy = "^1.7.1"
 
 [tool.poetry.scripts]
 scandeval = "scandeval.cli:benchmark"
diff --git a/src/scandeval/benchmark_config_factory.py b/src/scandeval/benchmark_config_factory.py
index 33441fffa..99c02e3fd 100644
--- a/src/scandeval/benchmark_config_factory.py
+++ b/src/scandeval/benchmark_config_factory.py
@@ -26,6 +26,7 @@ def build_benchmark_config(
     device: Device | None,
     trust_remote_code: bool,
     load_in_4bit: bool | None,
+    use_flash_attention: bool,
     testing: bool = False,
 ) -> BenchmarkConfig:
     """Create a benchmark configuration.
@@ -76,9 +77,11 @@ def build_benchmark_config(
         trust_remote_code:
             Whether to trust remote code when loading models from the Hugging Face
             Hub.
-        load_in_4bit (bool or None):
+        load_in_4bit:
             Whether to load models in 4-bit precision. If None then this will be done
             if CUDA is available and the model is a decoder model. Defaults to None.
+        use_flash_attention:
+            Whether to use Flash Attention.
         testing:
             Whether to run the benchmark in testing mode. Defaults to False.
     """
@@ -114,6 +117,7 @@ def build_benchmark_config(
         device=torch_device,
         trust_remote_code=trust_remote_code,
         load_in_4bit=load_in_4bit,
+        use_flash_attention=use_flash_attention,
         testing=testing,
     )
 
diff --git a/src/scandeval/benchmarker.py b/src/scandeval/benchmarker.py
index 4cc51a945..1d1c7c19a 100644
--- a/src/scandeval/benchmarker.py
+++ b/src/scandeval/benchmarker.py
@@ -95,9 +95,11 @@ class Benchmarker:
             Whether to output additional output. Defaults to False.
         trust_remote_code:
             Whether to trust remote code when loading models. Defaults to False.
-        load_in_4bit (bool or None, optional):
+        load_in_4bit:
             Whether to load models in 4-bit precision. If None then this will be done
             if CUDA is available and the model is a decoder model. Defaults to None.
+        use_flash_attention:
+            Whether to use Flash Attention. Defaults to False.
 
     Attributes:
         progress_bar: Whether progress bars should be shown.
@@ -130,6 +132,7 @@ def __init__(
         verbose: bool = False,
         trust_remote_code: bool = False,
         load_in_4bit: bool | None = None,
+        use_flash_attention: bool = False,
     ) -> None:
         self.benchmark_config = build_benchmark_config(
             language=language,
@@ -149,6 +152,7 @@ def __init__(
             device=device,
             trust_remote_code=trust_remote_code,
             load_in_4bit=load_in_4bit,
+            use_flash_attention=use_flash_attention,
         )
 
         # Set attributes from arguments
@@ -235,9 +239,9 @@ def benchmark(
                     error_msg = record["error"]
                     logger.info(
                         f"{m_id} could not be benchmarked on "
-                        f"{dataset_config.pretty_name}. Skipping."
+                        f"{dataset_config.pretty_name}. Skipping. The error message "
+                        f"raise dwas {error_msg!r}."
                     )
-                    logger.debug(f"The error message was {error_msg!r}.")
                     continue
 
                 # Add the record to the benchmark results
diff --git a/src/scandeval/cli.py b/src/scandeval/cli.py
index a073167cc..6b56c24b6 100644
--- a/src/scandeval/cli.py
+++ b/src/scandeval/cli.py
@@ -176,6 +176,12 @@
     show_default=True,
     help="Whether to load the model in 4-bit precision.",
 )
+@click.option(
+    "--use-flash-attention/--no-use-flash-attention",
+    default=False,
+    show_default=True,
+    help="Whether to use Flash Attention.",
+)
 def benchmark(
     model_id: tuple[str],
     dataset: tuple[str],
@@ -197,8 +203,9 @@ def benchmark(
     device: str | None,
     trust_remote_code: bool,
     load_in_4bit: bool | None,
+    use_flash_attention: bool,
 ) -> None:
-    """Benchmark pretrained language models on Scandinavian language tasks."""
+    """Benchmark pretrained language models on language tasks."""
 
     # Set up language variables
     model_ids = None if len(model_id) == 0 else list(model_id)
@@ -230,6 +237,7 @@ def benchmark(
         device=device,
         trust_remote_code=trust_remote_code,
         load_in_4bit=load_in_4bit,
+        use_flash_attention=use_flash_attention,
     )
 
     # Perform the benchmark evaluation
diff --git a/src/scandeval/config.py b/src/scandeval/config.py
index 49247c4c0..5c5401f6f 100644
--- a/src/scandeval/config.py
+++ b/src/scandeval/config.py
@@ -104,7 +104,7 @@ class BenchmarkConfig:
             The authentication token for the Hugging Face Hub. If a boolean value is
             specified then the token will be fetched from the Hugging Face CLI, where
             the user has logged in through `huggingface-cli login`. If a string is
-            specified then it will be used as the token. Defaults to False.
+            specified then it will be used as the token.
         openai_api_key:
             The API key for the OpenAI API. If None then OpenAI models will not be
             benchmarked.
@@ -122,7 +122,9 @@ class BenchmarkConfig:
             Hub.
         load_in_4bit:
             Whether to load models in 4-bit precision. If None then this will be done
-            if CUDA is available and the model is a decoder model. Defaults to None.
+            if CUDA is available and the model is a decoder model.
+        use_flash_attention:
+            Whether to use Flash Attention.
         testing:
             Whether a unit test is being run. Defaults to False.
     """
@@ -143,6 +145,7 @@ class BenchmarkConfig:
     verbose: bool
     trust_remote_code: bool
     load_in_4bit: bool | None
+    use_flash_attention: bool
     testing: bool = False
 
 
diff --git a/src/scandeval/dataset_configs.py b/src/scandeval/dataset_configs.py
index 72d1db1df..0e93666b2 100644
--- a/src/scandeval/dataset_configs.py
+++ b/src/scandeval/dataset_configs.py
@@ -1,7 +1,7 @@
 """All dataset configurations used in ScandEval."""
 
 from .config import DatasetConfig
-from .dataset_tasks import LA, NER, QA, SENT, SPEED, TEXT_TO_TEXT
+from .dataset_tasks import LA, NER, QA, SENT, SPEED, SUMMARIZATION
 from .languages import DA, DE, NL, FO, IS, NB, NN, SV, get_all_languages
 
 
@@ -444,7 +444,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     name="nordjylland-news",
     pretty_name="the truncated version of Nordjylland News",
     huggingface_id="ScandEval/nordjylland-news-mini",
-    task=TEXT_TO_TEXT,
+    task=SUMMARIZATION,
     languages=[DA],
     prompt_template="{text}\nOpsummering: {target_text}",
     num_few_shot_examples=2,
@@ -456,7 +456,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     name="mlsum",
     pretty_name="the truncated version of MLSum",
     huggingface_id="ScandEval/mlsum-mini",
-    task=TEXT_TO_TEXT,
+    task=SUMMARIZATION,
     languages=[DE],
     prompt_template="{text}\nZusammenfassung: {target_text}",
     num_few_shot_examples=2,
@@ -468,7 +468,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     name="rrn",
     pretty_name="the truncated version of RÚV Radio News",
     huggingface_id="ScandEval/rrn-mini",
-    task=TEXT_TO_TEXT,
+    task=SUMMARIZATION,
     languages=[IS],
     prompt_template="{text}\nSamantekt: {target_text}",
     num_few_shot_examples=2,
@@ -480,7 +480,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     name="no-sammendrag",
     pretty_name="the truncated version of the Norwegian Sammendrag dataset",
     huggingface_id="ScandEval/no-sammendrag-mini",
-    task=TEXT_TO_TEXT,
+    task=SUMMARIZATION,
     languages=[NB, NN],
     prompt_template="{text}\nSammendrag: {target_text}",
     num_few_shot_examples=2,
@@ -492,7 +492,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     name="wiki-lingua-nl",
     pretty_name="the Dutch part of the truncated version of WikiLingua",
     huggingface_id="ScandEval/wiki-lingua-nl-mini",
-    task=TEXT_TO_TEXT,
+    task=SUMMARIZATION,
     languages=[NL],
     prompt_template="{text}\nSamenvatting: {target_text}",
     num_few_shot_examples=2,
@@ -504,7 +504,7 @@ def get_dataset_config(dataset_name: str) -> DatasetConfig:
     name="swedn",
     pretty_name="the truncated version of SweDN",
     huggingface_id="ScandEval/swedn-mini",
-    task=TEXT_TO_TEXT,
+    task=SUMMARIZATION,
     languages=[SV],
     prompt_template="{text}\nSammanfattning: {target_text}",
     num_few_shot_examples=2,
diff --git a/src/scandeval/dataset_tasks.py b/src/scandeval/dataset_tasks.py
index a42ab671d..a0a0bdf23 100644
--- a/src/scandeval/dataset_tasks.py
+++ b/src/scandeval/dataset_tasks.py
@@ -110,30 +110,8 @@ def get_all_dataset_tasks() -> dict[str, DatasetTask]:
 )
 
 
-MULTIPLE_CHOICE = DatasetTask(
-    name="multiple-choice",
-    supertask="multiple-choice",
-    metrics=[  # Not sure
-        MetricConfig(
-            name="mcc",
-            pretty_name="Matthew's Correlation Coefficient",
-            huggingface_id="matthews_correlation",
-            results_key="matthews_correlation",
-        ),
-        MetricConfig(
-            name="macro_f1",
-            pretty_name="Macro-average F1-score",
-            huggingface_id="f1",
-            results_key="f1",
-            compute_kwargs=dict(average="macro"),
-        ),
-    ],
-    labels=[],  # Maybe add?
-)
-
-
-TEXT_TO_TEXT = DatasetTask(
-    name="text-to-text",
+SUMMARIZATION = DatasetTask(
+    name="summarization",
     supertask="text-to-text",
     metrics=[
         MetricConfig(
@@ -192,6 +170,28 @@ def get_all_dataset_tasks() -> dict[str, DatasetTask]:
 )
 
 
+MULTIPLE_CHOICE = DatasetTask(
+    name="multiple-choice",
+    supertask="multiple-choice",
+    metrics=[  # Not sure
+        MetricConfig(
+            name="mcc",
+            pretty_name="Matthew's Correlation Coefficient",
+            huggingface_id="matthews_correlation",
+            results_key="matthews_correlation",
+        ),
+        MetricConfig(
+            name="macro_f1",
+            pretty_name="Macro-average F1-score",
+            huggingface_id="f1",
+            results_key="f1",
+            compute_kwargs=dict(average="macro"),
+        ),
+    ],
+    labels=[],  # Maybe add?
+)
+
+
 TEXT_MODELLING = DatasetTask(
     name="text-modelling",
     supertask="text-modelling",
diff --git a/src/scandeval/model_loading.py b/src/scandeval/model_loading.py
index f4c2ba821..40d72418f 100644
--- a/src/scandeval/model_loading.py
+++ b/src/scandeval/model_loading.py
@@ -5,7 +5,7 @@
 from transformers import PreTrainedModel
 
 from .exceptions import InvalidBenchmark
-from .utils import GENERATIVE_DATASET_TASKS
+from .utils import GENERATIVE_DATASET_SUPERTASKS
 from .config import BenchmarkConfig, DatasetConfig, ModelConfig
 from .protocols import GenerativeModel, Tokenizer, ModelSetup
 from .model_setups import (
@@ -58,8 +58,9 @@ def load_model(
         error_to_raise = e
 
     # Refuse to benchmark non-generative models on generative tasks
-    if dataset_config.task.name in GENERATIVE_DATASET_TASKS and not isinstance(
-        model, GenerativeModel
+    if (
+        dataset_config.task.supertask in GENERATIVE_DATASET_SUPERTASKS
+        and not isinstance(model, GenerativeModel)
     ):
         raise InvalidBenchmark(error_message)
     elif error_to_raise is not None:
diff --git a/src/scandeval/model_setups/hf.py b/src/scandeval/model_setups/hf.py
index 4be5c9aa4..133ef977c 100644
--- a/src/scandeval/model_setups/hf.py
+++ b/src/scandeval/model_setups/hf.py
@@ -214,17 +214,44 @@ def load_model(
                 and self.benchmark_config.device == torch.device("cuda")
             )
 
+        bnb_config = (
+            BitsAndBytesConfig(
+                load_in_4bit=load_in_4bit,
+                bnb_4bit_compute_type=torch.float16,
+                bnb_4bit_use_double_quant=True,
+            )
+            if load_in_4bit
+            else None
+        )
+
+        use_flash_attention = (
+            self.benchmark_config.use_flash_attention or "mistral" in model_id.lower()
+        )
+
+        config = self._load_hf_model_config(
+            model_id=model_id,
+            num_labels=dataset_config.num_labels,
+            id2label=dataset_config.id2label,
+            label2id=dataset_config.label2id,
+            revision=model_config.revision,
+            model_cache_dir=model_config.model_cache_dir,
+        )
+
+        model_kwargs = dict(
+            config=config,
+            from_flax=from_flax,
+            ignore_mismatched_sizes=ignore_mismatched_sizes,
+            revision=model_config.revision,
+            token=self.benchmark_config.token,
+            cache_dir=model_config.model_cache_dir,
+            trust_remote_code=self.benchmark_config.trust_remote_code,
+            quantization_config=bnb_config,
+            torch_dtype="auto",
+            use_flash_attention_2=use_flash_attention,
+        )
+
         while True:
             try:
-                config = self._load_hf_model_config(
-                    model_id=model_id,
-                    num_labels=dataset_config.num_labels,
-                    id2label=dataset_config.id2label,
-                    label2id=dataset_config.label2id,
-                    revision=model_config.revision,
-                    model_cache_dir=model_config.model_cache_dir,
-                )
-
                 # Get the model class associated with the supertask
                 if model_config.task in ["text-generation", "conversational"]:
                     model_cls_supertask = "causal-l-m"
@@ -255,28 +282,18 @@ def load_model(
                     warnings.filterwarnings("ignore", category=FutureWarning)
                     with HiddenPrints():
                         try:
-                            bnb_config = (
-                                BitsAndBytesConfig(
-                                    load_in_4bit=load_in_4bit,
-                                    bnb_4bit_compute_type=torch.float16,
-                                    bnb_4bit_use_double_quant=True,
-                                )
-                                if load_in_4bit
-                                else None
-                            )
                             model_or_tuple = model_cls_or_none.from_pretrained(
-                                model_config.model_id,
-                                config=config,
-                                from_flax=from_flax,
-                                ignore_mismatched_sizes=ignore_mismatched_sizes,
-                                revision=model_config.revision,
-                                token=self.benchmark_config.token,
-                                cache_dir=model_config.model_cache_dir,
-                                trust_remote_code=(
-                                    self.benchmark_config.trust_remote_code
-                                ),
-                                quantization_config=bnb_config,
+                                model_config.model_id, **model_kwargs
                             )
+                        except ImportError as e:
+                            if "flash attention" in str(e).lower():
+                                raise InvalidBenchmark(
+                                    "The model you are trying to load requires Flash "
+                                    "Attention. To use Flash Attention, please install "
+                                    "the `flash-attn` package, which can be done by "
+                                    "running `pip install --no-build-isolation -U "
+                                    "flash-attn`."
+                                )
                         except (KeyError, RuntimeError) as e:
                             if not ignore_mismatched_sizes:
                                 ignore_mismatched_sizes = True
@@ -311,6 +328,7 @@ def load_model(
 
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=UserWarning)
+            warnings.filterwarnings("ignore", category=FutureWarning)
             tokenizer = self._load_tokenizer(model=model, model_id=model_id)
 
         model, tokenizer = align_model_and_tokenizer(
@@ -356,16 +374,19 @@ def _load_hf_model_config(
         """
         while True:
             try:
-                config = AutoConfig.from_pretrained(
-                    model_id,
-                    num_labels=num_labels,
-                    id2label=id2label,
-                    label2id=label2id,
-                    revision=revision,
-                    cache_dir=model_cache_dir,
-                    token=self.benchmark_config.token,
-                    trust_remote_code=self.benchmark_config.trust_remote_code,
-                )
+                with warnings.catch_warnings():
+                    warnings.filterwarnings("ignore", category=UserWarning)
+                    warnings.filterwarnings("ignore", category=FutureWarning)
+                    config = AutoConfig.from_pretrained(
+                        model_id,
+                        num_labels=num_labels,
+                        id2label=id2label,
+                        label2id=label2id,
+                        revision=revision,
+                        cache_dir=model_cache_dir,
+                        token=self.benchmark_config.token,
+                        trust_remote_code=self.benchmark_config.trust_remote_code,
+                    )
                 if config.eos_token_id is not None and config.pad_token_id is None:
                     config.pad_token_id = config.eos_token_id
                 return config
@@ -405,6 +426,7 @@ def _load_tokenizer(
         padding_side = "left" if model_is_generative(model=model) else "right"
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=UserWarning)
+            warnings.filterwarnings("ignore", category=FutureWarning)
             while True:
                 try:
                     return AutoTokenizer.from_pretrained(
diff --git a/src/scandeval/utils.py b/src/scandeval/utils.py
index 622b28a59..4d303a7d8 100644
--- a/src/scandeval/utils.py
+++ b/src/scandeval/utils.py
@@ -39,10 +39,10 @@
 ]
 
 
-GENERATIVE_DATASET_TASKS = [
+GENERATIVE_DATASET_SUPERTASKS = [
     "multiple-choice",
     "text-to-text",
-    "raw-text",
+    "text-modelling",
 ]
 
 


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 2e6a69669..38b898f3c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -37,6 +37,7 @@ def benchmark_config() -> Generator[BenchmarkConfig, None, None]:
         verbose=False,
         trust_remote_code=False,
         load_in_4bit=None,
+        use_flash_attention=False,
         testing=True,
     )
 
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 212c7db72..b2c8431cf 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -34,6 +34,7 @@ def test_cli_param_names(params):
         "device",
         "trust_remote_code",
         "load_in_4bit",
+        "use_flash_attention",
         "help",
     }
 
@@ -59,4 +60,5 @@ def test_cli_param_types(params):
     assert isinstance(params["device"], Choice)
     assert params["trust_remote_code"] == BOOL
     assert params["load_in_4bit"] == BOOL
+    assert params["use_flash_attention"] == BOOL
     assert params["help"] == BOOL
diff --git a/tests/test_config.py b/tests/test_config.py
index 08f3f9314..52835b847 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -108,6 +108,7 @@ def benchmark_config(self, language, dataset_task):
             verbose=False,
             trust_remote_code=False,
             load_in_4bit=None,
+            use_flash_attention=False,
             testing=True,
         )
 
@@ -133,6 +134,7 @@ def test_attributes_correspond_to_arguments(
         assert benchmark_config.verbose is False
         assert benchmark_config.trust_remote_code is False
         assert benchmark_config.load_in_4bit is None
+        assert benchmark_config.use_flash_attention is False
         assert benchmark_config.testing is True
 
 
diff --git a/tests/test_sequence_classification.py b/tests/test_sequence_classification.py
index 5825b2898..7a28a4ac2 100644
--- a/tests/test_sequence_classification.py
+++ b/tests/test_sequence_classification.py
@@ -17,12 +17,12 @@
 @pytest.mark.parametrize(
     argnames=["dataset", "correct_scores"],
     argvalues=[
-        (ANGRY_TWEETS_CONFIG, (5.32, 27.37)),
+        (ANGRY_TWEETS_CONFIG, (-4.43, 23.00)),
         (SWEREC_CONFIG, (0.00, 23.07)),
         (NOREC_CONFIG, (2.32, 25.12)),
         (SCALA_DA_CONFIG, (3.69, 36.61)),
-        (SCALA_SV_CONFIG, (9.36, 42.31)),
-        (SCALA_NB_CONFIG, (4.70, 33.04)),
+        (SCALA_SV_CONFIG, (2.85, 36.84)),
+        (SCALA_NB_CONFIG, (0.00, 30.71)),
         (SCALA_NN_CONFIG, (0.00, 33.39)),
     ],
     ids=[