Add Token Overlap setting to Controls (#106)

* Add Token Overlap setting to Controls

* Fix mypy issue

* Add more documentation

aleph_alpha_client/__init__.py

@@ -1,78 +1,82 @@
+from .prompt import (
+    ControlTokenOverlap,
+    Image,
+    ImageControl,
+    ImagePrompt,
+    Prompt,
+    Text,
+    TextControl,
+    TokenControl,
+    Tokens,
+)
 from .aleph_alpha_client import (
-    AlephAlphaClient,
-    QuotaError,
     POOLING_OPTIONS,
+    AlephAlphaClient,
     AsyncClient,
     Client,
+    QuotaError,
 )
 from .aleph_alpha_model import AlephAlphaModel
-from .image import Image, ImagePrompt, ImageControl
-from .prompt import Prompt, Tokens, TokenControl, Text, TextControl
+from .completion import CompletionRequest, CompletionResponse
+from .detokenization import DetokenizationRequest, DetokenizationResponse
+from .document import Document
+from .embedding import (
+    EmbeddingRequest,
+    EmbeddingResponse,
+    SemanticEmbeddingRequest,
+    SemanticEmbeddingResponse,
+    SemanticRepresentation,
+)
+from .evaluation import EvaluationRequest, EvaluationResponse
 from .explanation import (
-    ExplanationRequest,
+    CustomGranularity,
+    Explanation,
     ExplanationPostprocessing,
+    ExplanationRequest,
     ExplanationResponse,
-    TargetGranularity,
-    CustomGranularity,
-    TextScore,
+    ImagePromptItemExplanation,
     ImageScore,
+    TargetGranularity,
+    TargetPromptItemExplanation,
     TargetScore,
-    TokenScore,
-    ImagePromptItemExplanation,
     TextPromptItemExplanation,
-    TargetPromptItemExplanation,
+    TextScore,
     TokenPromptItemExplanation,
-    Explanation,
-)
-from .embedding import (
-    EmbeddingRequest,
-    EmbeddingResponse,
-    SemanticEmbeddingRequest,
-    SemanticEmbeddingResponse,
-    SemanticRepresentation,
+    TokenScore,
 )
-from .completion import CompletionRequest, CompletionResponse
 from .qa import QaRequest, QaResponse
-from .evaluation import EvaluationRequest, EvaluationResponse
-from .tokenization import TokenizationRequest, TokenizationResponse
-from .detokenization import DetokenizationRequest, DetokenizationResponse
-from .summarization import SummarizationRequest, SummarizationResponse
 from .search import SearchRequest, SearchResponse, SearchResult
-from .utils import load_base64_from_url, load_base64_from_file
-from .document import Document
+from .summarization import SummarizationRequest, SummarizationResponse
+from .tokenization import TokenizationRequest, TokenizationResponse
+from .utils import load_base64_from_file, load_base64_from_url
 from .version import __version__
 
 __all__ = [
-    "POOLING_OPTIONS",
     "AlephAlphaClient",
     "AlephAlphaModel",
     "AsyncClient",
     "Client",
     "CompletionRequest",
     "CompletionResponse",
+    "ControlTokenOverlap",
+    "CustomGranularity",
     "DetokenizationRequest",
     "DetokenizationResponse",
     "Document",
     "EmbeddingRequest",
     "EmbeddingResponse",
     "EvaluationRequest",
     "EvaluationResponse",
-    "CustomGranularity",
-    "TextScore",
-    "ImageScore",
-    "TargetScore",
-    "TokenScore",
-    "ImagePromptItemExplanation",
-    "TextPromptItemExplanation",
-    "TargetPromptItemExplanation",
-    "TokenPromptItemExplanation",
     "Explanation",
     "ExplanationPostprocessing",
     "ExplanationRequest",
     "ExplanationResponse",
     "Image",
     "ImageControl",
     "ImagePrompt",
+    "ImagePromptItemExplanation",
+    "ImageScore",
+    "POOLING_OPTIONS",
     "Prompt",
     "QaRequest",
     "QaResponse",
@@ -86,10 +90,16 @@
     "SummarizationRequest",
     "SummarizationResponse",
     "TargetGranularity",
+    "TargetPromptItemExplanation",
+    "TargetScore",
     "Text",
     "TextControl",
+    "TextPromptItemExplanation",
+    "TextScore",
+    "TokenControl",
     "TokenizationRequest",
     "TokenizationResponse",
-    "TokenControl",
+    "TokenPromptItemExplanation",
     "Tokens",
+    "TokenScore",
 ]

aleph_alpha_client/aleph_alpha_client.py

@@ -20,7 +20,7 @@
     ExplanationRequest,
     ExplanationResponse,
 )
-from aleph_alpha_client.image import Image
+from aleph_alpha_client import Image
 from aleph_alpha_client.prompt import _to_json, _to_serializable_prompt
 from aleph_alpha_client.summarization import SummarizationRequest, SummarizationResponse
 from aleph_alpha_client.qa import QaRequest, QaResponse

aleph_alpha_client/document.py

@@ -1,8 +1,7 @@
 import base64
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, Optional, Sequence, Union
 
-from aleph_alpha_client.image import Image
-from aleph_alpha_client.prompt import PromptItem, Text, _to_json
+from aleph_alpha_client.prompt import Image, PromptItem, Text, _to_json
 
 
 class Document:

