feat(transformers): Add HuggingFace Transformers integration

.github/workflows/checks.yaml

@@ -135,7 +135,36 @@ jobs:
         shell: bash
         run: |
           mise exec python@${{ matrix.python-version }} -- python ./py/scripts/nox-matrix.py ${{ matrix.shard }} 6 \
-            --exclude-static-checks
+            --exclude-static-checks \
+            --exclude-session-prefix test_transformers
+
+  transformers-py:
+    runs-on: ubuntu-24.04
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      # Serialize jobs so later Python versions can restore the model cache
+      # saved by the first job instead of downloading the same models in parallel.
+      max-parallel: 1
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - name: Setup Python environment
+        uses: ./.github/actions/setup-python-env
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Cache Hugging Face models
+        uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
+        with:
+          path: ~/.cache/huggingface/hub
+          key: huggingface-transformers-${{ runner.os }}-${{ runner.arch }}-${{ matrix.python-version }}-${{ hashFiles('py/src/braintrust/integrations/transformers/test_transformers.py') }}
+          restore-keys: |
+            huggingface-transformers-${{ runner.os }}-${{ runner.arch }}-
+      - name: Test Transformers integration
+        shell: bash
+        run: |
+          mise exec -- uv run --project ./py nox -f ./py/noxfile.py -s test_transformers
 
   adk-py:
     uses: ./.github/workflows/adk-py-test.yaml
@@ -174,6 +203,7 @@ jobs:
       - static_checks
       - smoke
       - nox
+      - transformers-py
       - adk-py
       - langchain-py
       - upload-wheel
@@ -200,6 +230,7 @@ jobs:
           check_result "static_checks" "${{ needs.static_checks.result }}"
           check_result "smoke" "${{ needs.smoke.result }}"
           check_result "nox" "${{ needs.nox.result }}"
+          check_result "transformers-py" "${{ needs['transformers-py'].result }}"
           check_result "adk-py" "${{ needs['adk-py'].result }}"
           check_result "langchain-py" "${{ needs['langchain-py'].result }}"
           check_result "upload-wheel" "${{ needs['upload-wheel'].result }}"

README.md

@@ -65,6 +65,7 @@ BRAINTRUST_API_KEY=<YOUR_API_KEY> braintrust eval tutorial_eval.py
 | [Google GenAI](py/src/braintrust/integrations/google_genai/) | Yes | `google-genai>=1.30.0` |
 | [Bedrock Runtime](py/src/braintrust/integrations/bedrock_runtime/) | Yes | `boto3>=1.34.116` |
 | [HuggingFace Hub](py/src/braintrust/integrations/huggingface_hub/) | Yes | `huggingface-hub>=0.32.0` |
+| [Hugging Face Transformers](py/src/braintrust/integrations/transformers/) | Yes | `transformers>=4.42.0` |
 | [Google ADK](py/src/braintrust/integrations/adk/) | Yes | `google-adk>=1.14.1` |
 | [Pydantic AI](py/src/braintrust/integrations/pydantic_ai/) | Yes | `pydantic_ai>=1.10.0` |
 | [LangChain](py/src/braintrust/integrations/langchain/) | Yes | `langchain-core>=0.3.28` |

py/noxfile.py

@@ -99,12 +99,19 @@ def _ensure_livekit_server(session: nox.Session) -> str:
     return str(install_dir.resolve())
 
 
-def _install_group_locked(session: nox.Session, *group_names: str) -> None:
+def _install_group_locked(
+    session: nox.Session,
+    *group_names: str,
+    indexes: tuple[str, ...] = (),
+) -> None:
     """Install deps from one or more dependency groups using the lockfile.
 
     Runs ``uv export --only-group <name>`` for each group, merges the output,
-    and installs the pre-resolved pins into the session venv.  This gives
+    and installs the pre-resolved pins into the session venv. This gives
     reproducible installs without ad-hoc resolution at install time.
+
+    ``indexes`` names explicit ``[[tool.uv.index]]`` entries that must remain
+    available while installing exported requirements.
     """
     with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
         req_file = f.name
@@ -122,7 +129,21 @@ def _install_group_locked(session: nox.Session, *group_names: str) -> None:
         for name in group_names:
             cmd.extend(["--only-group", name])
         session.run_install(*cmd, silent=SILENT_INSTALLS)
-        session.install("-r", req_file, silent=SILENT_INSTALLS)
+        install_args = ["-r", req_file]
+        configured_indexes = {
+            index.get("name"): index.get("url") for index in _PYPROJECT.get("tool", {}).get("uv", {}).get("index", [])
+        }
+        for index_name in indexes:
+            index_url = configured_indexes.get(index_name)
+            if not index_url:
+                session.error(f"Unknown [[tool.uv.index]] name: {index_name!r}")
+            install_args.extend(("--index", index_url))
+        if indexes:
+            # Exported requirements are fully pinned, so it is safe to search
+            # all configured indexes for each exact version. Without this,
+            # uv may stop at PyTorch's index for unrelated packages.
+            install_args.extend(("--index-strategy", "unsafe-best-match"))
+        session.install(*install_args, silent=SILENT_INSTALLS)
     finally:
         os.unlink(req_file)
 
