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

wolof_translate/utils/bucket_iterator.py

@@ -1,84 +1,153 @@
-
-import torch
 import numpy as np
-from typing import *
+from typing import Optional, List, Iterator, Union
 from torch.utils.data import Sampler
-from torch.nn.utils.rnn import pad_sequence
 from math import ceil
-
-class SequenceLengthBatchSampler(Sampler):
-    def __init__(self, dataset, boundaries, batch_sizes, input_key = None, label_key = None, drop_unique = True):
+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,
+        drop_unique: bool = True,
+    ):
         self.dataset = dataset
         self.boundaries = boundaries
         self.batch_sizes = batch_sizes
-        self.data_info = {}
         self.drop_unique = drop_unique
 
-        # Initialize dictionary with indices and element lengths
-        for i in range(len(dataset)):
-            data = dataset[i]
-            length = (
-                max(len(data[0]), len(data[2]))
-                if (input_key is None and label_key is None)
-                else max(len(data[input_key]), len(data[label_key]))
-            )
-            self.data_info[i] = {"index": i, "length": length}
-            
-        self.calculate_length()
-
-    def calculate_length(self):
-        self.batches = []
+        start_time = time.time()
+        tqdm.write("Computing sequence lengths...")
 
-        # Sort indices based on element length
-        sorted_indices = sorted(self.data_info.keys(), key=lambda i: self.data_info[i]["length"])
-        
-        # Group indices into batches of sequences with the same length
-        for boundary in self.boundaries:
-            batch = [i for i in sorted_indices if self.data_info[i]["length"] <= boundary]  # Filter indices based on length boundary
-            self.batches.append(batch)
-            sorted_indices = [i for i in sorted_indices if i not in batch]  # Remove processed indices
+        # 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")
+        ])
 
-        # Add remaining indices to the last batch
-        self.batches.append(sorted_indices)
+        tqdm.write(f"Sequence lengths computed in {time.time() - start_time:.2f} seconds.")
 
-        # Calculate the total length of the data loader
-        self.length = sum(ceil(len(batch) / batch_size) for batch, batch_size in zip(self.batches, self.batch_sizes) if len(batch) % batch_size != 1 or not self.drop_unique)
+        start_time = time.time()
+        tqdm.write("Assigning buckets...")
 
-    def __iter__(self):
-#         indices = list(self.data_info.keys())  # Get indices from the data_info dictionary
-#         np.random.shuffle(indices)  # Shuffle the indices
+        # Assign bucket ids using digitize (vectorized)
+        self.bucket_ids = np.digitize(self.lengths, bins=self.boundaries, right=True)
 
-        # Yield batches with the corresponding batch sizes
-        for batch_indices, batch_size in zip(self.batches, self.batch_sizes):
-            num_batches = len(batch_indices) // batch_size
+        # Create buckets of indices
+        self.buckets = [np.where(self.bucket_ids == i)[0] for i in range(len(boundaries) + 1)]
 
-            for i in range(num_batches):
-                # Recuperate the current bucket
-                current_bucket = batch_indices[i * batch_size: (i + 1) * batch_size]
+        tqdm.write(f"Buckets assigned in {time.time() - start_time:.2f} seconds.")
 
-                # Shuffle the current bucket
-                np.random.shuffle(current_bucket)
+        start_time = time.time()
+        tqdm.write("Preparing batches...")
 
-                # Yield the current bucket
-                yield [self.data_info[i]["index"] for i in current_bucket]
+        # Prepare batches from buckets
+        self.batches = []
+        for bucket, batch_size in zip(self.buckets, self.batch_sizes):
+            bucket = bucket.copy()
+            np.random.shuffle(bucket)
 
-            remaining_indices = len(batch_indices) % batch_size
-    
-            if remaining_indices > 0 and remaining_indices != 1 or not self.drop_unique:
-                
-                # Recuperate the current bucket
-                current_bucket = batch_indices[-remaining_indices:]
+            n_full_batches = len(bucket) // batch_size
+            leftover = len(bucket) % batch_size
 
-                # Shuffle the current bucket
-                np.random.shuffle(current_bucket)
+            for i in range(n_full_batches):
+                batch = bucket[i * batch_size : (i + 1) * batch_size].tolist()
+                self.batches.append(batch)
 