aleph_alpha_client/explanation.py

@@ -11,9 +11,8 @@
 
 # Import Literal with Python 3.7 fallback
 from typing_extensions import Literal
-from aleph_alpha_client.image import Image
 
-from aleph_alpha_client.prompt import Prompt, PromptItem
+from aleph_alpha_client.prompt import ControlTokenOverlap, Image, Prompt, PromptItem
 
 
 class ExplanationPostprocessing(Enum):
@@ -34,6 +33,17 @@ def to_json(self) -> str:
 
 
 class CustomGranularity(NamedTuple):
+    """
+    Allows for passing a custom delimiter to determine the granularity to
+    to explain the prompt by. The text of the prompt will be split by the
+    delimiter you provide.
+
+    Parameters:
+        delimiter (str, required):
+            String to split the text in the prompt by for generating
+            explanations for your prompt.
+    """
+
     delimiter: str
 
     def to_json(self) -> Mapping[str, Any]:
@@ -76,10 +86,78 @@ def to_json(self) -> str:
 
 
 class ExplanationRequest(NamedTuple):
+    """
+    Describes an Explanation request you want to make agains the API.
+
+    Parameters:
+        prompt (Prompt, required)
+            Prompt you want to generate explanations for a target completion.
+        target (str, required)
+            The completion string to be explained based on model probabilities.
+        contextual_control_threshold (float, default None)
+            If set to None, attention control parameters only apply to those tokens that have
+            explicitly been set in the request.
+            If set to a non-None value, we apply the control parameters to similar tokens as well.
+            Controls that have been applied to one token will then be applied to all other tokens
+            that have at least the similarity score defined by this parameter.
+            The similarity score is the cosine similarity of token embeddings.
+        control_factor (float, default None):
+            The amount to adjust model attention by.
+            For Explanation, you want to supress attention, and the API will default to 0.1.
+            Values between 0 and 1 will supress attention.
+            A value of 1 will have no effect.
+            Values above 1 will increase attention.
+        control_token_overlap (ControlTokenOverlap, default None)
+            What to do if a control partially overlaps with a text or image token.
+            If set to "partial", the factor will be adjusted proportionally with the amount
+            of the token it overlaps. So a factor of 2.0 of a control that only covers 2 of
+            4 token characters, would be adjusted to 1.5.
+            If set to "complete", the full factor will be applied as long as the control
+            overlaps with the token at all.
+        control_log_additive (bool, default None)
+            True: apply control by adding the log(control_factor) to attention scores.
+            False: apply control by (attention_scores - - attention_scores.min(-1)) * control_factor
+            If None, the API will default to True
+        prompt_granularity (PromptGranularity, default None)
+            At which granularity should the target be explained in terms of the prompt.
+            If you choose, for example, "sentence" then we report the importance score of each
+            sentence in the prompt towards generating the target output.
+
+            If you do not choose a granularity then we will try to find the granularity that
+            brings you closest to around 30 explanations. For large documents, this would likely
+            be sentences. For short prompts this might be individual words or even tokens.
+
+            If you choose a custom granularity then you must provide a custom delimiter. We then
+            split your prompt by that delimiter. This might be helpful if you are using few-shot
+            prompts that contain stop sequences.
+
+            For image prompt items, the granularities determine into how many tiles we divide
+            the image for the explanation.
+            "token" -> 12x12
+            "word" -> 6x6
+            "sentence" -> 3x3
+            "paragraph" -> 1
+        target_granularity (TargetGranularity, default None)
+            How many explanations should be returned in the output.
+
+            "complete" -> Return one explanation for the entire target. Helpful in many cases to determine which parts of the prompt contribute overall to the given completion.
+            "token" -> Return one explanation for each token in the target.
+
+            If None, API will default to "complete"
+        postprocessing (ExplanationPostprocessing, default None)
+            Optionally apply postprocessing to the difference in cross entropy scores for each token.
+            "none": Apply no postprocessing.
+            "absolute": Return the absolute value of each value.
+            "square": Square each value
+        normalize (bool, default None)
+            Return normalized scores. Minimum score becomes 0 and maximum score becomes 1. Applied after any postprocessing
+    """
+
     prompt: Prompt
     target: str
     contextual_control_threshold: Optional[float] = None
     control_factor: Optional[float] = None
+    control_token_overlap: Optional[ControlTokenOverlap] = None
     control_log_additive: Optional[bool] = None
     prompt_granularity: Optional[PromptGranularity] = None
     target_granularity: Optional[TargetGranularity] = None
@@ -93,8 +171,10 @@ def to_json(self) -> Dict[str, Any]:
         }
         if self.contextual_control_threshold is not None:
             payload["contextual_control_threshold"] = self.contextual_control_threshold
-        if self.control_factor is not None:
-            payload["control_factor"] = self.control_factor
+        if self.control_token_overlap is not None:
+            payload["control_token_overlap"] = self.control_token_overlap.to_json()
+        if self.postprocessing is not None:
+            payload["postprocessing"] = self.postprocessing.to_json()
         if self.control_log_additive is not None:
             payload["control_log_additive"] = self.control_log_additive
         if self.prompt_granularity is not None: