granite-common 0.1.dev19__py3-none-any.whl

granite_common/.gitignore

@@ -0,0 +1,3 @@
+
+# Generated version info
+_version.py

granite_common/__init__.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+
+__doc__ = f"""
+{__package__} is a Python library that provides enhanced prompt creation and output
+parsing for IBM Granite models
+"""
+
+# Local
+# This file explicitly imports all the symbols that we export at the top level of this
+# package's namespace.
+from .base.types import (
+    AssistantMessage,
+    ChatCompletion,
+    UserMessage,
+)
+from .granite3.granite32 import (
+    Granite3Point2ChatCompletion,
+    Granite3Point2InputProcessor,
+    Granite3Point2OutputProcessor,
+)
+from .granite3.granite33 import (
+    Granite3Point3ChatCompletion,
+    Granite3Point3InputProcessor,
+    Granite3Point3OutputProcessor,
+)
+
+# The contents of __all__ must be strings
+__all__ = (
+    obj.__name__
+    for obj in (
+        AssistantMessage,
+        ChatCompletion,
+        UserMessage,
+        Granite3Point2InputProcessor,
+        Granite3Point2OutputProcessor,
+        Granite3Point2ChatCompletion,
+        Granite3Point3ChatCompletion,
+        Granite3Point3InputProcessor,
+        Granite3Point3OutputProcessor,
+    )
+)

granite_common/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1.dev19'
+__version_tuple__ = version_tuple = (0, 1, 'dev19')

granite_common/base/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+
+__doc__ = f"""
+{__name__} data structures and functions that are shared across multiple releases of the
+Granite models.
+"""

granite_common/base/io.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+
+__doc__ = """
+Classes and functions that implement common aspects of input and output string 
+processing for all Granite models.
+"""
+
+# Standard
+import abc
+
+# Local
+from .types import AssistantMessage, ChatCompletion
+
+
+class InputProcessor(abc.ABC):
+    """
+    Interface for generic input processors. An input processor exposes an
+    API to transform a chat completion request into a string prompt.
+    """
+
+    @abc.abstractmethod
+    def transform(
+        self, chat_completion: ChatCompletion, add_generation_prompt: bool = True
+    ) -> str:
+        """
+        Convert the structured representation of the inputs to a completion request into
+        the string representation of the tokens that should be sent to the model to
+        implement said request.
+
+        :param chat_completion: Structured representation of the inputs
+        :param add_generation_prompt: If ``True``, the returned prompt string will
+            contain a prefix of the next assistant response for use as a prompt to a
+            generation request. Otherwise, the prompt will only contain the messages and
+            documents in ``input``.
+
+        :returns: String that can be passed to the model's tokenizer to create a prompt
+            for generation.
+        """
+
+
+class OutputProcessor(abc.ABC):
+    """
+    Base class for generic output processors. An output processor exposes an
+    API to transform model output into a structured representation of the
+    information.
+
+    This interface is very generic; see individual classes for more specific arguments
+    """
+
+    @abc.abstractmethod
+    def transform(
+        self, model_output: str, chat_completion: ChatCompletion | None = None
+    ) -> AssistantMessage:
+        """
+        Convert the model output generated into a structured representation of the
+        information.
+
+        :param model_output: String output of the a generation request, potentially
+            incomplete if it was a streaming request
+        :param chat_completion: The chat completion request that produced
+            ``model_output``. Parameters of the request can determine how the output
+            should be decoded.
+
+        :returns: The parsed output so far, as an instance of :class:`AssistantMessage`
+            possibly with model-specific extension fields.
+        """

