aviary.labbench 0.30.0__tar.gz

aviary_labbench-0.30.0/PKG-INFO

@@ -0,0 +1,155 @@
+Metadata-Version: 2.4
+Name: aviary.labbench
+Version: 0.30.0
+Summary: LAB-Bench environments implemented with aviary
+Author-email: FutureHouse technical staff <hello@futurehouse.org>
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: fhaviary>=0.14
+Requires-Dist: fhlmi
+Requires-Dist: ldp>=0.25.2
+Requires-Dist: paper-qa[pymupdf]>=2025
+Requires-Dist: pydantic~=2.0
+Requires-Dist: tenacity
+Requires-Dist: typing-extensions; python_version <= "3.12"
+Provides-Extra: datasets
+Requires-Dist: datasets>=2.15; extra == "datasets"
+Provides-Extra: dev
+Requires-Dist: aviary.labbench[datasets,typing]; extra == "dev"
+Requires-Dist: pandas; extra == "dev"
+Requires-Dist: paper-qa>=5.29.1; extra == "dev"
+Requires-Dist: tantivy>=0.25.0; python_version >= "3.14" and extra == "dev"
+Provides-Extra: typing
+Requires-Dist: pillow; extra == "typing"
+
+# aviary.labbench
+
+LAB-Bench environments implemented with aviary,
+allowing agents to perform question answering on scientific tasks.
+
+## Installation
+
+To install the LAB-Bench environment, run:
+
+```bash
+pip install 'fhaviary[labbench]'
+```
+
+## Usage
+
+In [`labbench/env.py`](src/aviary/envs/labbench/env.py), you will find:
+
+- `GradablePaperQAEnvironment`: an PaperQA-backed environment
+  that can grade answers given an evaluation function.
+- `ImageQAEnvironment`: an `GradablePaperQAEnvironment`
+  subclass for QA where image(s) are pre-added.
+
+And in [`labbench/task.py`](src/aviary/envs/labbench/task.py), you will find:
+
+- `TextQATaskDataset`: a task dataset designed to
+  pull down FigQA, LitQA2, or TableQA from Hugging Face,
+  and create one `GradablePaperQAEnvironment` per question.
+- `ImageQATaskDataset`: a task dataset that pairs with `ImageQAEnvironment`
+  for FigQA or TableQA.
+
+Here is an example of how to use them:
+
+```python
+import os
+
+from ldp.agent import SimpleAgent
+from ldp.alg import Evaluator, EvaluatorConfig, MeanMetricsCallback
+from paperqa import Settings
+
+from aviary.env import TaskDataset
+
+
+async def evaluate(folder_of_litqa_v2_papers: str | os.PathLike) -> None:
+    settings = Settings(paper_directory=folder_of_litqa_v2_papers)
+    dataset = TaskDataset.from_name("litqa2", settings=settings)
+    metrics_callback = MeanMetricsCallback(eval_dataset=dataset)
+
+    evaluator = Evaluator(
+        config=EvaluatorConfig(batch_size=3),
+        agent=SimpleAgent(),
+        dataset=dataset,
+        callbacks=[metrics_callback],
+    )
+    await evaluator.evaluate()
+    print(metrics_callback.eval_means)
+```
+
+### Image Question-Answer
+
+This is an environment/dataset for giving PaperQA a `Docs` object with
+the image(s) for one LAB-Bench question.
+It's designed to be a comparison with zero-shotting the question to a LLM,
+but instead of a singular prompt the image is put through the PaperQA agent loop.
+
+```python
+from typing import cast
+
+import litellm
+import pytest
+from ldp.agent import Agent
+from ldp.alg import (
+    Evaluator,
+    EvaluatorConfig,
+    MeanMetricsCallback,
+    StoreTrajectoriesCallback,
+)
+from paperqa.settings import AgentSettings, IndexSettings
+
+from aviary.envs.labbench import (
+    ImageQAEnvironment,
+    ImageQATaskDataset,
+    LABBenchDatasets,
+)
+
+
+@pytest.mark.asyncio
+async def test_image_qa(tmp_path) -> None:
+    litellm.num_retries = 8  # Mitigate connection-related failures
+    settings = ImageQAEnvironment.make_base_settings()
+    settings.agent = AgentSettings(
+        agent_type="ldp.agent.SimpleAgent",
+        index=IndexSettings(paper_directory=tmp_path),
+        # TODO: add image support for paper_search
+        tool_names={"gather_evidence", "gen_answer", "complete", "reset"},
+        agent_evidence_n=3,  # Bumped up to collect several perspectives
+    )
+    dataset = ImageQATaskDataset(dataset=LABBenchDatasets.TABLE_QA, settings=settings)
+    t_cb = StoreTrajectoriesCallback()
+    m_cb = MeanMetricsCallback(eval_dataset=dataset, track_tool_usage=True)
+    evaluator = Evaluator(
+        config=EvaluatorConfig(
+            batch_size=256,  # Use batch size greater than FigQA size and TableQA size
+            max_rollout_steps=18,  # Match aviary paper's PaperQA setting
+        ),
+        agent=cast(Agent, await settings.make_ldp_agent(settings.agent.agent_type)),
+        dataset=dataset,
+        callbacks=[t_cb, m_cb],
+    )
+    await evaluator.evaluate()
+    print(m_cb.eval_means)
+```
+
+## References
+
+[1] Skarlinski et al.
+[Language agents achieve superhuman synthesis of scientific knowledge](https://arxiv.org/abs/2409.13740).
+ArXiv:2409.13740, 2024.
+
+[2] Laurent et al.
+[LAB-Bench: Measuring Capabilities of Language Models for Biology Research](https://arxiv.org/abs/2407.10362).
+ArXiv:2407.10362, 2024.

aviary_labbench-0.30.0/README.md

@@ -0,0 +1,121 @@
+# aviary.labbench
+
+LAB-Bench environments implemented with aviary,
+allowing agents to perform question answering on scientific tasks.
+
+## Installation
+
+To install the LAB-Bench environment, run:
+
+```bash
+pip install 'fhaviary[labbench]'
+```
+
+## Usage
+
+In [`labbench/env.py`](src/aviary/envs/labbench/env.py), you will find:
+
+- `GradablePaperQAEnvironment`: an PaperQA-backed environment
+  that can grade answers given an evaluation function.
+- `ImageQAEnvironment`: an `GradablePaperQAEnvironment`
+  subclass for QA where image(s) are pre-added.
+
+And in [`labbench/task.py`](src/aviary/envs/labbench/task.py), you will find:
+
+- `TextQATaskDataset`: a task dataset designed to
+  pull down FigQA, LitQA2, or TableQA from Hugging Face,
+  and create one `GradablePaperQAEnvironment` per question.
+- `ImageQATaskDataset`: a task dataset that pairs with `ImageQAEnvironment`
+  for FigQA or TableQA.
+
+Here is an example of how to use them:
+
+```python
+import os
+
+from ldp.agent import SimpleAgent
+from ldp.alg import Evaluator, EvaluatorConfig, MeanMetricsCallback
+from paperqa import Settings
+
+from aviary.env import TaskDataset
+
+
+async def evaluate(folder_of_litqa_v2_papers: str | os.PathLike) -> None:
+    settings = Settings(paper_directory=folder_of_litqa_v2_papers)
+    dataset = TaskDataset.from_name("litqa2", settings=settings)
+    metrics_callback = MeanMetricsCallback(eval_dataset=dataset)
+
+    evaluator = Evaluator(
+        config=EvaluatorConfig(batch_size=3),
+        agent=SimpleAgent(),
+        dataset=dataset,
+        callbacks=[metrics_callback],
+    )
+    await evaluator.evaluate()
+    print(metrics_callback.eval_means)
+```
+
+### Image Question-Answer
+
+This is an environment/dataset for giving PaperQA a `Docs` object with
+the image(s) for one LAB-Bench question.
+It's designed to be a comparison with zero-shotting the question to a LLM,
+but instead of a singular prompt the image is put through the PaperQA agent loop.
+
+```python
+from typing import cast
+
+import litellm
+import pytest
+from ldp.agent import Agent
+from ldp.alg import (
+    Evaluator,
+    EvaluatorConfig,
+    MeanMetricsCallback,
+    StoreTrajectoriesCallback,
+)
+from paperqa.settings import AgentSettings, IndexSettings
+
+from aviary.envs.labbench import (
+    ImageQAEnvironment,
+    ImageQATaskDataset,
+    LABBenchDatasets,
+)
+
+
+@pytest.mark.asyncio
+async def test_image_qa(tmp_path) -> None:
+    litellm.num_retries = 8  # Mitigate connection-related failures
+    settings = ImageQAEnvironment.make_base_settings()
+    settings.agent = AgentSettings(
+        agent_type="ldp.agent.SimpleAgent",
+        index=IndexSettings(paper_directory=tmp_path),
+        # TODO: add image support for paper_search
+        tool_names={"gather_evidence", "gen_answer", "complete", "reset"},
+        agent_evidence_n=3,  # Bumped up to collect several perspectives
+    )
+    dataset = ImageQATaskDataset(dataset=LABBenchDatasets.TABLE_QA, settings=settings)
+    t_cb = StoreTrajectoriesCallback()
+    m_cb = MeanMetricsCallback(eval_dataset=dataset, track_tool_usage=True)
+    evaluator = Evaluator(
+        config=EvaluatorConfig(
+            batch_size=256,  # Use batch size greater than FigQA size and TableQA size
+            max_rollout_steps=18,  # Match aviary paper's PaperQA setting
+        ),
+        agent=cast(Agent, await settings.make_ldp_agent(settings.agent.agent_type)),
+        dataset=dataset,
+        callbacks=[t_cb, m_cb],
+    )
+    await evaluator.evaluate()
+    print(m_cb.eval_means)
+```
+
+## References
+
+[1] Skarlinski et al.
+[Language agents achieve superhuman synthesis of scientific knowledge](https://arxiv.org/abs/2409.13740).
+ArXiv:2409.13740, 2024.
+
+[2] Laurent et al.
+[LAB-Bench: Measuring Capabilities of Language Models for Biology Research](https://arxiv.org/abs/2407.10362).
+ArXiv:2407.10362, 2024.

aviary_labbench-0.30.0/pyproject.toml

@@ -0,0 +1,56 @@
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = ["setuptools>=64", "setuptools_scm>=8"]
+
+[project]
+authors = [
+    {email = "hello@futurehouse.org", name = "FutureHouse technical staff"},
+]
+classifiers = [
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Programming Language :: Python",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "fhaviary>=0.14",  # For MultipleChoiceQuestion
+    "fhlmi",
+    "ldp>=0.25.2",  # Pin for lmi migration
+    "paper-qa[pymupdf]>=2025",  # Pin for multimodal
+    "pydantic~=2.0",
+    "tenacity",
+    "typing-extensions; python_version <= '3.12'",  # For TypeVar default
+]
+description = "LAB-Bench environments implemented with aviary"
+dynamic = ["version"]
+name = "aviary.labbench"
+readme = "README.md"
+requires-python = ">=3.11"
+
+[project.optional-dependencies]
+datasets = [
+    "datasets>=2.15",  # Lower pin for https://github.com/huggingface/datasets/pull/6404
+]
+dev = [
+    "aviary.labbench[datasets,typing]",
+    "pandas",
+    "paper-qa>=5.29.1",  # Pin for gen_answer's EmptyDocsError, with fix
+    "tantivy>=0.25.0; python_version >= '3.14'",  # For Python 3.14 support
+]
+typing = ["pillow"]
+
+[tool.ruff]
+extend = "../../pyproject.toml"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools_scm]
+root = "../.."
+version_file = "src/aviary/envs/labbench/version.py"

aviary_labbench-0.30.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

aviary_labbench-0.30.0/src/aviary/envs/labbench/__init__.py

@@ -0,0 +1,31 @@
+from .env import (
+    DEFAULT_REWARD_MAPPING,
+    GradablePaperQAEnvironment,
+    ImageQAEnvironment,
+    make_discounted_returns,
+)
+from .task import (
+    DEFAULT_AVIARY_PAPER_HF_HUB_NAME,
+    DEFAULT_LABBENCH_HF_HUB_NAME,
+    ImageQATaskDataset,
+    LABBenchDatasets,
+    PaperQATaskDataset,
+    TextQATaskDataset,
+    TextQATaskSplit,
+    read_ds_from_hub,
+)
+
+__all__ = [
+    "DEFAULT_AVIARY_PAPER_HF_HUB_NAME",
+    "DEFAULT_LABBENCH_HF_HUB_NAME",
+    "DEFAULT_REWARD_MAPPING",
+    "GradablePaperQAEnvironment",
+    "ImageQAEnvironment",
+    "ImageQATaskDataset",
+    "LABBenchDatasets",
+    "PaperQATaskDataset",
+    "TextQATaskDataset",
+    "TextQATaskSplit",
+    "make_discounted_returns",
+    "read_ds_from_hub",
+]

aviary_labbench-0.30.0/src/aviary/envs/labbench/env.py

@@ -0,0 +1,276 @@
+import logging
+import sys
+import tempfile
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from copy import deepcopy
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generic, Self, cast
+from uuid import UUID
+
+from aviary.core import (
+    Messages,
+    MultipleChoiceEvaluation,
+    MultipleChoiceQuestion,
+    ToolRequestMessage,
+)
+from aviary.env import ENV_REGISTRY
+from ldp.utils import discounted_returns
+from lmi import EmbeddingModel, LiteLLMModel
+from paperqa.agents.env import POPULATE_FROM_SETTINGS, PaperQAEnvironment
+from paperqa.agents.search import SearchIndex, maybe_get_manifest
+from paperqa.docs import Docs
+from paperqa.settings import AnswerSettings, ParsingSettings, Settings
+
+if TYPE_CHECKING:
+    from PIL.Image import Image
+
+if sys.version_info >= (3, 13):
+    from typing import TypeVar
+else:
+    from typing_extensions import TypeVar  # For TypeVar.default backport
+
+logger = logging.getLogger(__name__)
+
+TEvaluation = TypeVar("TEvaluation", default=MultipleChoiceEvaluation)
+
+DEFAULT_REWARD_MAPPING = {"correct": 1.0, "unsure": 0.1, "incorrect": -1.0}
+
+
+def make_discounted_returns(
+    evaluation: MultipleChoiceEvaluation,
+    num_steps: int,
+    rewards: Mapping[str, float] = DEFAULT_REWARD_MAPPING,
+    discount: float = 1.0,
+) -> list[float]:
+    return discounted_returns(
+        # paper-qa has no intermediary rewards
+        [0] * (num_steps - 1) + [rewards[evaluation.value]],
+        terminated=[False] * (num_steps - 1) + [True],
+        discount=discount,
+    )
+
+
+class GradablePaperQAEnvironment(PaperQAEnvironment, Generic[TEvaluation]):
+    """Extended environment that can grade answers."""
+
+    def __init__(
+        self,
+        query: str | MultipleChoiceQuestion,
+        settings: Settings,
+        docs: Docs,
+        llm_model: LiteLLMModel | None = POPULATE_FROM_SETTINGS,
+        summary_llm_model: LiteLLMModel | None = POPULATE_FROM_SETTINGS,
+        embedding_model: EmbeddingModel | None = POPULATE_FROM_SETTINGS,
+        session_id: UUID | None = None,
+        sources: str | list[str] | None = None,
+        rewards: Mapping[str, float] = DEFAULT_REWARD_MAPPING,
+        evaluation_callback: Callable[[TEvaluation], Awaitable] | None = None,
+        **env_kwargs,
+    ):
+        super().__init__(
+            query,
+            settings,
+            docs,
+            llm_model,
+            summary_llm_model,
+            embedding_model,
+            session_id,
+            **env_kwargs,
+        )
+        # Enables checking an Index has the right DOI(s)
+        self.sources: list[str] | None = (
+            [sources] if isinstance(sources, str) else sources
+        )
+        self._evaluation_callback = evaluation_callback
+        self._rewards = rewards
+
+    async def validate_sources(
+        self, manifest_or_index: dict[str, dict[str, Any]] | SearchIndex | None = None
+    ) -> None:
+        """Validate the sources can be found in the input manifest or index."""
+        if not self.sources:
+            return
+        if manifest_or_index is None:  # Let's try to load in the manifest
+            manifest_or_index = await maybe_get_manifest(
+                filename=await self._settings.agent.index.finalize_manifest_file()
+            )
+        if isinstance(manifest_or_index, SearchIndex):
+            entity: str = "index"
+            file_names: set[str] = {k for k in await manifest_or_index.index_files if k}
+            lowercased_dois: set[str] = set()
+        else:
+            entity = "manifest"
+            file_names = {k for k in manifest_or_index if k}
+            lowercased_dois = {
+                v["doi"].lower() for v in manifest_or_index.values() if v["doi"]
+            }
+        if not file_names:  # File names being empty means something's wrong
+            logger.warning(
+                f"Can't validate sources {self.sources} without a correctly specified"
+                f" {entity}."
+            )
+            return
+        not_found = [
+            s
+            for s in self.sources
+            if s not in file_names and s.lower() not in lowercased_dois
+        ]
+        if not_found:
+            question = (
+                self._query
+                if isinstance(self._query, str)
+                else self._query.question_prompt
+            )
+            raise ValueError(
+                f"Sources {not_found} of {self.sources} not found in the {entity},"
+                f" the corresponding query was {question!r}."
+            )
+
+    async def _evaluate_answer(self) -> TEvaluation:
+        # If the ensuring evaluation fails (e.g. due to OpenAI being down), we can:
+        # - Suppress the exception and declare the evaluation as incorrect, which can
+        #   negatively reward what otherwise was a good trajectory containing a correct
+        #   answer. We don't want "bad" offline data, so it's not what we do.
+        # - Suppress the exception and just give super()'s reward, but again this could
+        #   incorrectly reward what otherwise was a good trajectory.
+        # - Don't suppress the exception, which leads to the trajectory failing, and
+        #   removes it from the learnable pool. This is the only safe default behavior.
+        evaluation, self.state.session.graded_answer = await cast(
+            "MultipleChoiceQuestion", self._query
+        ).grade(self.state.session.answer)
+        return evaluation  # type: ignore[return-value]
+
+    async def step(
+        self, action: ToolRequestMessage
+    ) -> tuple[Messages, float, bool, bool]:
+        messages, reward, done, truncated = await super().step(action)
+        if not done or not isinstance(self._query, MultipleChoiceQuestion):
+            return messages, reward, done, truncated
+        evaluation = await self._evaluate_answer()
+        if evaluation_callback := self._evaluation_callback:
+            await evaluation_callback(evaluation)
+
+        return (
+            messages,
+            reward + self._rewards[cast("MultipleChoiceEvaluation", evaluation).value],
+            done,
+            truncated,
+        )
+
+    async def get_id(self) -> str:
+        if (
+            isinstance(self._query, str)
+            or self._query.question_id
+            == MultipleChoiceQuestion.model_fields["question_id"].default
+        ):
+            details = (
+                ", as just a question was configured"
+                if isinstance(self._query, str)
+                else ", as the default ID remains present"
+            )
+            raise ValueError(f"No question ID was configured{details}.")
+        return str(self._query.question_id)
+
+    def __deepcopy__(self, memo) -> Self:
+        copy_state = deepcopy(self.state, memo)
+        # We don't know the side effects of deep copying a litellm.Router,
+        # so we force a shallow copy of these LiteLLMModels
+        env_model_kwargs: dict[str, Any] = {
+            name: model if model is None else type(model)(**model.model_dump())
+            for name, model in (
+                ("llm_model", self._llm_model),
+                ("summary_llm_model", self._summary_llm_model),
+                ("embedding_model", self._embedding_model),
+            )
+        }
+        copy_self = type(self)(
+            query=self._query,  # No need to copy since we read only
+            settings=deepcopy(self._settings, memo),  # Deepcopy just to be safe
+            docs=copy_state.docs,
+            sources=self.sources,
+            rewards=self._rewards,
+            evaluation_callback=self._evaluation_callback,
+            **env_model_kwargs,
+        )
+        copy_self.state = copy_state
+        # Because we shallow copied the LiteLLMModels, we need to re-make the
+        # tool functions within the tools
+        copy_self.tools = copy_self.make_tools()
+        return copy_self
+
+
+ENV_REGISTRY["paperqa-local"] = (
+    GradablePaperQAEnvironment.__module__,
+    GradablePaperQAEnvironment.__name__,
+)
+
+
+class ImageQAEnvironment(GradablePaperQAEnvironment):
+    """Image question-answer environment useful for LAB-Bench's FigQA and TableQA."""
+
+    @classmethod
+    def make_base_settings(cls, **kwargs) -> Settings:
+        """Make a settings object that takes into account image-based QA restrictions."""
+        return Settings(
+            # PaperQA doesn't support image embeddings yet, so disable embedding
+            # Disable doc details since we just have images here (not a PDF with metadata)
+            parsing=ParsingSettings(defer_embedding=True, use_doc_details=False),
+            answer=AnswerSettings(evidence_retrieval=False),
+            **kwargs,
+        )
+
+    def __init__(
+        self,
+        *args,
+        images: "bytes | Image | Sequence[bytes | Image]",
+        image_paths: str | Sequence[str],
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        if not isinstance(self._query, MultipleChoiceQuestion):
+            raise TypeError(
+                f"{type(self).__name__} requires a {MultipleChoiceQuestion.__name__}"
+                f" as the query, not {type(self._query)}."
+            )
+        # FigQA has 1 image with paths, TableQA has 1+ images with paths
+        if not isinstance(image_paths, str):  # Assume TableQA
+            self._images_with_names: "list[tuple[bytes | Image, str]]" = [  # noqa: UP037
+                (image, Path(image_path).name)
+                for image, image_path in zip(
+                    cast("Sequence[bytes | Image]", images), image_paths, strict=True
+                )
+            ]
+        else:  # Assume FigQA
+            self._images_with_names = [
+                (cast("bytes | Image", images), Path(image_paths).name)
+            ]
+
+    def get_images(self) -> "list[bytes | Image]":
+        """
+        Get the image(s) used in the environment, helpful for recall measurement.
+
+        NOTE: FigQA has 1 image with paths, TableQA has 1+ images with paths.
+        """
+        return [image for image, _ in self._images_with_names]
+
+    async def _reset_docs(self) -> None:
+        """Hook to reset the docs when creating the initial state."""
+        self._docs.clear_docs()
+
+        # Now add the image(s) to the docs
+        with tempfile.TemporaryDirectory() as tmpdir:
+            for image, image_name in self._images_with_names:
+                tmp_image_path = Path(tmpdir) / image_name
+                if isinstance(image, bytes):
+                    tmp_image_path.write_bytes(image)
+                else:
+                    image.save(tmp_image_path)
+                await self._docs.aadd(
+                    tmp_image_path,
+                    citation=(
+                        f"Row ID {self._query.question_id} filename {tmp_image_path.name}"
+                        if isinstance(self._query, MultipleChoiceQuestion)
+                        else f"Filename {tmp_image_path.name}"
+                    ),
+                    settings=self._settings,
+                )