Metadata-Version: 2.1
Name: datachain
Version: 0.2.7
Summary: Wrangle unstructured AI data at scale
Author-email: Dmitry Petrov <support@dvc.org>
License: Apache-2.0
Project-URL: Documentation, https://datachain.dvc.ai
Project-URL: Issues, https://github.com/iterative/dvcx/issues
Project-URL: Source, https://github.com/iterative/dvcx
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Development Status :: 2 - Pre-Alpha
Requires-Python: >=3.9
Description-Content-Type: text/x-rst
License-File: LICENSE
Requires-Dist: pyyaml
Requires-Dist: tomlkit
Requires-Dist: tqdm
Requires-Dist: numpy
Requires-Dist: pandas >=2.0.0
Requires-Dist: pyarrow
Requires-Dist: typing-extensions
Requires-Dist: python-dateutil >=2
Requires-Dist: attrs >=21.3.0
Requires-Dist: s3fs >=2024.2.0
Requires-Dist: gcsfs >=2024.2.0
Requires-Dist: adlfs >=2024.2.0
Requires-Dist: dvc-data <4,>=3.10
Requires-Dist: dvc-objects <6,>=4
Requires-Dist: shtab <2,>=1.3.4
Requires-Dist: sqlalchemy >=2
Requires-Dist: multiprocess ==0.70.16
Requires-Dist: dill ==0.3.8
Requires-Dist: ujson >=5.9.0
Requires-Dist: pydantic <3,>=2
Requires-Dist: jmespath >=1.0
Requires-Dist: datamodel-code-generator >=0.25
Requires-Dist: numpy <2,>=1 ; sys_platform == "win32"
Provides-Extra: cv
Requires-Dist: Pillow <11,>=10.0.0 ; extra == 'cv'
Requires-Dist: torch >=2.1.0 ; extra == 'cv'
Requires-Dist: torchvision ; extra == 'cv'
Requires-Dist: transformers >=4.36.0 ; extra == 'cv'
Provides-Extra: dev
Requires-Dist: datachain[docs,tests] ; extra == 'dev'
Requires-Dist: mypy ==1.10.1 ; extra == 'dev'
Requires-Dist: types-python-dateutil ; extra == 'dev'
Requires-Dist: types-PyYAML ; extra == 'dev'
Requires-Dist: types-requests ; extra == 'dev'
Requires-Dist: types-ujson ; extra == 'dev'
Provides-Extra: docs
Requires-Dist: mkdocs >=1.5.2 ; extra == 'docs'
Requires-Dist: mkdocs-gen-files >=0.5.0 ; extra == 'docs'
Requires-Dist: mkdocs-material >=9.3.1 ; extra == 'docs'
Requires-Dist: mkdocs-section-index >=0.3.6 ; extra == 'docs'
Requires-Dist: mkdocstrings-python >=1.6.3 ; extra == 'docs'
Requires-Dist: mkdocs-literate-nav >=0.6.1 ; extra == 'docs'
Provides-Extra: remote
Requires-Dist: lz4 ; extra == 'remote'
Requires-Dist: msgpack <2,>=1.0.4 ; extra == 'remote'
Requires-Dist: requests >=2.22.0 ; extra == 'remote'
Provides-Extra: tests
Requires-Dist: datachain[cv,remote,vector] ; extra == 'tests'
Requires-Dist: pytest <9,>=8 ; extra == 'tests'
Requires-Dist: pytest-sugar >=0.9.6 ; extra == 'tests'
Requires-Dist: pytest-cov >=4.1.0 ; extra == 'tests'
Requires-Dist: pytest-mock >=3.12.0 ; extra == 'tests'
Requires-Dist: pytest-servers[all] >=0.5.4 ; extra == 'tests'
Requires-Dist: pytest-benchmark[histogram] ; extra == 'tests'
Requires-Dist: pytest-asyncio >=0.23.2 ; extra == 'tests'
Requires-Dist: pytest-xdist >=3.3.1 ; extra == 'tests'
Requires-Dist: virtualenv ; extra == 'tests'
Requires-Dist: dulwich ; extra == 'tests'
Requires-Dist: hypothesis ; extra == 'tests'
Requires-Dist: open-clip-torch ; extra == 'tests'
Requires-Dist: aiotools >=1.7.0 ; extra == 'tests'
Requires-Dist: requests-mock ; extra == 'tests'
Provides-Extra: vector
Requires-Dist: usearch ; extra == 'vector'

