Add __all__ for some files

fgnt · Jul 11, 2023 · 510d6e0 · 510d6e0
1 parent 6fd2410
commit 510d6e0
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 3 deletions.
diff --git a/mms_msg/sampling/environment/noise.py b/mms_msg/sampling/environment/noise.py
@@ -4,6 +4,10 @@
 
 from mms_msg.sampling.utils.rng import get_rng_example
 
+__all__ = [
+    'sample_uniform_snr',
+    'UniformSNRSampler',
+]
 
 def sample_uniform_snr(example, *, min_snr: float = 20, max_snr: float = 30):
     example['snr'] = float(

diff --git a/mms_msg/sampling/environment/rir.py b/mms_msg/sampling/environment/rir.py
@@ -1,10 +1,13 @@
 from dataclasses import dataclass
-from pathlib import Path
 
 import lazy_dataset
-from lazy_dataset.database import JsonDatabase
 from mms_msg.databases.reverberation.sms_wsj import SMSWSJRIRDatabase
 
+__all__ = [
+    'sample_rirs',
+    'RIRSampler',
+]
+
 
 def sample_rirs(example: dict, *, rir_dataset: lazy_dataset.Dataset):
     # Assume the examples have a running index

diff --git a/mms_msg/sampling/pattern/classical/offset.py b/mms_msg/sampling/pattern/classical/offset.py
@@ -4,6 +4,16 @@
 from mms_msg.sampling.utils.rng import get_rng_example
 from mms_msg.sampling.utils.utils import update_num_samples
 
+__all__ = [
+    'assign_offset',
+    'sample_offsets_sms_wsj',
+    'SMSWSJOffsetSampler',
+    'sample_offsets_constant',
+    'ConstantOffsetSampler',
+    'sample_partial_overlap',
+    'PartialOverlapOffsetSampler',
+]
+
 
 def assign_offset(example, offset):
     assert keys.OFFSET not in example
@@ -56,7 +66,7 @@ def sample_partial_overlap(example, *, minimum_overlap, maximum_overlap):
     overlap = rng.uniform(minimum_overlap, maximum_overlap)
     num_samples = example[keys.NUM_SAMPLES][keys.ORIGINAL_SOURCE]
     assert len(num_samples) == 2, (len(num_samples), num_samples)
-    overlap_samples = sum(num_samples)*overlap / (1 + overlap)
+    overlap_samples = sum(num_samples) * overlap / (1 + overlap)
     offset = [0, int(max(num_samples[0] - overlap_samples, 0))]
     assign_offset(example, offset)
     return example

diff --git a/mms_msg/sampling/pattern/meeting/meeting_sampler.py b/mms_msg/sampling/pattern/meeting/meeting_sampler.py
@@ -18,6 +18,10 @@
 
 logger = logging.getLogger('meeting')
 
+__all__ = [
+    'MeetingSampler',
+    'sample_meeting_from_full_overlap',
+]
 
 @dataclass
 class _MeetingSampler:

diff --git a/mms_msg/sampling/pattern/meeting/scenario_sequence_sampler.py b/mms_msg/sampling/pattern/meeting/scenario_sequence_sampler.py
@@ -3,6 +3,17 @@
 from paderbox.utils.mapping import Dispatcher
 
 
+__all__ = [
+    'sample_random',
+    'sample_round_robin',
+    'sample_random_round_robin',
+    'sample_balanced',
+    'sample_balanced_no_repeat',
+    'sample_asymmetric',
+    'scenario_sequence_samplers',
+]
+
+
 def sample_random(scenarios: list, examples: list, rng: np.random.Generator):
     """
     Samples the scenarios uniformly, independent of the already sampled