@@ -669,6 +690,27 @@ def test_huggingface_hub(session, version):
     _run_tests(session, f"{INTEGRATION_DIR}/huggingface_hub/test_huggingface_hub.py", version=version)
 
 
+TRANSFORMERS_VERSIONS = _get_matrix_versions("transformers")
+
+
+@nox.session()
+@nox.parametrize("version", TRANSFORMERS_VERSIONS, ids=TRANSFORMERS_VERSIONS)
+def test_transformers(session, version):
+    """Test local Hugging Face Transformers pipeline instrumentation."""
+    # The 4.42.0 floor pins tokenizers 0.19, whose wheels stop at Python 3.12.
+    if version != LATEST and sys.version_info >= (3, 13):
+        session.skip(f"Transformers {version} does not support Python 3.13+")
+    _install_test_deps(session)
+    _install_matrix_dep(session, "transformers", version)
+    _install_group_locked(session, "test-transformers", indexes=("pytorch-cpu",))
+    _run_tests(
+        session,
+        f"{INTEGRATION_DIR}/transformers/test_transformers.py",
+        version=version,
+        env={"HF_HUB_DISABLE_PROGRESS_BARS": "1"},
+    )
+
+
 TEMPORAL_VERSIONS = _get_matrix_versions("temporalio")
 
 

py/pyproject.toml

@@ -218,6 +218,14 @@ test-huggingface-hub = [
     "aiohttp",
 ]
 
+test-transformers = [
+    {include-group = "test"},
+    # Local pipeline execution is intentionally isolated from other provider
+    # sessions because PyTorch is a large optional runtime dependency.
+    "torch",
+    "sentencepiece",
+]
+
 test-cli = [
     {include-group = "test"},
     "httpx==0.28.1",
@@ -298,6 +306,18 @@ dev = [
 ]
 
 
+# ---------------------------------------------------------------------------
+# UV package sources
+# ---------------------------------------------------------------------------
+
+[tool.uv.sources]
+torch = { index = "pytorch-cpu" }
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
 # ---------------------------------------------------------------------------
 # UV conflicts — auxiliary session groups that pin different versions of the
 # same package cannot coexist in a single resolution.
@@ -477,6 +497,10 @@ latest = "mistralai==2.9.1"
 latest = "huggingface-hub==1.27.0"
 "0.32.0" = "huggingface-hub==0.32.0"
 
+[tool.braintrust.matrix.transformers]
+latest = "transformers==5.15.0"
+"4.42.0" = "transformers==4.42.0"
+
 [tool.braintrust.matrix.temporalio]
 latest = "temporalio==1.31.0"
 "1.20.0" = "temporalio==1.20.0"
@@ -581,3 +605,4 @@ openrouter = "openrouter"
 pipecat-ai = "pipecat"
 strands-agents = "strands"
 temporalio = "temporalio"
+transformers = "transformers"

py/scripts/nox-matrix.py

@@ -92,6 +92,12 @@ def main() -> None:
         default=[],
         help="Exclude a nox session from shard assignment. May be passed multiple times.",
     )
+    parser.add_argument(
+        "--exclude-session-prefix",
+        action="append",
+        default=[],
+        help="Exclude nox sessions whose names start with this prefix. May be passed multiple times.",
+    )
     parser.add_argument(
         "--exclude-static-checks",
         action="store_true",
@@ -129,7 +135,12 @@ def main() -> None:
     excluded_sessions = set(args.exclude_session)
     if args.exclude_static_checks:
         excluded_sessions.update(STATIC_CHECK_SESSIONS)
-    all_sessions = [session for session in all_sessions if session not in excluded_sessions]
+    all_sessions = [
+        session
+        for session in all_sessions
+        if session not in excluded_sessions
+        and not any(session.startswith(prefix) for prefix in args.exclude_session_prefix)
+    ]
     weights, default_weight = load_weights(weights_file)
     shard_assignments = assign_shards(all_sessions, args.num_shards, weights, default_weight)
 

py/src/braintrust/auto.py

@@ -35,6 +35,7 @@
     PydanticAIIntegration,
     StrandsIntegration,
     TemporalIntegration,
+    TransformersIntegration,
 )
 from braintrust.integrations.base import BaseIntegration
 