|PyPI| |Python Version| |Codecov| |Tests|

.. |PyPI| image:: https://img.shields.io/pypi/v/datachain.svg
   :target: https://pypi.org/project/datachain/
   :alt: PyPI
.. |Python Version| image:: https://img.shields.io/pypi/pyversions/datachain
   :target: https://pypi.org/project/datachain
   :alt: Python Version
.. |Codecov| image:: https://codecov.io/gh/iterative/dvcx/branch/main/graph/badge.svg?token=VSCP2T9R5X
   :target: https://app.codecov.io/gh/iterative/dvcx
   :alt: Codecov
.. |Tests| image:: https://github.com/iterative/dvcx/workflows/Tests/badge.svg
   :target: https://github.com/iterative/dvcx/actions?workflow=Tests
   :alt: Tests

AI 🔗 DataChain
----------------

DataChain is an open-source Python data processing library for wrangling unstructured AI data at scale.

Datachain enables multimodal API calls and local AI inferences to run in parallel over many samples as chained operations. The resulting datasets can be saved, versioned, and sent directly to PyTorch and TensorFlow for training. Datachain can persist features of Python objects returned by AI models, and enables vectorized analytical operations over them.

The typical use cases are data curation, LLM analytics and validation, image segmentation, pose detection, and GenAI alignment. Datachain is especially helpful if batch operations can be optimized – for instance, when synchronous API calls can be parallelized  or where an LLM API offers batch processing.

.. code:: console

   $ pip install datachain

Operation basics
----------------

DataChain is built by composing wrangling operations.

For example, let us consider a dataset from Karlsruhe Institute of Technology detailing dialogs between users and customer service chatbots. We can use the chain to read data from the cloud, map it onto the parallel API calls for LLM evaluation, and organize the output into a dataset :

.. code:: py

      # pip install mistralai
      # this example requires a free Mistral API key, get yours at https://console.mistral.ai
      # add the key to your shell environment: $ export MISTRAL_API_KEY= your key

      # pip install mistralai
      # this example requires a free Mistral API key, get yours at https://console.mistral.ai
      # add the key to your shell environment: $ export MISTRAL_API_KEY= your key

      import os

      from mistralai.client import MistralClient
      from mistralai.models.chat_completion import ChatMessage

      from datachain.lib.dc import DataChain, Column

      PROMPT = "Was this bot dialog successful? Describe the 'result' as 'Yes' or 'No' in a short JSON"

      model = "mistral-large-latest"
      api_key = os.environ["MISTRAL_API_KEY"]

      chain = (
          DataChain.from_storage("gs://datachain-demo/chatbot-KiT/")
          .limit(5)
          .settings(cache=True, parallel=5)
          .map(
              mistral_response=lambda file: MistralClient(api_key=api_key)
              .chat(
                  model=model,
                  response_format={"type": "json_object"},
                  messages=[
                      ChatMessage(role="user", content=f"{PROMPT}: {file.get_value()}")
                  ],
              )
              .choices[0]
              .message.content,
          )
          .save()
      )

      try:
          print(chain.select("mistral_response").results())
      except Exception as e:
          print(f"do you have the right Mistral API key? {e}")


.. code:: shell

      [('{"result": "Yes"}',), ('{"result": "No"}',), ... , ('{"result": "Yes"}',)]

Now we have parallel-processed an LLM API-based query over cloud data and persisted the results.

Vectorized analytics
--------------------

Datachain internally represents datasets as tables, so analytical queries on the chain are automatically vectorized:

.. code:: py

      failed_dialogs = chain.filter(Column("mistral_response") == '{"result": "No"}')
      success_rate = failed_dialogs.count() / chain.count()
      print(f"Chatbot dialog success rate: {100*success_rate:.2f}%")


.. code:: shell

      "40.00%"

Note that DataChain represents file samples as pointers into their respective storage locations. This means a newly created dataset version does not duplicate files in storage, and storage remains the single source of truth for the original samples

