Limit categories in output multinouli distribution

metasyn/distribution/categorical.py

@@ -68,6 +68,25 @@ def _param_schema(cls):
             "probs": {"type": "array", "items": {"type": "number"}},
         }
 
+    @property
+    def _params_formatted(self) -> str:
+        n_categories = len(self.labels)
+        top_bottom_counts = (7, 3)
+        if n_categories <= top_bottom_counts[0]:
+            return super()._params_formatted
+        sorted_indices = np.argsort(-self.probs)
+        top_indices = sorted_indices[:top_bottom_counts[0]]
+        bottom_indices = sorted_indices[-top_bottom_counts[1]:]
+        top_labels = ", ".join(repr(label) for label in self.labels[top_indices].tolist())
+        top_probs = ", ".join(f"{prob:.3f}" for prob in self.probs[top_indices])
+        bottom_labels = ", ".join(repr(label) for label in self.labels[bottom_indices].tolist())
+        bottom_probs = ", ".join(f"{prob:.3f}" for prob in self.probs[bottom_indices])
+        return (
+            f"\t- labels: [{top_labels}, ..., {bottom_labels}]\n"
+            f"\t- probs: [{top_probs}, ..., {bottom_probs}]\n"
+            f"\tShowing {top_bottom_counts[0]+top_bottom_counts[1]} of {n_categories} categories"
+        )
+
     def draw(self):
         return np.random.choice(self.labels, p=self.probs)