-                # Yield the current bucket
-                yield [self.data_info[i]["index"] for i in batch_indices[-remaining_indices:]]
+            if leftover > 0 and (leftover != 1 or not self.drop_unique):
+                batch = bucket[-leftover:].tolist()
+                self.batches.append(batch)
 
-    def __len__(self):
+        self.length = len(self.batches)
+        tqdm.write(f"Batches prepared in {time.time() - start_time:.2f} seconds.")
+
+    def __iter__(self) -> Iterator[List[int]]:
+        # Shuffle all batches globally to add randomness between buckets
+        np.random.shuffle(self.batches)
+        for batch in self.batches:
+            yield batch
+
+    def __len__(self) -> int:
         return self.length
 
 
+
+
+# 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,
+#         drop_unique: bool = True,
+#     ):
+#         self.dataset = dataset
+#         self.boundaries = boundaries
+#         self.batch_sizes = batch_sizes
+#         self.drop_unique = drop_unique
+#         self.data_info = {}
+
+#         # Extract lengths
+#         for i in range(len(dataset)):
+#             data = dataset[i]
+#             if input_key is None or label_key is None:
+#                 length = max(len(data[0]), len(data[2]))
+#             else:
+#                 length = max(len(data[input_key]), len(data[label_key]))
+#             self.data_info[i] = {"index": i, "length": length}
+
+#         self.calculate_length()
+
+#     def calculate_length(self):
+#         self.batches = []
+#         sorted_indices = sorted(self.data_info.keys(), key=lambda i: self.data_info[i]["length"])
+
+#         prev_boundary = 0
+#         for boundary in self.boundaries:
+#             batch = [i for i in sorted_indices if prev_boundary < self.data_info[i]["length"] <= boundary]
+#             self.batches.append(batch)
+#             sorted_indices = [i for i in sorted_indices if i not in batch]
+#             prev_boundary = boundary
+
+#         # Remaining sequences > last boundary
+#         self.batches.append(sorted_indices)
+
+#         total_batches = 0
+#         for batch, batch_size in zip(self.batches, self.batch_sizes):
+#             n_full_batches = len(batch) // batch_size
+#             leftover = len(batch) % batch_size
+#             total_batches += n_full_batches
+#             if leftover > 0 and (leftover != 1 or not self.drop_unique):
+#                 total_batches += 1
+#         self.length = total_batches
+
+#     def __iter__(self) -> Iterator[List[int]]:
+#         for batch_indices, batch_size in zip(self.batches, self.batch_sizes):
+#             num_batches = len(batch_indices) // batch_size
+
+#             for i in range(num_batches):
+#                 current_bucket = batch_indices[i * batch_size: (i + 1) * batch_size]
+#                 np.random.shuffle(current_bucket)
+#                 yield [self.data_info[idx]["index"] for idx in current_bucket]
+
+#             remaining = len(batch_indices) % batch_size
+#             if remaining > 0 and (remaining != 1 or not self.drop_unique):
+#                 current_bucket = batch_indices[-remaining:]
+#                 np.random.shuffle(current_bucket)
+#                 yield [self.data_info[idx]["index"] for idx in current_bucket]
+
+#     def __len__(self) -> int:
+#         return self.length
+
+
+
 class BucketSampler(Sampler):
     def __init__(self, dataset, batch_size, sort_key=lambda x, index_1, index_2: max(len(x[index_1]), len(x[index_2])), input_key: Union[str, int] = 0, label_key: Union[str, int] = 1):
         self.dataset = dataset

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: wolof-translate
-Version: 0.0.1
+Version: 0.0.3
 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.1.dist-info → wolof_translate-0.0.3.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=Hglii1Hj6H_K51JunTjUAxuLd4ehPb6LeeMVhsmhNxQ,6248
+wolof_translate/utils/bucket_iterator.py,sha256=JctNzZI1yU1FjZUKY-zNiRMXS3JpRiXOOSKdObRdPbg,8511
 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.1.dist-info/METADATA,sha256=itQMCA-zGM3gSDiKco5dMtY5qfpPHvKNJbw9KtffxzI,818
-wolof_translate-0.0.1.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-wolof_translate-0.0.1.dist-info/top_level.txt,sha256=YG-kBnOwUZyQ7SofNvMxNYjzCreH2PVcW2UaEg1-Reg,16
-wolof_translate-0.0.1.dist-info/RECORD,,
+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,,