granite_common/base/types.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Common shared types
+"""
+
+# Standard
+from typing import Literal, TypeAlias
+
+# Third Party
+from typing_extensions import Any
+import pydantic
+
+
+class NoDefaultsMixin:
+    """
+    Mixin so that we don't need to copy and paste the code to avoid filling JSON values
+    with a full catalog of the default values of rarely-used fields.
+    """
+
+    @pydantic.model_serializer(mode="wrap")
+    def _workaround_for_design_flaw_in_pydantic(self, nxt):
+        """
+        Workaround for a design flaw in Pydantic that forces users to accept
+        unnecessary garbage in their serialized JSON data or to override
+        poorly-documented serialization hooks repeatedly.  Automates overriding said
+        poorly-documented serialization hooks for a single dataclass.
+
+        See https://github.com/pydantic/pydantic/issues/4554 for the relevant dismissive
+        comment from the devs. This comment suggests overriding :func:`dict()`, but that
+        method was disabled a year later. Now you need to add a custom serializer method
+        with a ``@model_serializer`` decorator.
+
+        See the docs at
+        https://docs.pydantic.dev/latest/api/functional_serializers/
+        for some dubious information on how this API works.
+        See comments below for important gotchas that aren't in the documentation.
+        """
+        # Start with the value that self.model_dump() would return without this mixin.
+        # Otherwise serialization of sub-records will be inconsistent.
+        serialized_value = nxt(self)
+
+        # Figure out which fields are set. Pydantic does not make this easy.
+        # Start with fields that are set in __init__() or in the JSON parser.
+        fields_to_retain_set = self.model_fields_set
+
+        # Add in fields that were set during validation and extra fields added by
+        # setattr().  These fields all go to self.model.extra
+        if self.model_extra is not None:  # model_extra is sometimes None. Not sure why.
+            # model_extra is a dictionary. There is no self.model_extra_fields_set.
+            fields_to_retain_set |= set(list(self.model_extra))
+
+        # Use a subclass hook for the additional fields that fall through the cracks.
+        fields_to_retain_set |= set(self._keep_these_fields())
+
+        # Avoid changing Pydantic's field order or downstream code that computes a
+        # diff over JSON strings will break.
+        fields_to_retain = [k for k in serialized_value if k in fields_to_retain_set]
+
+        # Fields that weren't in the original serialized value should be in a consistent
+        # order to ensure consistent serialized output.
+        # Use alphabetical order for now and hope for the best.
+        fields_to_retain.extend(sorted(fields_to_retain_set - self.model_fields_set))
+
+        result = {}
+        for f in fields_to_retain:
+            if f in serialized_value:
+                result[f] = serialized_value[f]
+            else:
+                # Sometimes Pydantic adds fields to self.model_fields_set without adding
+                # them to the output of self.model_dump()
+                result[f] = getattr(self, f)
+        return result
+
+    def _keep_these_fields(self) -> tuple[str]:
+        """
+        Dataclasses that include this mixin can override this method to add specific
+        default values to serialized JSON.
+
+        This is necessary for round-tripping to JSON when there are fields that
+        determine which dataclass to use for deserialization.
+        """
+        return ()
+
+
+class _ChatMessageBase(pydantic.BaseModel, NoDefaultsMixin):
+    """Base class for all message types.
+
+    Due to the vagaries of Pydantic's JSON parser, we use this class only for common
+    functionality, and NOT for defining a common dataclass base type. Use the
+    :class:`ChatMessage` type alias to annotate a field or argument as accepting all
+    subclasses of this one."""
+
+    content: str
+    """Every message has raw string content, even if it also contains parsed structured
+    content such as a JSON record."""
+
+    def _keep_these_fields(self):
+        return ("role",)
+
+
+class UserMessage(_ChatMessageBase):
+    """User message for an IBM Granite model chat completion request."""
+
+    role: Literal["user"] = "user"
+
+
+class ToolCall(pydantic.BaseModel, NoDefaultsMixin):
+    """Format of an entry in the ``tool_calls`` list of an assistant message"""
+
+    id: str | None = None
+    name: str
+
+    # This field should adhere to the argument schema from the  associated
+    # FunctionDefinition in the generation request that produced it.
+    arguments: dict[str, Any] | None
+
+
+class AssistantMessage(_ChatMessageBase):
+    """
+    Lowest-common-denominator assistant message for an IBM Granite model chat
+    completion request.
+    """
+
+    role: Literal["assistant"] = "assistant"
+    tool_calls: list[ToolCall] | None = None
+
+
+class ToolResultMessage(_ChatMessageBase):
+    """
+    Message containing the result of a tool call in an IBM Granite model chat completion
+    request.
+    """
+
+    role: Literal["tool"] = "tool"
+    tool_call_id: str
+
+
+class SystemMessage(_ChatMessageBase):
+    """System message for an IBM Granite model chat completion request."""
+
+    role: Literal["system"] = "system"
+
+
+ChatMessage: TypeAlias = (
+    UserMessage | AssistantMessage | ToolResultMessage | SystemMessage
+)
+"""Type alias for all message types. We use this Union instead of the actual base class
+:class:`_ChatMessageBase` so that Pydantic can parse the message list from JSON."""
+
+
+class ToolDefinition(pydantic.BaseModel, NoDefaultsMixin):
+    """
+    An entry in the ``tools`` list in an IBM Granite model chat completion request.
+    """
+
+    name: str
+    description: str | None = None
+
+    # This field holds a JSON schema for a record, but the `jsonschema` package doesn't
+    # define an object type for such a schema, instead using a dictionary.
+    parameters: dict[str, Any] | None = None
+
+
+class Document(pydantic.BaseModel, NoDefaultsMixin):
+    """RAG documents, which in practice are usually snippets drawn from larger
+    documents."""
+
+    text: str
+    doc_id: str | None = None
+
+
+class ChatTemplateKwargs(pydantic.BaseModel, NoDefaultsMixin):
+    """
+    Values that can appear in the ``chat_template_kwargs`` portion of a valid chat
+    completion request for a Granite model.
+    """
+
+    documents: list[Document] | None = None
+
+
+class ChatCompletion(pydantic.BaseModel, NoDefaultsMixin):
+    """
+    Lowest-common-denominator inputs to a chat completion request for an IBM Granite
+    model.
+
+    The schema of this object mirrors that of a chat completion request in vLLM's
+    OpenAI-compatible inference API.
+    """
+
+    messages: list[ChatMessage]
+    model: str | None = None
+    tools: list[ToolDefinition] | None = None
+
+    chat_template_kwargs: ChatTemplateKwargs | None = pydantic.Field(
+        default=None,
+        description=(
+            "Additional kwargs to pass to the template renderer. "
+            "Will be accessible by the chat template. "
+            "Restricted to fields that at least one Granite model "
+            "supports."
+        ),
+    )
+
+    model_config = pydantic.ConfigDict(
+        # Pass through arbitrary additional keyword arguments for handling by
+        # model-specific I/O processors.
+        arbitrary_types_allowed=True,
+        extra="allow",
+    )
+
+    def __getattr__(self, name: str) -> any:
+        """Allow attribute access for unknown attributes"""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return None

granite_common/granite3/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+
+__doc__ = """
+Input and output processing for the Granite 3 family of models.
+"""
+
+# Local
+from .types import ControlsRecord
+
+__all__ = ("ControlsRecord",)

granite_common/granite3/constants.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+
+__doc__ = """
+Constants used in code that is specific to the Granite 3 family of models, but not 
+specific to a particular point release.
+"""
+
+
+# String that a Granite 3.x model must receive immediately after _SYSTEM_MESSAGE_START
+# if there are documents in the current request but there are no tools in the current
+# request.
+NO_TOOLS_AND_DOCS_SYSTEM_MESSAGE_PART = """\
+ Write the response to the user's input by strictly aligning with the facts in the \
+provided documents. If the information needed to answer the question is not available \
+in the documents, inform the user that the question cannot be answered based on the \
+available data."""
+
+
+# String that a Granite 3.x model must receive immediately after _SYSTEM_MESSAGE_START
+# if there are no tools or documents in the current request and the "thinking" flag is
+# set to `False`.
+NO_TOOLS_NO_DOCS_NO_THINKING_SYSTEM_MESSAGE_PART = """\
+ You are a helpful AI assistant."""