@@ -67,6 +68,7 @@ def auto_instrument(
     openrouter: bool = True,
     mistral: bool = True,
     huggingface_hub: bool = True,
+    transformers: bool = True,
     agno: bool = True,
     agentscope: bool = True,
     claude_agent_sdk: bool = True,
@@ -105,6 +107,7 @@ def auto_instrument(
         openrouter: Enable OpenRouter instrumentation (default: True)
         mistral: Enable Mistral instrumentation (default: True)
         huggingface_hub: Enable HuggingFace Hub instrumentation (default: True)
+        transformers: Enable local Hugging Face Transformers pipeline instrumentation (default: True)
         agno: Enable Agno instrumentation (default: True)
         agentscope: Enable AgentScope instrumentation (default: True)
         claude_agent_sdk: Enable Claude Agent SDK instrumentation (default: True)
@@ -189,6 +192,8 @@ def auto_instrument(
         results["mistral"] = _instrument_integration(MistralIntegration)
     if huggingface_hub:
         results["huggingface_hub"] = _instrument_integration(HuggingFaceHubIntegration)
+    if transformers:
+        results["transformers"] = _instrument_integration(TransformersIntegration)
     if agno:
         results["agno"] = _instrument_integration(AgnoIntegration)
     if agentscope:

py/src/braintrust/integrations/__init__.py

@@ -25,6 +25,7 @@
 from .pydantic_ai import PydanticAIIntegration
 from .strands import StrandsIntegration
 from .temporal import TemporalIntegration
+from .transformers import TransformersIntegration
 
 
 __all__ = [
@@ -55,4 +56,5 @@
     "PydanticAIIntegration",
     "StrandsIntegration",
     "TemporalIntegration",
+    "TransformersIntegration",
 ]

py/src/braintrust/integrations/auto_test_scripts/test_auto_transformers.py

@@ -0,0 +1,32 @@
+"""Subprocess coverage for Transformers auto-instrumentation."""
+
+# Keep the large Transformers/PyTorch dependencies isolated to their nox job.
+# pylint: disable=import-error
+
+from braintrust.auto import auto_instrument
+from braintrust.integrations.test_utils import autoinstrument_test_context
+
+
+results = auto_instrument(transformers=True)
+assert results.get("transformers") is True
+assert auto_instrument(transformers=True).get("transformers") is True
+
+from transformers import pipeline  # noqa: E402
+
+
+with autoinstrument_test_context("test_auto_transformers", use_vcr=False) as memory_logger:
+    generator = pipeline(
+        "text-generation",
+        model="hf-internal-testing/tiny-random-LlamaForCausalLM",
+        device=-1,
+    )
+    response = generator("Hello", do_sample=False, max_new_tokens=1)
+    assert response
+
+    spans = memory_logger.pop()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span["span_attributes"]["name"] == "huggingface.transformers.text_generation"
+    assert span["context"]["span_origin"]["instrumentation"]["name"] == "transformers-auto"
+
+print("SUCCESS")

py/src/braintrust/integrations/huggingface_hub/tracing.py

@@ -31,6 +31,7 @@
     _log_and_end_span,
     _log_error_and_end_span,
     _normalize_chat_messages,
+    _tensor_shape,
     _timing_metrics,
 )
 from braintrust.logger import start_span as _bt_start_span
@@ -313,22 +314,17 @@ def _feature_extraction_output(result: Any) -> dict[str, Any] | None:
     if result is None:
         return None
 
-    shape = getattr(result, "shape", None)
-    if shape is not None:
-        try:
-            dims = len(shape)
-            if dims == 1:
-                return {"embedding_length": int(shape[0])}
-            if dims == 2:
-                return {"embedding_count": int(shape[0]), "embedding_length": int(shape[1])}
-            if dims >= 3:
-                return {
-                    "embedding_batch_count": int(shape[0]),
-                    "embedding_count": int(shape[1]),
-                    "embedding_length": int(shape[-1]),
-                }
-        except (TypeError, ValueError):
-            pass
+    shape = _tensor_shape(result)
+    if shape:
+        if len(shape) == 1:
+            return {"embedding_length": shape[0]}
+        if len(shape) == 2:
+            return {"embedding_count": shape[0], "embedding_length": shape[1]}
+        return {
+            "embedding_batch_count": shape[0],
+            "embedding_count": shape[1],
+            "embedding_length": shape[-1],
+        }
 
     if isinstance(result, list):
         if not result:

py/src/braintrust/integrations/transformers/__init__.py

@@ -0,0 +1,16 @@
+"""Braintrust integration for local Hugging Face Transformers pipelines."""
+
+from .integration import TransformersIntegration
+from .patchers import wrap_transformers
+
+
+__all__ = [
+    "TransformersIntegration",
+    "setup_transformers",
+    "wrap_transformers",
+]
+
+
+def setup_transformers() -> bool:
+    """Instrument supported Transformers pipeline classes."""
+    return TransformersIntegration.setup()

py/src/braintrust/integrations/transformers/integration.py

@@ -0,0 +1,15 @@
+"""Hugging Face Transformers integration orchestration."""
+
+from braintrust.integrations.base import BaseIntegration
+
+from .patchers import PIPELINE_PATCHERS
+
+
+class TransformersIntegration(BaseIntegration):
+    """Instrument supported local ``transformers`` text pipelines."""
+
+    name = "transformers"
+    import_names = ("transformers",)
+    distribution_names = ("transformers",)
+    min_version = "4.42.0"
+    patchers = PIPELINE_PATCHERS