Handling Python objects
-----------------------
In addition to storing primitive Python data types, chain is also capable of using data models.

For example, instead of collecting just a text response from Mistral API, we might be interested in more fields of the Mistral response object. For this task, we can define a Pydantic-like model and populate it from the API replies:

.. code:: py

      import os

      from mistralai.client import MistralClient
      from mistralai.models.chat_completion import ChatMessage

      from datachain.lib.dc import DataChain
      from datachain.lib.feature import Feature


      PROMPT = (
          "Was this dialog successful? Describe the 'result' as 'Yes' or 'No' in a short JSON"
      )

      model = "mistral-large-latest"
      api_key = os.environ["MISTRAL_API_KEY"]


      ## define the data model ###
      class Usage(Feature):
          prompt_tokens: int = 0
          completion_tokens: int = 0


      class MyChatMessage(Feature):
          role: str = ""
          content: str = ""


      class CompletionResponseChoice(Feature):
          message: MyChatMessage = MyChatMessage()


      class MistralModel(Feature):
          id: str = ""
          choices: list[CompletionResponseChoice]
          usage: Usage = Usage()


      ## Populate model instances ###
      chain = (
          DataChain.from_storage("gs://datachain-demo/chatbot-KiT/")
          .limit(5)
          .settings(cache=True, parallel=5)
          .map(
              mistral_response=lambda file: MistralModel(
                  **MistralClient(api_key=api_key)
                  .chat(
                      model=model,
                      response_format={"type": "json_object"},
                      messages=[
                          ChatMessage(role="user", content=f"{PROMPT}: {file.get_value()}")
                      ],
                  )
                  .dict()
              ),
              output=MistralModel,
          )
          .save("dialog-eval")
      )

After the chain execution, we can collect the objects:

.. code:: py

      for obj in responses:
          assert isinstance(obj, MistralModel)
          print(obj.dict())

.. code:: shell

      {'choices': [{'message': {'role': 'assistant', 'content': '{"result": "Yes"}'}}], 'usage': {'prompt_tokens': 610, 'completion_tokens': 6}}
      {'choices': [{'message': {'role': 'assistant', 'content': '{"result": "No"}'}}], 'usage': {'prompt_tokens': 3983, 'completion_tokens': 6}}
      {'choices': [{'message': {'role': 'assistant', 'content': '{"result": "Yes"}'}}], 'usage': {'prompt_tokens': 706, 'completion_tokens': 6}}
      {'choices': [{'message': {'role': 'assistant', 'content': '{"result": "No"}'}}], 'usage': {'prompt_tokens': 1250, 'completion_tokens': 6}}
      {'choices': [{'message': {'role': 'assistant', 'content': '{"result": "Yes"}'}}], 'usage': {'prompt_tokens': 1217, 'completion_tokens': 6}}


Dataset persistence
--------------------

The “save” operation makes chain dataset persistent in the current (working) directory of the query. A hidden folder .datachain/ holds the records. A persistent dataset can be accessed later to start a derivative chain:

.. code:: py

         DataChain.from_dataset("dialog-eval").limit(2).save("dialog-eval")

Persistent datasets are immutable and automatically versioned. Versions can be listed from shell:

.. code:: shell

      $ datachain ls-datasets

      dialog-rate (v1)
      dialog-rate (v2)

By default, when a persistent dataset is loaded, the latest version is fetched but another version can be requested:

.. code:: py

      ds = DataChain.from_dataset("dialog-eval", version = 1)

Chain optimization and execution
--------------------------------

Datachain avoids redundant operations. Execution is triggered only when a downstream operation requests the processed results. However, it would be inefficient to run, say, LLM queries again every time you just want to collect several objects.

“Save” operation nails execution results and automatically refers to them every time the downstream functions ask for data. Saving without an explicit name generates an auto-named dataset which serves the same purpose.


Matching data with metadata
----------------------------
It is common for AI data to come with pre-computed metadata (annotations, classes, etc).

DataChain library understands common metadata formats (JSON, CSV and parquet), and can unite data samples from storage with side-loaded metadata. The schema for metadata can be set explicitly or be inferred.

Here is an example of reading a CSV file where schema is heuristically derived from the header:

