only show the distribution for top-10 tokens

sunlabuiuc · Nov 10, 2022 · 02ae89f · 02ae89f
1 parent ae67753
commit 02ae89f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py
@@ -501,8 +501,10 @@ def task_stat(self) -> str:
  f"{sum(num_events) / len(num_events):.4f}")
  lines.append(
  f"\t\t- Number of unique {key}: {len(self.get_all_tokens(key))}")
+ distribution = self.get_distribution_tokens(key)
+ top10 = sorted(distribution.items(), key=lambda x: x[1], reverse=True)[:10]
  lines.append(
- f"\t\t- Distribution of {key}: {self.get_distribution_tokens(key)}")
+ f"\t\t- Distribution of {key} (Top-10): {top10}")
  return "\n".join(lines)
 
  @staticmethod