Merge pull request #416 from explosion/develop

Synch `main` with `develop`

Author: rmitsch

.github/workflows/test.yml

@@ -24,8 +24,6 @@ jobs:
         os: [ubuntu-latest, windows-latest, macos-latest]
         python_version: ["3.11"]
         include:
-          - os: ubuntu-20.04
-            python_version: "3.6"
           - os: windows-latest
             python_version: "3.7"
           - os: macos-latest

.github/workflows/test_gpu.yml

@@ -1,8 +1,8 @@
 name: GPU tests
 
 on:
-  schedule:
-    - cron: "0 0 * * *"
+#  schedule:
+#    - cron: "0 0 * * *"
   issue_comment:
     types: [created]
   workflow_dispatch:

README.md

@@ -1,6 +1,11 @@
 <a href="https://explosion.ai"><img src="https://explosion.ai/assets/img/logo.svg" width="125" height="125" align="right" /></a>
+<a href="https://explosion.ai"><img src="assets/logo.png" width="125" height="125" align="left" style="margin-right:30px" /></a>
 
-# spacy-llm: Integrating LLMs into structured NLP pipelines
+<h1 align="center">
+<span style="font: bold 38pt'Courier New';">spacy-llm</span>
+<br>Structured NLP with LLMs
+</h1>
+<br><br>
 
 [![GitHub Workflow Status](https://img.shields.io/github/actions/workflow/status/explosion/spacy-llm/test.yml?branch=main)](https://github.com/explosion/spacy-llm/actions/workflows/test.yml)
 [![pypi Version](https://img.shields.io/pypi/v/spacy-llm.svg?style=flat-square&logo=pypi&logoColor=white)](https://pypi.org/project/spacy-llm/)
@@ -16,7 +21,8 @@ This package integrates Large Language Models (LLMs) into [spaCy](https://spacy.
   - **[OpenAI](https://platform.openai.com/docs/api-reference/)**
   - **[Cohere](https://docs.cohere.com/reference/generate)**
   - **[Anthropic](https://docs.anthropic.com/claude/reference/)**
-  - **[PaLM](https://ai.google/discover/palm2/)**
+  - **[Google PaLM](https://ai.google/discover/palm2/)**
+  - **[Microsoft Azure AI](https://azure.microsoft.com/en-us/solutions/ai)**
 - Supports open-source LLMs hosted on Hugging Face 🤗:
   - **[Falcon](https://huggingface.co/tiiuae)**
   - **[Dolly](https://huggingface.co/databricks)**
@@ -33,10 +39,13 @@ This package integrates Large Language Models (LLMs) into [spaCy](https://spacy.
   - Sentiment analysis
   - Span categorization
   - Summarization
+  - Entity linking
+  - Translation
+  - Raw prompt execution for maximum flexibility
   - Soon:
-    - Entity linking
     - Semantic role labeling
 - Easy implementation of **your own functions** via [spaCy's registry](https://spacy.io/api/top-level#registry) for custom prompting, parsing and model integrations
+- Map-reduce approach for splitting prompts too long for LLM's context window and fusing the results back together
 
 ## 🧠 Motivation
 

assets/logo.png

@@ -16,7 +16,6 @@ classifiers =
     Operating System :: MacOS :: MacOS X
     Operating System :: Microsoft :: Windows
     Programming Language :: Python :: 3
-    Programming Language :: Python :: 3.6
     Programming Language :: Python :: 3.7
     Programming Language :: Python :: 3.8
     Programming Language :: Python :: 3.9
@@ -30,7 +29,7 @@ project_urls =
 [options]
 zip_safe = false
 include_package_data = true
-python_requires = >=3.6
+python_requires = >=3.7
 install_requires =
     spacy>=3.5,<4.0
     jinja2

setup.cfg

@@ -8,7 +8,7 @@
 from spacy.vocab import Vocab
 
 from .registry import registry
-from .ty import LLMTask, PromptTemplateProvider
+from .ty import PromptTemplateProvider, ShardingLLMTask
 
 
 @registry.llm_misc("spacy.BatchCache.v1")
@@ -68,11 +68,11 @@ def __init__(
 
         self._init_cache_dir()
 
-    def initialize(self, vocab: Vocab, task: LLMTask) -> None:
+    def initialize(self, vocab: Vocab, task: ShardingLLMTask) -> None:
         """
         Initialize cache with data not available at construction time.
         vocab (Vocab): Vocab object.
-        task (LLMTask): Task.
+        task (ShardingLLMTask): Task.
         """
         self._vocab = vocab
         if isinstance(task, PromptTemplateProvider):

spacy_llm/cache.py

@@ -57,5 +57,7 @@
 if PYDANTIC_V2:
     from pydantic.v1 import BaseModel, ExtraError, ValidationError  # noqa: F401
     from pydantic.v1 import validator
+    from pydantic.v1.generics import GenericModel  # noqa: F401
 else:
     from pydantic import BaseModel, ExtraError, ValidationError, validator  # noqa: F401
+    from pydantic.generics import GenericModel  # noqa: F401

spacy_llm/compat.py

@@ -17,16 +17,18 @@ def __init__(
         name: str,
         config_init: Optional[Dict[str, Any]],
         config_run: Optional[Dict[str, Any]],
+        context_length: Optional[int],
     ):
         """Initializes HF model instance.
         query (Callable[[Any, Iterable[Any]], Iterable[Any]): Callable executing LLM prompts when
             supplied with the `integration` object.
         name (str): Name of HF model to load (without account name).
         config_init (Optional[Dict[str, Any]]): HF config for initializing the model.
         config_run (Optional[Dict[str, Any]]): HF config for running the model.
-        inference_config (Dict[Any, Any]): HF config for model run.
+        context_length (Optional[int]): Context length for this model. Necessary for sharding.
         """
         self._name = name if self.hf_account in name else f"{self.hf_account}/{name}"
+        self._context_length = context_length
         default_cfg_init, default_cfg_run = self.compile_default_configs()
         self._config_init, self._config_run = default_cfg_init, default_cfg_run
 
@@ -73,10 +75,10 @@ def __init__(
         self._model = self.init_model()
 
     @abc.abstractmethod
-    def __call__(self, prompts: Iterable[Any]) -> Iterable[Any]:
+    def __call__(self, prompts: Iterable[Iterable[Any]]) -> Iterable[Iterable[Any]]:
         """Executes prompts on specified API.
-        prompts (Iterable[Any]): Prompts to execute.
-        RETURNS (Iterable[Any]): API responses.
+        prompts (Iterable[Iterable[Any]]): Prompts to execute per doc.
+        RETURNS (Iterable[Iterable[Any]]): API responses per doc.
         """
 
     def _check_model(self) -> None:
@@ -93,6 +95,13 @@ def get_model_names(cls) -> Tuple[str, ...]:
         """
         return tuple(str(arg) for arg in cls.MODEL_NAMES.__args__)  # type: ignore[attr-defined]
 
+    @property
+    def context_length(self) -> Optional[int]:
+        """Returns context length in number of tokens for this model.
+        RETURNS (Optional[int]): Max. number of tokens allowed in prompt for the current model.
+        """
+        return self._context_length
+
     @property
     @abc.abstractmethod
     def hf_account(self) -> str:

spacy_llm/models/hf/base.py

@@ -18,14 +18,18 @@ def init_model(self) -> Any:
             model=self._name, return_full_text=False, **self._config_init
         )
 
-    def __call__(self, prompts: Iterable[str]) -> Iterable[str]:  # type: ignore[override]
+    def __call__(self, prompts: Iterable[Iterable[str]]) -> Iterable[Iterable[str]]:  # type: ignore[override]
         """Queries Dolly HF model.
         pipeline (transformers.pipeline): Transformers pipeline to query.
-        prompts (Iterable[str]): Prompts to query Dolly model with.
-        RETURNS (Iterable[str]): Prompt responses.
+        prompts (Iterable[Iterable[str]]): Prompts per doc to query Dolly model with.
+        RETURNS (Iterable[Iterable[str]]): Prompt responses per doc.
         """
         return [
-            self._model(pr, **self._config_run)[0]["generated_text"] for pr in prompts
+            [
+                self._model(pr, **self._config_run)[0]["generated_text"]
+                for pr in prompts_for_doc
+            ]
+            for prompts_for_doc in prompts
         ]
 
     @property
@@ -52,12 +56,14 @@ def dolly_hf(
     name: Dolly.MODEL_NAMES,
     config_init: Optional[Dict[str, Any]] = SimpleFrozenDict(),
     config_run: Optional[Dict[str, Any]] = SimpleFrozenDict(),
-) -> Callable[[Iterable[str]], Iterable[str]]:
+) -> Callable[[Iterable[Iterable[str]]], Iterable[Iterable[str]]]:
     """Generates Dolly instance that can execute a set of prompts and return the raw responses.
     name (Literal): Name of the Dolly model. Has to be one of Dolly.get_model_names().
     config_init (Optional[Dict[str, Any]]): HF config for initializing the model.
     config_run (Optional[Dict[str, Any]]): HF config for running the model.
     RETURNS (Callable[[Iterable[str]], Iterable[str]]): Dolly instance that can execute a set of prompts and return
         the raw responses.
     """
-    return Dolly(name=name, config_init=config_init, config_run=config_run)
+    return Dolly(
+        name=name, config_init=config_init, config_run=config_run, context_length=2048
+    )

spacy_llm/models/hf/dolly.py

@@ -17,9 +17,15 @@ def __init__(
         name: MODEL_NAMES,
         config_init: Optional[Dict[str, Any]],
         config_run: Optional[Dict[str, Any]],
+        context_length: Optional[int],
     ):
         self._tokenizer: Optional["transformers.AutoTokenizer"] = None
-        super().__init__(name=name, config_init=config_init, config_run=config_run)
+        super().__init__(
+            name=name,
+            config_init=config_init,
+            config_run=config_run,
+            context_length=context_length,
+        )
 
         assert isinstance(self._tokenizer, transformers.PreTrainedTokenizerBase)
         self._config_run["pad_token_id"] = self._tokenizer.pad_token_id
@@ -45,10 +51,15 @@ def init_model(self) -> Any:
     def hf_account(self) -> str:
         return "tiiuae"
 
-    def __call__(self, prompts: Iterable[str]) -> Iterable[str]:  # type: ignore[override]
+    def __call__(self, prompts: Iterable[Iterable[str]]) -> Iterable[Iterable[str]]:  # type: ignore[override]
         return [
-            self._model(pr, generation_config=self._hf_config_run)[0]["generated_text"]
-            for pr in prompts
+            [
+                self._model(pr, generation_config=self._hf_config_run)[0][
+                    "generated_text"
+                ]
+                for pr in prompts_for_doc
+            ]
+            for prompts_for_doc in prompts
         ]
 
     @staticmethod
@@ -68,12 +79,14 @@ def falcon_hf(
     name: Falcon.MODEL_NAMES,
     config_init: Optional[Dict[str, Any]] = SimpleFrozenDict(),
     config_run: Optional[Dict[str, Any]] = SimpleFrozenDict(),
-) -> Callable[[Iterable[str]], Iterable[str]]:
+) -> Callable[[Iterable[Iterable[str]]], Iterable[Iterable[str]]]:
     """Generates Falcon instance that can execute a set of prompts and return the raw responses.
     name (Literal): Name of the Falcon model. Has to be one of Falcon.get_model_names().
     config_init (Optional[Dict[str, Any]]): HF config for initializing the model.
     config_run (Optional[Dict[str, Any]]): HF config for running the model.
     RETURNS (Callable[[Iterable[str]], Iterable[str]]): Falcon instance that can execute a set of prompts and return
         the raw responses.
     """
-    return Falcon(name=name, config_init=config_init, config_run=config_run)
+    return Falcon(
+        name=name, config_init=config_init, config_run=config_run, context_length=2048
+    )