Add model utils tests

src/alignment/configs.py

@@ -17,7 +17,7 @@ import dataclasses
 import os
 import sys
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, NewType, Optional, Tuple, Union
+from typing import Any, Dict, List, NewType, Optional, Tuple
 
 import transformers
 from transformers import MODEL_FOR_CAUSAL_LM_MAPPING, HfArgumentParser
@@ -87,7 +87,7 @@ class H4ArgumentParser(HfArgumentParser):
 
         return outputs
 
-    def parse(self) -> Union[DataClassType, Tuple[DataClassType]]:
+    def parse(self) -> DataClassType | Tuple[DataClassType]:
         if len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
             # If we pass only one argument to the script and it's the path to a YAML file,
             # let's parse it to get our arguments.

src/alignment/data.py

@@ -1,5 +1,5 @@
 import re
-from typing import List, Literal, Optional, Union
+from typing import List, Literal, Optional
 
 from datasets import DatasetDict, concatenate_datasets, load_dataset
 
@@ -67,7 +67,7 @@ def apply_chat_template(
 
 
 def get_datasets(
-    data_config: Union[DataArguments, dict],
+    data_config: DataArguments | dict,
     splits: List[str] = ["train", "test"],
     shuffle: bool = True,
 ) -> DatasetDict:

src/alignment/model_utils.py

@@ -1,4 +1,4 @@
-from typing import Dict, Union
+from typing import Dict
 
 import torch
 from transformers import AutoTokenizer, BitsAndBytesConfig, PreTrainedTokenizer
@@ -62,7 +62,7 @@ def get_tokenizer(model_args: ModelArguments, data_args: DataArguments) -> PreTr
     return tokenizer
 
 
-def get_peft_config(model_args: ModelArguments) -> Union[PeftConfig, None]:
+def get_peft_config(model_args: ModelArguments) -> PeftConfig | None:
     if model_args.use_peft is False:
         return None