update: language_abbr parameter to a string

DOCS/gettingstarted.md

@@ -36,7 +36,7 @@
 huggingface_token = " "  # make sure token has write permissions
 dataset_name = "mozilla-foundation/common_voice_16_1" # Also supports "google/fleurs" and "facebook/multilingual_librispeech".
                                                       # For custom datasets, ensure the text key is one of the following: "sentence", "transcript", or "transcription".
-language_abbr= [ ]                                    # Example `["af"]`. see specific dataset for language code.
+language_abbr= " "                                    # Example `"af"`. see specific dataset for language code.
 model_id= "model-id"                                  # Example openai/whisper-small, openai/whisper-medium
 processing_task= "translate"                          # translate or transcribe
 wandb_api_key = " "     

src/training/data_prep.py

@@ -20,7 +20,7 @@
 from .load_data import Dataset
 from .whisper_model_prep import WhisperModelPrep
 from .audio_data_processor import AudioDataProcessor
-from typing import Tuple, List
+from typing import Tuple
 import warnings
 warnings.filterwarnings("ignore")
 
@@ -39,7 +39,7 @@ def __init__(
         self,
         huggingface_token: str,
         dataset_name: str,
-        language_abbr: List[str],
+        language_abbr: str,
         model_id: str,
         processing_task: str,
         use_peft: bool,

src/training/load_data.py

@@ -1,6 +1,5 @@
-from datasets import load_dataset, IterableDatasetDict, concatenate_datasets
+from datasets import load_dataset, IterableDatasetDict
 import warnings
-from typing import List
 from datasets import DatasetDict
 from huggingface_hub import HfFolder
 warnings.filterwarnings("ignore")
@@ -15,7 +14,7 @@ class Dataset:
         language_abbr (str): Abbreviation of the language for the dataset.
     """
 
-    def __init__(self, huggingface_token: str, dataset_name: str, language_abbr: List[str]):
+    def __init__(self, huggingface_token: str, dataset_name: str, language_abbr: str):
         """
         Initializes the DatasetManager with necessary details for dataset operations.
 
@@ -46,39 +45,44 @@ def load_dataset(self, streaming: bool = True, train_num_samples: int = None, te
             dict: A dictionary containing concatenated train and test splits for each language.
         """
         data = {}
-        for lang in self.language_abbr:
-            train_dataset = load_dataset(self.dataset_name,
-                                         lang,
-                                         split='train',
-                                         streaming=streaming,
-                                         token=self.huggingface_token,
-                                         trust_remote_code=True)
-            test_dataset = load_dataset(self.dataset_name,
-                                        lang,
-                                        split='test',
-                                        streaming=streaming,
-                                        token=self.huggingface_token,
-                                        trust_remote_code=True)
-            if streaming:
-                train_split = train_dataset.take(train_num_samples) if train_num_samples else train_dataset
-                test_split = test_dataset.take(test_num_samples) if test_num_samples else test_dataset
-
-            else:
-
-                train_split = train_dataset if not train_num_samples or len(train_dataset) < train_num_samples else \
-                    train_dataset.select(range(train_num_samples))
-
-                test_split = test_dataset if not test_num_samples or len(test_dataset) < test_num_samples else \
-                    test_dataset.select(range(test_num_samples))
-
-            if "train" in data:
-                data["train"] = concatenate_datasets([data["train"], train_split])
-            else:
-                data["train"] = train_split
-            if "test" in data:
-                data["test"] = concatenate_datasets([data["test"], test_split])
-            else:
-                data["test"] = test_split
+
+        train_dataset = load_dataset(
+            self.dataset_name,
+            self.language_abbr,
+            split="train",
+            streaming=streaming,
+            token=self.huggingface_token,
+            trust_remote_code=True,
+        )
+        test_dataset = load_dataset(
+            self.dataset_name,
+            self.language_abbr,
+            split="test",
+            streaming=streaming,
+            token=self.huggingface_token,
+            trust_remote_code=True,
+        )
+
+        if streaming:
+            train_split = train_dataset.take(train_num_samples) if train_num_samples else train_dataset
+        else:
+            train_split = (
+                train_dataset
+                if not train_num_samples or len(train_dataset) < train_num_samples
+                else train_dataset.select(range(train_num_samples))
+            )
+
+        if streaming:
+            test_split = test_dataset.take(test_num_samples) if test_num_samples else test_dataset
+        else:
+            test_split = (
+                test_dataset
+                if not test_num_samples or len(test_dataset) < test_num_samples
+                else test_dataset.select(range(test_num_samples))
+            )
+
+        data["train"] = train_split
+        data["test"] = test_split
 
         return data
 

src/training/main.py

@@ -41,9 +41,8 @@ def parse_args():
     )
     parser.add_argument(
         "--language_abbr",
-        nargs='+',
         required=True,
-        help="Abbreviation(s) of the language(s) for the dataset.",
+        help="Abbreviation of the language(s) for the dataset.",
     )
     parser.add_argument(
         "--model_id",

src/training/whisper_model_prep.py

@@ -29,7 +29,7 @@ class WhisperModelPrep:
     def __init__(
         self,
         model_id: str,
-        language: list,
+        language: str,
         processing_task: str,
         use_peft: bool,
     ):
@@ -45,7 +45,7 @@ def __init__(
         """
 
         self.model_id = model_id
-        self.language = language[0]
+        self.language = language
         self.processing_task = processing_task
         self.use_peft = use_peft
 

tests/test_audio_processor.py

@@ -19,7 +19,7 @@ def setUp(self):
         self.data_loader = Dataset(
             huggingface_token = os.environ.get("HF_TOKEN"),
             dataset_name="mozilla-foundation/common_voice_16_1",
-            language_abbr=["af"]
+            language_abbr="af"
         )
         self.dataset_streaming = self.data_loader.load_dataset(streaming=True, train_num_samples=10, test_num_samples=10)
         self.dataset_batch = self.data_loader.load_dataset(streaming=False, train_num_samples=10, test_num_samples=10)
@@ -38,7 +38,7 @@ def setUp(self):
 
         # Initialize model preparation
         self.model_prep = WhisperModelPrep(
-            language = ["af"],
+            language = "af",
             model_id="openai/whisper-tiny",
             processing_task="transcribe",
             use_peft=False

tests/test_data_prep.py

@@ -16,13 +16,13 @@ def setUp(self):
         self.data_prep = DataPrep(
             huggingface_token= os.environ.get("HF_TOKEN"),
             dataset_name="mozilla-foundation/common_voice_16_1",
-            language_abbr=["af"],
+            language_abbr="af",
             model_id="openai/whisper-small",
             processing_task="transcribe",
             use_peft=False,
         )
         self.model_prep=WhisperModelPrep(
-            language= ["af"],
+            language= "af",
             model_id="openai/whisper-small",
             processing_task="transcribe",
             use_peft=False),

tests/test_load_dataset.py

@@ -12,7 +12,7 @@ def setUp(self):
         self.dataset_manager = Dataset(
             huggingface_token=os.environ.get("HF_TOKEN"),
             dataset_name="mozilla-foundation/common_voice_16_1",
-            language_abbr=["af"]
+            language_abbr="af"
         )
 
     def test_load_dataset(self):

tests/test_model_prep.py

@@ -8,7 +8,7 @@ class TestDatasetManager(unittest.TestCase):
     def setUp(self):
         """Initialize the test setup with an instance of WhisperModelPrep."""
         self.model_prep = WhisperModelPrep(
-            language=["af"],
+            language="af",
             model_id="openai/whisper-small",
             processing_task="transcribe",
             use_peft=False

tests/test_model_trainer_cpu.py

@@ -14,7 +14,7 @@ def setUp(self) -> None:
         process = DataPrep(
             huggingface_token=os.environ.get("HF_TOKEN"),
             dataset_name="mozilla-foundation/common_voice_16_1",
-            language_abbr=["af"],
+            language_abbr="af",
             model_id=self.model_id,
             processing_task="transcribe",
             use_peft=False,

tests/test_model_trainer_gpu.py

@@ -14,7 +14,7 @@ def setUp(self) -> None:
         process = DataPrep(
             huggingface_token=os.environ.get("HF_TOKEN"),
             dataset_name="mozilla-foundation/common_voice_16_1",
-            language_abbr=["af"],
+            language_abbr="af",
             model_id=self.model_id,
             processing_task="transcribe",
             use_peft=True,