Style fix for dspy/datasets (#8175)

* enable style check

* update rules

* fix style for dspy/datasets

Author: chenmoneygithub

dspy/clients/lm.py

@@ -298,7 +298,8 @@ async def async_stream_completion():
         return async_stream_completion
 
 
-def litellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cache": True, "no-store": True}):
+def litellm_completion(request: Dict[str, Any], num_retries: int, cache: Optional[Dict[str, Any]] = None):
+    cache = cache or {"no-cache": True, "no-store": True}
     stream_completion = _get_stream_completion_fn(request, cache, sync=True)
     if stream_completion is None:
         return litellm.completion(
@@ -311,7 +312,8 @@ def litellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cac
     return stream_completion()
 
 
-def litellm_text_completion(request: Dict[str, Any], num_retries: int, cache={"no-cache": True, "no-store": True}):
+def litellm_text_completion(request: Dict[str, Any], num_retries: int, cache: Optional[Dict[str, Any]] = None):
+    cache = cache or {"no-cache": True, "no-store": True}
     # Extract the provider and model from the model string.
     # TODO: Not all the models are in the format of "provider/model"
     model = request.pop("model").split("/", 1)
@@ -336,7 +338,8 @@ def litellm_text_completion(request: Dict[str, Any], num_retries: int, cache={"n
     )
 
 
-async def alitellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cache": True, "no-store": True}):
+async def alitellm_completion(request: Dict[str, Any], num_retries: int, cache: Optional[Dict[str, Any]] = None):
+    cache = cache or {"no-cache": True, "no-store": True}
     stream_completion = _get_stream_completion_fn(request, cache, sync=False)
     if stream_completion is None:
         return await litellm.acompletion(
@@ -349,9 +352,8 @@ async def alitellm_completion(request: Dict[str, Any], num_retries: int, cache={
     return await stream_completion()
 
 
-async def alitellm_text_completion(
-    request: Dict[str, Any], num_retries: int, cache={"no-cache": True, "no-store": True}
-):
+async def alitellm_text_completion(request: Dict[str, Any], num_retries: int, cache: Optional[Dict[str, Any]] = None):
+    cache = cache or {"no-cache": True, "no-store": True}
     model = request.pop("model").split("/", 1)
     provider, model = model[0] if len(model) > 1 else "openai", model[-1]
 

dspy/datasets/__init__.py

@@ -1,9 +1,9 @@
+from dspy.datasets.alfworld import AlfWorld
 from dspy.datasets.colors import Colors
 from dspy.datasets.dataloader import DataLoader
 from dspy.datasets.dataset import Dataset
 from dspy.datasets.hotpotqa import HotPotQA
 from dspy.datasets.math import MATH
-from dspy.datasets.alfworld import AlfWorld
 
 __all__ = [
     "Colors",

dspy/datasets/alfworld/__init__.py

@@ -1 +1 @@
-from dspy.datasets.alfworld.alfworld import AlfWorld
+from dspy.datasets.alfworld.alfworld import AlfWorld

dspy/datasets/alfworld/alfworld.py

@@ -2,6 +2,7 @@
 import queue
 import random
 
+
 def env_worker(inq, outq):
     """
     Worker process: creates a single AlfredTWEnv instance,
@@ -10,16 +11,19 @@ def env_worker(inq, outq):
 
     try:
         import io
-        import yaml
+        from contextlib import redirect_stderr, redirect_stdout
+
         import alfworld.agents.environment as environment
-        from contextlib import redirect_stdout, redirect_stderr
+        import yaml
     except ImportError:
-        raise ImportError("alfworld is not installed. " \
-            "Please install it via `pip install alfworld==0.3.5` then run `alfworld-download`.")
+        raise ImportError(
+            "alfworld is not installed. "
+            "Please install it via `pip install alfworld==0.3.5` then run `alfworld-download`."
+        )
 
     buf = io.StringIO()
     base_dir = os.path.dirname(os.path.abspath(__file__))
-    config_path = os.path.join(base_dir, 'base_config.yml')
+    config_path = os.path.join(base_dir, "base_config.yml")
 
     with open(config_path) as f:
         config = yaml.safe_load(f)
@@ -30,19 +34,19 @@ def env_worker(inq, outq):
     env = None
     while True:
         cmd, data = inq.get()
-        if cmd == 'init':
+        if cmd == "init":
             env = base_env.init_env(batch_size=1)
             env.skip(data)
             task_def, info = env.reset()
             outq.put((task_def[0], info))
-        elif cmd == 'step':
+        elif cmd == "step":
             obs, rew, done, info = env.step([data])
             outq.put((obs, rew, done, info))
-        elif cmd == 'close':
-            outq.put('CLOSED')
+        elif cmd == "close":
+            outq.put("CLOSED")
             break
         else:
-            outq.put('UNKNOWN_CMD')
+            outq.put("UNKNOWN_CMD")
 
 
 class EnvPool:
@@ -54,6 +58,7 @@ class EnvPool:
             obs, rew, done, info = sess.step("go north")
             ...
     """
+
     def __init__(self, size=2):
         self.size = size
         self.workers = []
@@ -62,8 +67,7 @@ def __init__(self, size=2):
         try:
             import multiprocess as mp
         except ImportError:
-            raise ImportError("multiprocess is not installed. " \
-                "Please install it via `pip install multiprocess`.")
+            raise ImportError("multiprocess is not installed. " "Please install it via `pip install multiprocess`.")
 
         # Must call set_start_method('spawn') here, before creating any processes
         try:
@@ -93,7 +97,7 @@ def close_all(self):
         while not self.available.empty():
             wid = self.available.get()
             inq, outq, proc = self.workers[wid]
-            inq.put(('close', None))
+            inq.put(("close", None))
             outq.get()  # Wait 'CLOSED'
             inq.close()
             outq.close()
@@ -109,6 +113,7 @@ class _EnvSession:
     A context manager that acquires a worker from the pool,
     provides .init(idx) and .step(action), then releases the worker.
     """
+
     def __init__(self, pool: EnvPool):
         self.pool = pool
         self.wid = None
@@ -123,11 +128,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self.pool._release(self.wid)
 
     def init(self, idx):
-        self.inq.put(('init', idx))
+        self.inq.put(("init", idx))
         return self.outq.get()  # (task_def, info)
 
     def step(self, action):
-        self.inq.put(('step', action))
+        self.inq.put(("step", action))
         return self.outq.get()  # (obs, rew, done, info)
 
 
@@ -136,7 +141,8 @@ def __init__(self, max_threads=20):
         self.POOL = EnvPool(size=max_threads)
 
         import dspy
-        dataset = [dspy.Example(idx=idx).with_inputs('idx') for idx in range(3500)]
+
+        dataset = [dspy.Example(idx=idx).with_inputs("idx") for idx in range(3500)]
         random.Random(0).shuffle(dataset)
 
         trainset, devset = dataset[:3000], dataset[-500:]

dspy/datasets/colors.py

@@ -3,7 +3,146 @@
 from dspy.datasets.dataset import Dataset
 
 ### A bunch of colors, originally from matplotlib
-all_colors = ['alice blue', 'dodger blue', 'light sky blue', 'deep sky blue', 'sky blue', 'steel blue', 'light steel blue', 'medium blue', 'navy blue', 'blue', 'royal blue', 'cadet blue', 'cornflower blue', 'medium slate blue', 'slate blue', 'dark slate blue', 'powder blue', 'turquoise', 'dark turquoise', 'medium turquoise', 'pale turquoise', 'light sea green', 'medium sea green', 'sea green', 'forest green', 'green yellow', 'lime green', 'dark green', 'green', 'lime', 'chartreuse', 'lawn green', 'yellow green', 'olive green', 'dark olive green', 'medium spring green', 'spring green', 'medium aquamarine', 'aquamarine', 'aqua', 'cyan', 'dark cyan', 'teal', 'medium orchid', 'dark orchid', 'orchid', 'blue violet', 'violet', 'dark violet', 'plum', 'thistle', 'magenta', 'fuchsia', 'dark magenta', 'medium purple', 'purple', 'rebecca purple', 'dark red', 'fire brick', 'indian red', 'light coral', 'dark salmon', 'light salmon', 'salmon', 'red', 'crimson', 'tomato', 'coral', 'orange red', 'dark orange', 'orange', 'yellow', 'gold', 'light goldenrod yellow', 'pale goldenrod', 'goldenrod', 'dark goldenrod', 'beige', 'moccasin', 'blanched almond', 'navajo white', 'antique white', 'bisque', 'burlywood', 'dark khaki', 'khaki', 'tan', 'wheat', 'snow', 'floral white', 'old lace', 'ivory', 'linen', 'seashell', 'honeydew', 'mint cream', 'azure', 'lavender', 'ghost white', 'white smoke', 'gainsboro', 'light gray', 'silver', 'dark gray', 'gray', 'dim gray', 'slate gray', 'light slate gray', 'dark slate gray', 'black', 'medium violet red', 'pale violet red', 'deep pink', 'hot pink', 'light pink', 'pink', 'peach puff', 'rosy brown', 'saddle brown', 'sandy brown', 'chocolate', 'peru', 'sienna', 'brown', 'maroon', 'white', 'misty rose', 'lavender blush', 'papaya whip', 'lemon chiffon', 'light yellow', 'corn silk', 'pale green', 'light green', 'olive drab', 'olive', 'dark sea green']
+all_colors = [
+    "alice blue",
+    "dodger blue",
+    "light sky blue",
+    "deep sky blue",
+    "sky blue",
+    "steel blue",
+    "light steel blue",
+    "medium blue",
+    "navy blue",
+    "blue",
+    "royal blue",
+    "cadet blue",
+    "cornflower blue",
+    "medium slate blue",
+    "slate blue",
+    "dark slate blue",
+    "powder blue",
+    "turquoise",
+    "dark turquoise",
+    "medium turquoise",
+    "pale turquoise",
+    "light sea green",
+    "medium sea green",
+    "sea green",
+    "forest green",
+    "green yellow",
+    "lime green",
+    "dark green",
+    "green",
+    "lime",
+    "chartreuse",
+    "lawn green",
+    "yellow green",
+    "olive green",
+    "dark olive green",
+    "medium spring green",
+    "spring green",
+    "medium aquamarine",
+    "aquamarine",
+    "aqua",
+    "cyan",
+    "dark cyan",
+    "teal",
+    "medium orchid",
+    "dark orchid",
+    "orchid",
+    "blue violet",
+    "violet",
+    "dark violet",
+    "plum",
+    "thistle",
+    "magenta",
+    "fuchsia",
+    "dark magenta",
+    "medium purple",
+    "purple",
+    "rebecca purple",
+    "dark red",
+    "fire brick",
+    "indian red",
+    "light coral",
+    "dark salmon",
+    "light salmon",
+    "salmon",
+    "red",
+    "crimson",
+    "tomato",
+    "coral",
+    "orange red",
+    "dark orange",
+    "orange",
+    "yellow",
+    "gold",
+    "light goldenrod yellow",
+    "pale goldenrod",
+    "goldenrod",
+    "dark goldenrod",
+    "beige",
+    "moccasin",
+    "blanched almond",
+    "navajo white",
+    "antique white",
+    "bisque",
+    "burlywood",
+    "dark khaki",
+    "khaki",
+    "tan",
+    "wheat",
+    "snow",
+    "floral white",
+    "old lace",
+    "ivory",
+    "linen",
+    "seashell",
+    "honeydew",
+    "mint cream",
+    "azure",
+    "lavender",
+    "ghost white",
+    "white smoke",
+    "gainsboro",
+    "light gray",
+    "silver",
+    "dark gray",
+    "gray",
+    "dim gray",
+    "slate gray",
+    "light slate gray",
+    "dark slate gray",
+    "black",
+    "medium violet red",
+    "pale violet red",
+    "deep pink",
+    "hot pink",
+    "light pink",
+    "pink",
+    "peach puff",
+    "rosy brown",
+    "saddle brown",
+    "sandy brown",
+    "chocolate",
+    "peru",
+    "sienna",
+    "brown",
+    "maroon",
+    "white",
+    "misty rose",
+    "lavender blush",
+    "papaya whip",
+    "lemon chiffon",
+    "light yellow",
+    "corn silk",
+    "pale green",
+    "light green",
+    "olive drab",
+    "olive",
+    "dark sea green",
+]
+
 
 class Colors(Dataset):
     def __init__(self, sort_by_suffix=True, *args, **kwargs) -> None:
@@ -12,22 +151,24 @@ def __init__(self, sort_by_suffix=True, *args, **kwargs) -> None:
         self.sort_by_suffix = sort_by_suffix
         colors = self.sorted_by_suffix(all_colors)
 
-        train_size = int(len(colors) * 0.6) # chosen to ensure that similar colors aren't repeated between train and dev
+        train_size = int(
+            len(colors) * 0.6
+        )  # chosen to ensure that similar colors aren't repeated between train and dev
         train_colors, dev_colors = colors[:train_size], colors[train_size:]
 
-        self._train = [dict(color=color) for color in train_colors]
-        self._dev = [dict(color=color) for color in dev_colors]
+        self._train = [{"color": color} for color in train_colors]
+        self._dev = [{"color": color} for color in dev_colors]
 
         random.Random(0).shuffle(self._train)
         random.Random(0).shuffle(self._dev)
-    
+
     def sorted_by_suffix(self, colors):
         if not self.sort_by_suffix:
             return colors
 
         if isinstance(colors[0], str):
             sorted_colors = sorted(colors, key=lambda x: x[::-1])
         else:
-            sorted_colors = sorted(colors, key=lambda x: x['color'][::-1])
+            sorted_colors = sorted(colors, key=lambda x: x["color"][::-1])
 
         return sorted_colors