.. code:: py

      from datachain.lib.dc import DataChain
      csv_dataset = DataChain.from_csv("gs://datachain-demo/chatbot-csv/")

      print(csv_dataset.to_pandas())

Reading metadata from JSON format is a more complicated scenario because a JSON-annotated dataset typically references data samples (e.g. images) in annotation arrays somewhere within JSON files.

Here is an example from MS COCO “captions” JSON which employs separate sections for image meta and captions:

.. code:: json


      {
        "images": [
          {
            "license": 4,
            "file_name": "000000397133.jpg",
            "coco_url": "http://images.cocodataset.org/val2017/000000397133.jpg",
            "height": 427,
            "width": 640,
            "date_captured": "2013-11-14 17:02:52",
            "flickr_url": "http://farm7.staticflickr.com/6116/6255196340_da26cf2c9e_z.jpg",
            "id": 397133
          },
          ...
        ],
        "annotations": [
          {
            "image_id"	:	"179765",
            "id"	:	38,
            "caption"	:	"A black Honda motorcycle parked in front of a garage."
          },
          ...
        ],
        ...
      }

To deal with this layout, we can take the following steps:

1. Generate a dataset of raw image files from storage
2. Generate a meta-information dataset from the JSON section “images”
3. Join these datasets via the matching id keys

.. code:: python


   from datachain.lib.dc import DataChain

   images = DataChain.from_storage("gs://datachain-demo/coco2017/images/val/")
   meta = DataChain.from_json("gs://datachain-demo/coco2017/annotations_captions", jmespath = "images")

   images_with_meta = images.merge(meta, on="file.name", right_on="images.file_name")

   print(images_with_meta.limit(1).results())

.. code:: shell


      Processed: 5000 rows [00:00, 15481.66 rows/s]
      Processed: 1 rows [00:00, 1291.75 rows/s]
      Processed: 1 rows [00:00,  4.70 rows/s]
      Generated: 5000 rows [00:00, 27128.67 rows/s]
      [(1, 2336066478558845549, '', 0, 'coco2017/images/val', '000000000139.jpg', 'CNvXoemj8IYDEAE=', '1719096046021595', 1, datetime.datetime(2024, 6, 22, 22, 40, 46, 70000, tzinfo=datetime.timezone.utc), 161811, '', '', None, 'gs://datachain-demo', 'gs://datachain-demo', 'coco2017/images/val', '000000000139.jpg', 161811, '1719096046021595', 'CNvXoemj8IYDEAE=', 1, datetime.datetime(1970, 1, 1, 0, 0, tzinfo=datetime.timezone.utc), None, '', 4146, 6967063844996569113, 2, '000000000139.jpg', 'http://images.cocodataset.org/val2017/000000000139.jpg', 426, 640, '2013-11-21 01:34:01', 'http://farm9.staticflickr.com/8035/8024364858_9c41dc1666_z.jpg', 139)]

Passing data to training
------------------------

Chain results can be exported or passed directly to Pytorch dataloader. For example, if we are interested in passing three columns to training, the following Pytorch code will do it:

.. code:: py

      ds = train.select("file", "caption_choices", "label_ind").to_pytorch(
          transform=preprocess,
          tokenizer=clip.tokenize,
      )

      loader = DataLoader(ds, batch_size=2)
      optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
      train(loader, model, optimizer)

Tutorials
------------------

* `Multimodal <examples/multimodal/clip_fine_tuning.ipynb>`_ (try in `Colab <https://colab.research.google.com/github/iterative/dvclive/blob/main/examples/multimodal/clip_fine_tuning.ipynb>`__)

Contributions
--------------------

Contributions are very welcome.
To learn more, see the `Contributor Guide`_.


License
-------

Distributed under the terms of the `Apache 2.0 license`_,
*DataChain* is free and open source software.


Issues
------

If you encounter any problems,
please `file an issue`_ along with a detailed description.


.. _Apache 2.0 license: https://opensource.org/licenses/Apache-2.0
.. _PyPI: https://pypi.org/
.. _file an issue: https://github.com/iterative/dvcx/issues
.. _pip: https://pip.pypa.io/
.. github-only
.. _Contributor Guide: CONTRIBUTING.rst
.. _Pydantic: https://github.com/pydantic/pydantic
