Metadata-Version: 2.5
Name: distribatch
Version: 0.1.0a1
Summary: Batch LLM generation, embeddings and speech transcription over DataFrames with Spark or multiprocessing — Unity Catalog, local/Volume weights or OpenAI-compatible APIs
Project-URL: Homepage, https://github.com/byteway-pl/distribatch
Project-URL: Issues, https://github.com/byteway-pl/distribatch/issues
Author: Adrian Galik
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: batch,databricks,embeddings,inference,llm,spark,speech-to-text,unity-catalog,vllm,whisper
Classifier: Development Status :: 2 - Pre-Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.10
Requires-Dist: cloudpickle>=2.2
Requires-Dist: pandas>=1.5
Requires-Dist: pyarrow>=12
Requires-Dist: pydantic>=2.0
Provides-Extra: all
Requires-Dist: accelerate>=0.33; extra == 'all'
Requires-Dist: databricks-sdk>=0.30; extra == 'all'
Requires-Dist: faster-whisper>=1.0; extra == 'all'
Requires-Dist: mlflow-skinny>=2.16; extra == 'all'
Requires-Dist: openai>=1.40; extra == 'all'
Requires-Dist: pyspark>=3.5; extra == 'all'
Requires-Dist: scipy>=1.10; extra == 'all'
Requires-Dist: soundfile>=0.12; extra == 'all'
Requires-Dist: torch>=2.1; extra == 'all'
Requires-Dist: transformers>=4.44; extra == 'all'
Requires-Dist: vllm>=0.6; extra == 'all'
Provides-Extra: audio
Requires-Dist: scipy>=1.10; extra == 'audio'
Requires-Dist: soundfile>=0.12; extra == 'audio'
Provides-Extra: databricks
Requires-Dist: databricks-sdk>=0.30; extra == 'databricks'
Provides-Extra: dev
Requires-Dist: mypy>=1.10; extra == 'dev'
Requires-Dist: openai>=1.40; extra == 'dev'
Requires-Dist: pandas-stubs; extra == 'dev'
Requires-Dist: pyspark>=3.5; extra == 'dev'
Requires-Dist: pytest>=8; extra == 'dev'
Requires-Dist: ruff>=0.6; extra == 'dev'
Provides-Extra: mlflow
Requires-Dist: mlflow-skinny>=2.16; extra == 'mlflow'
Provides-Extra: openai
Requires-Dist: openai>=1.40; extra == 'openai'
Provides-Extra: spark
Requires-Dist: pyspark>=3.5; extra == 'spark'
Provides-Extra: transformers
Requires-Dist: accelerate>=0.33; extra == 'transformers'
Requires-Dist: torch>=2.1; extra == 'transformers'
Requires-Dist: transformers>=4.44; extra == 'transformers'
Provides-Extra: vllm
Requires-Dist: vllm>=0.6; extra == 'vllm'
Provides-Extra: whisper
Requires-Dist: faster-whisper>=1.0; extra == 'whisper'
Description-Content-Type: text/markdown

# distribatch

Batch model inference over DataFrames: **LLM generation, text embeddings and speech
transcription**, on Databricks or anywhere Python runs.

```python
import distribatch as dib
from pydantic import BaseModel

class Sentiment(BaseModel):
    label: str
    confidence: float

job = dib.GenerateJob(
    model=dib.models.databricks("databricks-meta-llama-3-3-70b-instruct"),
    prompt="Classify the sentiment of this review:\n{review}",
    output_schema=Sentiment,
)
scored = job.run(reviews_df, backend="spark", output_col="sentiment")
```

- **One job, three backends:** `local`, `multiprocessing` (one worker per GPU) and
  `spark` (`mapInArrow`, also on Spark Connect / serverless).
- **Any model source:** Unity Catalog registered models, Hugging Face weights in a
  directory or UC Volume, or any OpenAI-compatible API (OpenAI, Databricks Model Serving,
  vLLM servers, …).
- **Engines:** vLLM on GPUs, Hugging Face transformers on CPU/GPU, faster-whisper for
  speech, and an async OpenAI client with retries, rate limits and concurrency control.
- **Built for long runs:**
  - per-row errors instead of failed jobs;
  - structured output validated against pydantic models and typed as Spark structs;
  - checkpoints that resume where a run stopped;
  - deduplication of identical inputs;
  - run statistics with cost estimates.

## Installation

```bash
pip install "distribatch[openai]"                       # APIs and Databricks Model Serving
pip install "distribatch[vllm,mlflow]"                  # self-hosted LLMs/embedders from Unity Catalog
pip install "distribatch[whisper,transformers,audio]"   # local speech-to-text
pip install "distribatch[all]"
```

Python 3.10+. The core install is light: heavy dependencies come with the extras and are
imported only when used.

## Embeddings and transcription

```python
vectors = dib.EmbedJob(
    model=dib.models.path("/Volumes/main/ml/weights/bge-m3", copy_to_local=True),
    text="{title}\n\n{body}",
    normalize=True,
).run(docs_df, backend="spark", output_col="embedding")          # array<float>

calls = spark.read.format("binaryFile").load("/Volumes/main/raw/calls/").select("path")
transcripts = dib.TranscribeJob(
    model=dib.models.uc("main.audio.whisper_large_v3", alias="prod"),
    audio_col="path",
    timestamps=True,
).run(calls, backend="spark", output_col="transcript")
```

## Long runs

```python
out = job.run(
    reviews_df,
    backend="spark",
    checkpoint="/Volumes/main/ml/checkpoints/reviews",   # resume after interruptions
    id_col="review_id",
    dedupe=True,                                        # identical inputs sent once
)
stats = dib.summarize(out, output_col="sentiment")
print(stats)                                            # success rate, errors, tokens, latency
print(stats.cost(input_per_million=0.5, output_per_million=1.5))
```

On a single machine use `backend="multiprocessing"`. In plain Python scripts, put the
call under `if __name__ == "__main__":`, because worker processes are spawned.

## Documentation

The **[wiki](https://github.com/byteway-pl/distribatch/wiki)** covers every job, engine
option and backend, plus the output contract, checkpoints, a Databricks guide
(credentials, GPU clusters, Unity Catalog) and troubleshooting. Runnable examples are in
[`examples/`](examples/).

> **Status:** pre-1.0. The API may still change between minor versions, see
> [CHANGELOG.md](CHANGELOG.md).

## License

Apache-2.0
