Merge pull request #845 from PrefectHQ/mapping

Add mapping for automatic parallelization over inputs

docs/docs/text/classification.md

@@ -187,4 +187,21 @@ result = await marvin.classify_async(
 ) 
 
 assert result == "bug"
-```
+```
+
+## Mapping
+
+To classify a list of inputs at once, use `.map`:
+
+```python
+inputs = [
+    "The app crashes when I try to upload a file.",
+    "How do change my password?"
+]
+result = marvin.classify.map(inputs, ["bug", "feature request", "inquiry"])
+assert result == ["bug", "inquiry"]
+```
+
+(`marvin.classify_async.map` is also available for async environments.)
+
+Mapping automatically issues parallel requests to the API, making it a highly efficient way to classify multiple inputs at once. The result is a list of classifications in the same order as the inputs.

docs/docs/text/extraction.md

@@ -167,4 +167,21 @@ result = await marvin.extract_async(
 ) 
 
 assert result == ["NY", "CA"]
-```
+```
+
+## Mapping
+
+To extract from a list of inputs at once, use `.map`:
+
+```python
+inputs = [
+    "I drove from New York to California.",
+    "I took a flight from NYC to BOS."
+]
+result = marvin.extract.map(inputs, target=str, instructions="2-letter state codes")
+assert result  == [["NY", "CA"], ["NY", "MA"]]
+```
+
+(`marvin.extract_async.map` is also available for async environments.)
+
+Mapping automatically issues parallel requests to the API, making it a highly efficient way to work with multiple inputs at once. The result is a list of outputs in the same order as the inputs.

docs/docs/text/transformation.md

@@ -123,4 +123,21 @@ If you are using `marvin` in an async environment, you can use `cast_async`:
 result = await marvin.cast_async("one", int) 
 
 assert result == 1
-```
+```
+
+## Mapping
+
+To transform a list of inputs at once, use `.map`:
+
+```python
+inputs = [
+    "I bought two donuts.",
+    "I bought six hot dogs."
+]
+result = marvin.cast.map(inputs, int)
+assert result  == [2, 6]
+```
+
+(`marvin.cast_async.map` is also available for async environments.)
+
+Mapping automatically issues parallel requests to the API, making it a highly efficient way to work with multiple inputs at once. The result is a list of outputs in the same order as the inputs.

docs/docs/vision/classification.md

@@ -74,4 +74,21 @@ result = await marvin.beta.classify_async(
 ) 
 
 assert result == "bug"
-```
+```
+
+## Mapping
+
+To classify a list of inputs at once, use `.map`:
+
+```python
+inputs = [
+    "The app crashes when I try to upload a file.",
+    "How do change my password?"
+]
+result = marvin.beta.classify.map(inputs, ["bug", "feature request", "inquiry"])
+assert result == ["bug", "inquiry"]
+```
+
+(`marvin.beta.classify_async.map` is also available for async environments.)
+
+Mapping automatically issues parallel requests to the API, making it a highly efficient way to classify multiple inputs at once. The result is a list of classifications in the same order as the inputs.

docs/docs/vision/extraction.md

@@ -64,4 +64,21 @@ result = await marvin.beta.extract_async(
 ) 
 
 assert result == ["NY", "CA"]
-```
+```
+
+## Mapping
+
+To extract from a list of inputs at once, use `.map`:
+
+```python
+inputs = [
+    "I drove from New York to California.",
+    "I took a flight from NYC to BOS."
+]
+result = marvin.beta.extract.map(inputs, target=str, instructions="2-letter state codes")
+assert result  == [["NY", "CA"], ["NY", "MA"]]
+```
+
+(`marvin.beta.extract_async.map` is also available for async environments.)
+
+Mapping automatically issues parallel requests to the API, making it a highly efficient way to work with multiple inputs at once. The result is a list of outputs in the same order as the inputs.

docs/docs/vision/transformation.md

@@ -124,4 +124,21 @@ If you are using `marvin` in an async environment, you can use `cast_async`:
 result = await marvin.beta.cast_async("one", int) 
 
 assert result == 1
