feat: cleaned up diversity calculator tool

Author: Wermutton

examples/diversity_calculation.py

@@ -1,5 +1,17 @@
 from langdiversity.utils import DiversityCalculator
+from langdiversity.measures import ShannonEntropyMeasure, GiniImpurityMeasure
 
-# Calculate diversity measures
-diversity_calculator = DiversityCalculator(default_measures=["entropy", "gini"])
-print("DIVERSITY MEASURES:", diversity_calculator.calculate([1, 1, 1, 2]))
+# Define a custom measure class
+class CustomMeasure:
+    def generate(self, values):
+        return sum(values) // len(values)
+
+# Create instances
+entropy = ShannonEntropyMeasure()
+gini = GiniImpurityMeasure()
+custom = CustomMeasure()
+
+# Use built-in and custom measures
+diversity_calculator = DiversityCalculator(measures=[entropy, gini, custom])
+
+print("Diversity Measures:", diversity_calculator.calculate([1, 1, 1, 2]))

langdiversity/utils/calculate_measures.py

@@ -1,24 +1,14 @@
-from ..measures import ShannonEntropyMeasure, GiniImpurityMeasure
+from langdiversity.extras.spinner import loading_spinner
 
 class DiversityCalculator:
-    # by default, include entropy as part of diversity measures
-    def __init__(self, default_measures=["entropy"]):
-        self.default_measures = default_measures
-
-    # returns diversity measure results for each question
-    def calculate(self, values, measures=None):
-        if measures is None:
-            measures = self.default_measures
+    def __init__(self, measures=[]):
+        self.measures = measures
 
+    def calculate(self, values):
         results = {}
-        
-        if "entropy" in measures:
-            entropy_measure = ShannonEntropyMeasure()
-            results["entropy"] = entropy_measure.generate(values)
-        
-        if "gini" in measures:
-            gini_measure = GiniImpurityMeasure()
-            results["gini"] = gini_measure.generate(values)
-        
+        for measure in self.measures:
+            display_name = measure.__class__.__name__
+            with loading_spinner(f"Calculating '{display_name}' for {len(values)} values..."):
+                results[display_name] = measure.generate(values)
+                
         return results
-    

setup.py

@@ -8,7 +8,7 @@
 setup(
     name='langdiversity',
     packages=find_packages(exclude=['tests']),
-    version='1.0.3',
+    version='1.0.5',
     description='A tool to elevate your language models with insightful diversity metrics.',
     long_description=long_description,  
     long_description_content_type="text/markdown",