updated get_samples by Tom S

src/get_labels_from_samples.py

@@ -1,19 +1,16 @@
 from segments import SegmentsClient
+import math
 
 def get_samples(client, dataset_identifier):
+    dataset = client.get_dataset(dataset_identifier)
+    samples_count = getattr(dataset, "samples_count", 0)
+    total_pages = math.ceil(samples_count / 1000)
     page = 1
     per_page = 1000
     samples = []
-    while True:
+    while page <= total_pages:
         response = client.get_samples(dataset_identifier, per_page=per_page, page=page)
-        # Handle both paginated and direct list responses
-        if isinstance(response, list):
-            samples.extend(response)
-            break  # No pagination in this case
-        else:
-            samples.extend(response.results)
-            if not response.has_next:
-                break
+        samples.extend(response)
         page += 1
     return samples