-```
+```
+
+## Mapping
+
+To transform a list of inputs at once, use `.map`:
+
+```python
+inputs = [
+    "I bought two donuts.",
+    "I bought six hot dogs."
+]
+result = marvin.beta.cast.map(inputs, int)
+assert result  == [2, 6]
+```
+
+(`marvin.beta.cast_async.map` is also available for async environments.)
+
+Mapping automatically issues parallel requests to the API, making it a highly efficient way to work with multiple inputs at once. The result is a list of outputs in the same order as the inputs.

src/marvin/ai/text.py

@@ -40,6 +40,7 @@
 from marvin.utilities.context import ctx
 from marvin.utilities.jinja import Transcript
 from marvin.utilities.logging import get_logger
+from marvin.utilities.mapping import map_async
 from marvin.utilities.python import PythonFunction
 from marvin.utilities.strings import count_tokens
 
@@ -680,7 +681,7 @@ def __init__(self, *args, **kwargs):
     return decorator
 
 
-### Sync versions of the above functions
+# --- Sync versions of the above functions
 
 
 def cast(
@@ -839,3 +840,123 @@ def generate(
             client=client,
         )
     )
+
+
+# --- Mapping
+async def classify_async_map(
+    data: list[str],
+    labels: Union[Enum, list[T], type],
+    instructions: Optional[str] = None,
+    model_kwargs: Optional[dict] = None,
+    client: Optional[AsyncMarvinClient] = None,
+) -> list[T]:
+    return await map_async(
+        fn=classify_async,
+        map_kwargs=dict(data=data),
+        unmapped_kwargs=dict(
+            labels=labels,
+            instructions=instructions,
+            model_kwargs=model_kwargs,
+            client=client,
+        ),
+    )
+
+
+def classify_map(
+    data: list[str],
+    labels: Union[Enum, list[T], type],
+    instructions: Optional[str] = None,
+    model_kwargs: Optional[dict] = None,
+    client: Optional[AsyncMarvinClient] = None,
+) -> list[T]:
+    return run_sync(
+        classify_async_map(
+            data=data,
+            labels=labels,
+            instructions=instructions,
+            model_kwargs=model_kwargs,
+            client=client,
+        )
+    )
+
+
+async def cast_async_map(
+    data: list[str],
+    target: type[T],
+    instructions: Optional[str] = None,
+    model_kwargs: Optional[dict] = None,
+    client: Optional[AsyncMarvinClient] = None,
+) -> list[T]:
+    return await map_async(
+        fn=cast_async,
+        map_kwargs=dict(data=data),
+        unmapped_kwargs=dict(
+            target=target,
+            instructions=instructions,
+            model_kwargs=model_kwargs,
+            client=client,
+        ),
+    )
+
+
+def cast_map(
+    data: list[str],
+    target: type[T],
+    instructions: Optional[str] = None,
+    model_kwargs: Optional[dict] = None,
+    client: Optional[AsyncMarvinClient] = None,
+) -> list[T]:
+    return run_sync(
+        cast_async_map(
+            data=data,
+            target=target,
+            instructions=instructions,
+            model_kwargs=model_kwargs,
+            client=client,
+        )
+    )
+
+
+async def extract_async_map(
+    data: list[str],
+    target: Optional[type[T]] = None,
+    instructions: Optional[str] = None,
+    model_kwargs: Optional[dict] = None,
+    client: Optional[AsyncMarvinClient] = None,
+) -> list[list[T]]:
+    return await map_async(
+        fn=extract_async,
+        map_kwargs=dict(data=data),
+        unmapped_kwargs=dict(
+            target=target,
+            instructions=instructions,
+            model_kwargs=model_kwargs,
+            client=client,
+        ),
+    )
+
+
+def extract_map(
+    data: list[str],
+    target: Optional[type[T]] = None,
+    instructions: Optional[str] = None,
+    model_kwargs: Optional[dict] = None,
+    client: Optional[AsyncMarvinClient] = None,
+) -> list[list[T]]:
+    return run_sync(
+        extract_async_map(
+            data=data,
+            target=target,
+            instructions=instructions,
+            model_kwargs=model_kwargs,
+            client=client,
+        )
+    )
+
+
+cast_async.map = cast_async_map
+cast.map = cast_map
+classify_async.map = classify_async_map
+classify.map = classify_map
+extract_async.map = extract_async_map
+extract.map = extract_map