wolof-translate 0.0.3__py3-none-any.whl → 0.0.4__py3-none-any.whl

wolof_translate/utils/bucket_iterator.py

@@ -5,29 +5,43 @@ from math import ceil
 from tqdm import tqdm
 import time
 
+
 class SequenceLengthBatchSampler(Sampler[List[int]]):
     def __init__(
         self,
         dataset,
         boundaries: List[int],
         batch_sizes: List[int],
-        input_key: Optional[int] = None,
-        label_key: Optional[int] = None,
+        input_key: Optional[Union[int, str]] = None,
+        label_key: Optional[Union[int, str]] = None,
         drop_unique: bool = True,
     ):
+        """
+        Sampler that batches sequences of similar lengths together to minimize padding.
+
+        Args:
+            dataset: Dataset to sample from.
+            boundaries: List of length boundaries to bucket sequences.
+            batch_sizes: List of batch sizes per bucket (length must be len(boundaries)+1).
+            input_key: Key or index to access input sequence in dataset item.
+            label_key: Key or index to access label sequence in dataset item.
+            drop_unique: Whether to drop batches with a single leftover element.
+        """
         self.dataset = dataset
         self.boundaries = boundaries
         self.batch_sizes = batch_sizes
         self.drop_unique = drop_unique
 
+        assert len(batch_sizes) == len(boundaries) + 1, (
+            f"batch_sizes length ({len(batch_sizes)}) must be one more than boundaries length ({len(boundaries)})"
+        )
+
         start_time = time.time()
         tqdm.write("Computing sequence lengths...")
 
-        # Compute lengths with tqdm progress bar
         self.lengths = np.array([
-            max(len(data[0]), len(data[2])) if input_key is None or label_key is None
-            else max(len(data[input_key]), len(data[label_key]))
-            for data in tqdm(dataset, desc="Lengths", unit="seq")
+            self._get_length(data, input_key, label_key)
+            for data in tqdm(self.dataset, desc="Lengths", unit="seq")
         ])
 
         tqdm.write(f"Sequence lengths computed in {time.time() - start_time:.2f} seconds.")
@@ -35,7 +49,7 @@ class SequenceLengthBatchSampler(Sampler[List[int]]):
         start_time = time.time()
         tqdm.write("Assigning buckets...")
 
-        # Assign bucket ids using digitize (vectorized)
+        # Assign bucket ids (0-based)
         self.bucket_ids = np.digitize(self.lengths, bins=self.boundaries, right=True)
 
         # Create buckets of indices
@@ -46,7 +60,6 @@ class SequenceLengthBatchSampler(Sampler[List[int]]):
         start_time = time.time()
         tqdm.write("Preparing batches...")
 
-        # Prepare batches from buckets
         self.batches = []
         for bucket, batch_size in zip(self.buckets, self.batch_sizes):
             bucket = bucket.copy()
@@ -66,8 +79,26 @@ class SequenceLengthBatchSampler(Sampler[List[int]]):
         self.length = len(self.batches)
         tqdm.write(f"Batches prepared in {time.time() - start_time:.2f} seconds.")
 
+    def _get_length(self, data, input_key, label_key) -> int:
+        """
+        Helper to get the max length of input and label sequences in a dataset item.
+
+        Supports dict-like or tuple/list-like dataset items.
+        """
+        try:
+            if input_key is None or label_key is None:
+                # Assume tuple/list with input at 0, label at 2
+                input_seq = data[0]
+                label_seq = data[2]
+            else:
+                input_seq = data[input_key]
+                label_seq = data[label_key]
+            return max(len(input_seq), len(label_seq))
+        except Exception as e:
+            raise ValueError(f"Error accessing lengths with input_key={input_key}, label_key={label_key}: {e}")
+
     def __iter__(self) -> Iterator[List[int]]:
-        # Shuffle all batches globally to add randomness between buckets
+        # Shuffle batches globally for randomness
         np.random.shuffle(self.batches)
         for batch in self.batches:
             yield batch

{wolof_translate-0.0.3.dist-info → wolof_translate-0.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: wolof-translate
-Version: 0.0.3
+Version: 0.0.4
 Summary: Contain function and classes to process corpora for making translation between wolof text and other languages.
 Author: Oumar Kane
 Author-email: oumar.kane@univ-thies.sn

{wolof_translate-0.0.3.dist-info → wolof_translate-0.0.4.dist-info}/RECORD

@@ -22,7 +22,7 @@ wolof_translate/trainers/transformer_trainer_custom.py,sha256=hHUBcU4YK6wuRUMiwX
 wolof_translate/trainers/transformer_trainer_ml.py,sha256=WgggaugkVHSJlwIAZT-QwI90Fl-_zT8Clhb-7M0m8gM,33561
 wolof_translate/trainers/transformer_trainer_ml_.py,sha256=QaN9DB5pqhBxV4WlFmJCmUyfwlX-UyAzKRwL6rVEr4Q,38199
 wolof_translate/utils/__init__.py,sha256=Nl3300H-Xd3uTHDR8y-rYa-UUR9FqbqZPwUKJUpQOb4,64
-wolof_translate/utils/bucket_iterator.py,sha256=JctNzZI1yU1FjZUKY-zNiRMXS3JpRiXOOSKdObRdPbg,8511
+wolof_translate/utils/bucket_iterator.py,sha256=sGSBCGPn8NzZ32mfEKh0cfH1Z0WbNJWbDjghR8-u5tU,9847
 wolof_translate/utils/database_manager.py,sha256=7yhgBN1LvVFNEQikxCjSCva82h5nX44Nx2zh8cpFWyA,3543
 wolof_translate/utils/display_predictions.py,sha256=y5H5lfgIODl6E5Zfb1YIwiAxIlHUxRBoChfQR5kjh24,5145
 wolof_translate/utils/download_model.py,sha256=x92KpfVPvNK8Suen1qnOcPtZOlB4kXTfqWgoVuuMUEM,1241
@@ -43,7 +43,7 @@ wolof_translate/utils/training.py,sha256=5vPVuqHL6_gqLkh4PTxXqW4UvAJBWNWVDDXC9Fk
 wolof_translate/utils/trunc_hg_training.py,sha256=mMGrU7Mjr9vYd7eLc8nbFRhRXwSWMKyg35lGf0L6RtQ,6418
 wolof_translate/utils/improvements/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wolof_translate/utils/improvements/end_marks.py,sha256=scmhMMYguZmrZTPozx1ZovizKrrPfPpMLXbU2-IOdGs,1194
-wolof_translate-0.0.3.dist-info/METADATA,sha256=IztySbGGsGKbrtFqCcLpppA07bO8kt2_dxdd4hdMVOI,818
-wolof_translate-0.0.3.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-wolof_translate-0.0.3.dist-info/top_level.txt,sha256=YG-kBnOwUZyQ7SofNvMxNYjzCreH2PVcW2UaEg1-Reg,16
-wolof_translate-0.0.3.dist-info/RECORD,,
+wolof_translate-0.0.4.dist-info/METADATA,sha256=Adyexcw1wpc80mDnoeLxFuaWe4Pg0ZAZroqx6eHnSz4,818
+wolof_translate-0.0.4.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+wolof_translate-0.0.4.dist-info/top_level.txt,sha256=YG-kBnOwUZyQ7SofNvMxNYjzCreH2PVcW2UaEg1-Reg,16
+wolof_translate-0.0.4.dist-info/RECORD,,