feat: added loading effect when script runs

Author: Wermutton

langdiversity/extras/__init__.py

@@ -0,0 +1,30 @@
+import itertools
+import sys
+import threading
+import time
+from contextlib import contextmanager
+
+@contextmanager
+def loading_spinner(label, current_step=None, total_steps=None):
+    spinner_done = False
+
+    def spinner_function():
+        spinner = itertools.cycle(["-", "/", "|", "\\"])
+        while not spinner_done:
+            prefix = f"[{current_step}/{total_steps}] " if current_step is not None and total_steps is not None else ""
+            sys.stdout.write(f'\r{prefix}{label} {next(spinner)}')
+            sys.stdout.flush()
+            time.sleep(0.1)
+        print("\r", end="")
+        print(f"{prefix}{label} ✓")
+
+    spinner_thread = threading.Thread(target=spinner_function)
+    spinner_thread.start()
+
+    try:
+        yield
+    finally:
+        spinner_done = True
+        spinner_thread.join()
+
+

langdiversity/extras/spinner.py

@@ -1,5 +1,7 @@
 from typing import List
 
+from langdiversity.extras.spinner import loading_spinner
+
 from langdiversity.measures import AbstractMeasure
 from langdiversity.models import AbstractBaseModel
 
@@ -29,9 +31,13 @@ def generate(self, prompts: List[str]):
         selected_diversity = float("inf") if self.selection == "min" else float("-inf")
 
         info = []
-        for prompt in prompts:
-            responses = self.model.generate(prompt, self.num_responses)
-            diversity = self.diversity_measure.generate(responses)
+
+        total_prompts = len(prompts)
+        for i, prompt in enumerate(prompts):
+            with loading_spinner(f"Collecting {self.num_responses} responses...", current_step=i+1, total_steps=total_prompts):
+                responses = self.model.generate(prompt, self.num_responses)
+            with loading_spinner("Performing diversity measure calculations...", current_step=i+1, total_steps=total_prompts):
+                diversity = self.diversity_measure.generate(responses)
 
             if self.selection == "max" and diversity > selected_diversity:
                 selected_diversity = diversity

langdiversity/utils/prompt_selection.py

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