granite_common/granite3/granite32/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+
+__doc__ = """
+Input and output processing for the Granite 3.2 family of models.
+"""
+
+# Local
+from .input import Granite3Point2InputProcessor
+from .output import Granite3Point2OutputProcessor
+from .types import Granite3Point2ChatCompletion
+
+__all__ = (
+    "Granite3Point2ChatCompletion",
+    "Granite3Point2InputProcessor",
+    "Granite3Point2OutputProcessor",
+)

granite_common/granite3/granite32/constants.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+
+__doc__ = """
+Constants used in code that is specific to the Granite 3.2 family of models.
+"""
+
+
+# Delimiters for chain of thought output of Granite 3.2
+COT_START = "Here is my thought process:"
+COT_END = "Here is my response:"
+
+# Some versions of the model are known to shorten "Here is" to "Here's", so we
+# provide alternate forms of these strings for those versions.
+COT_START_ALTERNATIVES = [
+    COT_START,
+    "Here's my thought process:",
+]
+COT_END_ALTERNATIVES = [
+    COT_END,
+    "Here's my response:",
+]
+
+# Delimiters for hallucination and citation output of Granite 3.2
+CITATION_START = "# Citations:"
+HALLUCINATION_START = "# Hallucinations:"
+
+
+# String that a Granite 3.2 model must receive immediately after _SYSTEM_MESSAGE_START
+# if there are both tools and RAG documents in the current request.
+TOOLS_AND_DOCS_SYSTEM_MESSAGE_PART = """\
+ You are a helpful AI assistant with access to the following tools. When a tool is \
+required to answer the user's query, respond with <|tool_call|> followed by a JSON \
+list of tools used. If a tool does not exist in the provided list of tools, notify the \
+user that you do not have the ability to fulfill the request.
+
+Write the response to the user's input by strictly aligning with the facts in the \
+provided documents. If the information needed to answer the question is not available \
+in the documents, inform the user that the question cannot be answered based on the \
+available data."""
+
+# String that a Granite 3.2 model must receive immediately after _SYSTEM_MESSAGE_START
+# if there are no tools or documents in the current request and the "thinking" flag is
+# set to `True`.
+NO_TOOLS_AND_NO_DOCS_AND_THINKING_SYSTEM_MESSAGE_PART = f"""\
+ You are a helpful AI assistant.
+Respond to every user query in a comprehensive and detailed way. You can write down \
+your thoughts and reasoning process before responding. In the thought process, engage \
+in a comprehensive cycle of analysis, summarization, exploration, reassessment, \
+reflection, backtracing, and iteration to develop well-considered thinking process. \
+In the response section, based on various attempts, explorations, and reflections from \
+the thoughts section, systematically present the final solution that you deem correct. \
+The response should summarize the thought process. Write your thoughts after '\
+{COT_START}' and write your response after '{COT_END}' \
+for each user query."""
+
+
+# String that a Granite 3.2 model must receive immediately after either
+# _TOOLS_AND_DOCS_SYSTEM_MESSAGE_MIDDLE  (if there are tools) or
+# _NO_TOOLS_AND_DOCS_SYSTEM_MESSAGE_MIDDLE (if there are no tools) in the system prompt
+# if the "citations" flag is `True` and there are documents.
+DOCS_AND_CITATIONS_SYSTEM_MESSAGE_PART = """\
+
+
+In your response, use the symbols <co> and </co> to indicate when a fact comes from a \
+document in the search result, e.g <co>0</co> for a fact from document 0. Afterwards, \
+list all the citations with their corresponding documents in an ordered list."""
+
+# String that a Granite 3.2 model must receive immediately after either
+# _TOOLS_AND_DOCS_SYSTEM_MESSAGE_MIDDLE (if there are tools and no citations) or
+# _NO_TOOLS_AND_DOCS_SYSTEM_MESSAGE_MIDDLE (if there are no tools or citations) or
+# _DOCS_AND_CITATIONS_SYSTEM_MESSAGE_PART in the system prompt
+# if the "hallucinations" flag is `True` and there are documents.
+# Note that a list of zero documents counts as "having documents".
+DOCS_AND_HALLUCINATIONS_SYSTEM_MESSAGE_PART = """\
+
+
+Finally, after the response is written, include a numbered list of sentences from the \
+response that are potentially hallucinated and not based in the documents."""
+
+# String that a Granite 3.2 model must receive immediately after _SYSTEM_MESSAGE_START
+# if there are tools in the current request but there are no documents in the current
+# request.
+TOOLS_AND_NO_DOCS_SYSTEM_MESSAGE_PART = """\
+ You are a helpful AI assistant with access to the following tools. When a tool is \
+required to answer the user's query, respond with <|tool_call|> followed by a JSON \
+list of tools used. If a tool does not exist in the provided list of tools, notify the \
+user that you do not have the ability to fulfill the request."""
+
+MODEL_NAME = "Granite 3.2"
+MODEL_HF_PATH_2B = "ibm-granite/granite-3.2-2b-instruct"