Use prompt template to ensure data consistency in cache directory (#192)

* Fix cache and nlp.pipe(..., as_tuples=True) (#188)

* ensure doc from cache has the same context as the passed in doc. Needed for nlp.pipe(egs, as_tuples=True)

* add test for as_tuples and nlp.pipe

* Update README.md (#189)

* Update README.md (#190)

* Update README.md

* Update README.md

* Update README.md

---------

Co-authored-by: Raphael Mitsch <r.mitsch@outlook.com>

* Fix issue of last cached element in batch not having IO data (#191)

* Fix issue of last cached element in batch not having IO data.

* Minor formatting change.

* Bump version.

* Use prompt template to ensure data consistency in cache directory.

* Fix invalid config test.

* Fix remaining test tasks without prompt_template().

* Apply suggestions from code review

Co-authored-by: Madeesh Kannan <shadeMe@users.noreply.github.com>

* Fix formatting issues.

* Replace normalization with hash.

* Add Cache.initialize().

* Update README.md

Co-authored-by: Sofie Van Landeghem <svlandeg@users.noreply.github.com>

* Move prompt template availability to different protocol. Extend tests.

* Add section on prompt_template() in readme.

* Update spacy_llm/cache.py

Co-authored-by: Sofie Van Landeghem <svlandeg@users.noreply.github.com>

---------

Co-authored-by: Kabir Khan <kabir@explosion.ai>
Co-authored-by: vincent d warmerdam <vincentwarmerdam@gmail.com>
Co-authored-by: Madeesh Kannan <shadeMe@users.noreply.github.com>
Co-authored-by: svlandeg <svlandeg@github.com>
Co-authored-by: Sofie Van Landeghem <svlandeg@users.noreply.github.com>

Author: 5 people

README.md

@@ -1474,23 +1474,37 @@ The default `query` (`spacy.CallLangChain.v1`) executes the prompts by running `
 
 Interacting with LLMs, either through an external API or a local instance, is costly.
 Since developing an NLP pipeline generally means a lot of exploration and prototyping,
-`spacy-llm` implements a built-in cache to avoid reprocessing the same documents at each run.
+`spacy-llm` implements a built-in cache to avoid reprocessing the same documents at each run
+that keeps batches of documents stored on disk. 
+
+The cache implementation also ensures that documents in one cache directory were all produced using the same prompt 
+template. This is only possible however if the specified task implements 
+```python
+@property
+def prompt_template() -> str:
+    ...
+``` 
+which returns the raw prompt template as string. If `prompt_template()` isn't implemented, the cache will emit a warning
+and not check for prompt template consistency.
 
 Example config block:
 
 ```ini
 [components.llm.cache]
-@llm_misc = "spacy.BatchCache.v1",
+@llm_misc = "spacy.BatchCache.v1"
 path = "path/to/cache"
 batch_size = 64
 max_batches_in_mem = 4
 ```
 
-| Argument             | Type                         | Default | Description                                          |
-| -------------------- | ---------------------------- | ------- | ---------------------------------------------------- |
-| `path`               | `Optional[Union[str, Path]]` | `None`  | Cache directory. If `None`, no caching is performed. |
-| `batch_size`         | `int`                        | 64      | Number of docs in one batch (file).                  |
-| `max_batches_in_mem` | `int`                        | 4       | Max. number of batches to hold in memory.            |
+| Argument             | Type                         | Default | Description                                                                                                               |
+| -------------------- | ---------------------------- | ------- | ------------------------------------------------------------------------------------------------------------------------- |
+| `path`               | `Optional[Union[str, Path]]` | `None`  | Cache directory. If `None`, no caching is performed, and this component will act as a NoOp.                               |
+| `batch_size`         | `int`                        | 64      | Number of docs in one batch (file). Once a batch is full, it will be persisted to disk.                                   |
+| `max_batches_in_mem` | `int`                        | 4       | Max. number of batches to hold in memory. Allows you to limit the effect on your memory if you're handling a lot of docs. |
+
+When retrieving a document, the `BatchCache` will first figure out what batch the document belongs to. If the batch
+isn't in memory it will try to load the batch from disk and then move it into memory. 
 
 Note that since the cache is generated by a registered function, you can also provide your own registered function
 returning your own cache implementation. If you wish to do so, ensure that your cache object adheres to the

spacy_llm/cache.py

@@ -1,3 +1,4 @@
+import warnings
 from pathlib import Path
 from typing import Dict, Iterable, List, Optional, Union
 
@@ -7,6 +8,7 @@
 from spacy.vocab import Vocab
 
 from .registry import registry
+from .ty import LLMTask, PromptTemplateProvider
 
 
 @registry.llm_misc("spacy.BatchCache.v1")
@@ -31,7 +33,7 @@ def __init__(
         batch_size: int,
         max_batches_in_mem: int,
     ):
-        """Initialize Cache instance.
+        """Initialize Cache instance. Acts as NoOp if path is None.
         path (Optional[Union[str,Path]]): Cache directory.
         batch_size (int): Number of docs in one batch (file).
         max_batches_in_mem (int): Max. number of batches to hold in memory.
@@ -43,6 +45,8 @@ def __init__(
         # Max. number of batches to keep in memory.
         self.max_batches_in_mem = max_batches_in_mem
         self._vocab: Optional[Vocab] = None
+        self._prompt_template: Optional[str] = None
+        self._prompt_template_checked: bool = False
 
         # Stores doc hash -> batch hash to allow efficient lookup of available Docs.
         self._doc2batch: Dict[int, int] = {}
@@ -62,31 +66,71 @@ def __init__(
             "persisted": 0,
         }
 
-        self._init_cache_index()
+        self._init_cache_dir()
+
+    def initialize(self, vocab: Vocab, task: LLMTask) -> None:
+        """
+        Initialize cache with data not available at construction time.
+        vocab (Vocab): Vocab object.
+        task (LLMTask): Task.
+        """
+        self._vocab = vocab
+        if isinstance(task, PromptTemplateProvider):
+            self.prompt_template = task.prompt_template
+        else:
+            self.prompt_template = ""
+            if self._path:
+                warnings.warn(
+                    "The specified task does not provide its prompt template via `prompt_template()`. This means that "
+                    "the cache cannot verify whether all cached documents were generated using the same prompt "
+                    "template."
+                )
 
     @property
-    def vocab(self) -> Optional[Vocab]:
-        """Vocab used for deserializing docs.
-        RETURNS (Vocab): Vocab used for deserializing docs.
+    def prompt_template(self) -> Optional[str]:
+        """Get prompt template.
+        RETURNS (Optional[str]): Prompt template string used for docs to cache/cached docs.
         """
-        return self._vocab
+        return self._prompt_template
 
-    @vocab.setter
-    def vocab(self, vocab: Vocab) -> None:
-        """Set vocab to use for deserializing docs.
-        vocab (Vocab): Vocab to use for deserializing docs.
+    @prompt_template.setter
+    def prompt_template(self, prompt_template: str) -> None:
+        """Set prompt template.
+        prompt_template (str): Prompt template string used for docs to cache/cached docs.
         """
-        self._vocab = vocab
+        self._prompt_template = prompt_template
+        if not self._path:
+            return
+
+        # Store prompt template as file.
+        prompt_template_path = self._path / "prompt_template.txt"
+        # If no file exists: store in plain text for easier debugging.
+        if not prompt_template_path.exists():
+            with open(prompt_template_path, "w") as file:
+                file.write(self._prompt_template)
+
+        else:
+            # If file exists and prompt template is not equal to new prompt template: raise, as this indicates we are
+            # trying to reuse a cache directory with a changed prompt.
+            with open(prompt_template_path, "r") as file:
+                existing_prompt_template = "".join(file.readlines())
+            if hash(existing_prompt_template) != hash(self._prompt_template):
+                raise ValueError(
+                    f"Prompt template in cache directory ({prompt_template_path}) is not equal with "
+                    f"current prompt template. Reset your cache if you are using a new prompt "
+                    f"template."
+                )
 
-    def _init_cache_index(self) -> None:
-        """Init cache index and directory."""
+    def _init_cache_dir(self) -> None:
+        """Init cache directory with index and prompt template file."""
         if self._path is None:
             return
 
         if self._path.exists() and not self._path.is_dir():
             raise ValueError("Cache directory exists and is not a directory.")
         self._path.mkdir(parents=True, exist_ok=True)
 
+        # Read from index file, if it exists.
         index_path = self._index_path
         if index_path.exists():
             for rec in srsly.read_jsonl(index_path):
@@ -136,6 +180,13 @@ def add(self, doc: Doc) -> None:
         """
         if self._path is None:
             return
+        if not self._prompt_template_checked and self._prompt_template is None:
+            warnings.warn(
+                "No prompt template set for Cache object, entailing that consistency of prompt template used "
+                "to generate docs cannot be checked. Be mindful to reset your cache whenever you change your "
+                "prompt template."
+            )
+            self._prompt_template_checked = True
 
         self._cache_queue.append(doc)
         self._stats["added"] += 1

spacy_llm/pipeline/llm.py

@@ -116,8 +116,8 @@ def __init__(
         self._task = task
         self._model = model
         self._cache = cache
-        self._cache.vocab = vocab
         self._save_io = save_io
+        self._cache.initialize(vocab, self._task)
 
         # This is done this way because spaCy's `validate_init_settings` function
         # does not support `**kwargs: Any`.
@@ -212,11 +212,10 @@ def _process_docs(self, docs: List[Doc]) -> List[Doc]:
             if is_cached[i]:
                 cached_doc = self._cache[doc]
                 assert cached_doc is not None
+                cached_doc._context = doc._context
                 final_docs.append(cached_doc)
             else:
                 doc = next(modified_docs)
-                self._cache.add(doc)
-                final_docs.append(doc)
 
                 if self._save_io:
                     # Make sure the `llm_io` field is set
@@ -227,6 +226,9 @@ def _process_docs(self, docs: List[Doc]) -> List[Doc]:
                     llm_io["prompt"] = str(next(prompts_iters[2]))
                     llm_io["response"] = str(next(responses_iters[2]))
 
+                self._cache.add(doc)
+                final_docs.append(doc)
+
         return final_docs
 
     def to_bytes(

spacy_llm/tasks/lemma.py

@@ -74,6 +74,10 @@ def initialize(
             if n_prompt_examples < 0 or len(self._prompt_examples) < n_prompt_examples:
                 self._prompt_examples.append(self._create_prompt_example(eg))
 
+    @property
+    def prompt_template(self) -> str:
+        return self._template
+
     def generate_prompts(self, docs: Iterable[Doc]) -> Iterable[str]:
         environment = jinja2.Environment()
         _template = environment.from_string(self._template)

spacy_llm/tasks/noop.py

@@ -22,3 +22,10 @@ def parse_responses(
     ) -> Iterable[Doc]:
         # Not doing anything
         return docs
+
+    @property
+    def prompt_template(self) -> str:
+        return """
+        This is the NoOp
+        prompt template
+        """

spacy_llm/tasks/rel.py

@@ -143,6 +143,10 @@ def _check_rel_extension(cls):
     def labels(self) -> Tuple[str, ...]:
         return tuple(self._label_dict.values())
 
+    @property
+    def prompt_template(self) -> str:
+        return self._template
+
     def generate_prompts(self, docs: Iterable[Doc]) -> Iterable[str]:
         environment = jinja2.Environment()
         _template = environment.from_string(self._template)

spacy_llm/tasks/span.py

@@ -90,6 +90,10 @@ def _check_label_consistency(self) -> List[SpanExample]:
     def labels(self) -> Tuple[str, ...]:
         return tuple(self._label_dict.values())
 
+    @property
+    def prompt_template(self) -> str:
+        return self._template
+
     def generate_prompts(self, docs: Iterable[Doc]) -> Iterable[str]:
         environment = jinja2.Environment()
         _template = environment.from_string(self._template)

spacy_llm/tasks/textcat.py

@@ -255,6 +255,10 @@ def __init__(
     def labels(self) -> Tuple[str, ...]:
         return tuple(self._label_dict.values())
 
+    @property
+    def prompt_template(self) -> str:
+        return self._template
+
     def generate_prompts(self, docs: Iterable[Doc]) -> Iterable[str]:
         environment = jinja2.Environment()
         _template = environment.from_string(self._template)

spacy_llm/tests/models/test_rest.py

@@ -21,15 +21,21 @@
 
 @registry.llm_tasks("spacy.Count.v1")
 class _CountTask:
+    _PROMPT_TEMPLATE = "Count the number of characters in this string: '{text}'."
+
     def generate_prompts(self, docs: Iterable[Doc]) -> Iterable[str]:
         for doc in docs:
-            yield f"Count the number of characters in this string: '{doc.text}'."
+            yield _CountTask._PROMPT_TEMPLATE.format(text=doc.text)
 
     def parse_responses(
         self, docs: Iterable[Doc], responses: Iterable[str]
     ) -> Iterable[Doc]:
         return docs
 
+    @property
+    def prompt_template(self) -> str:
+        return _CountTask._PROMPT_TEMPLATE
+
 
 def test_initialization():
     """Test initialization and simple run"""

spacy_llm/tests/pipeline/test_llm.py

@@ -101,6 +101,11 @@ def test_llm_pipe_with_cache(tmp_path: Path, n_process: int):
     docs = list(nlp.pipe(texts=texts, n_process=n_process))
     assert [doc.text for doc in docs] == texts
 
+    egs = [(text, i) for i, text in enumerate(texts)]
+    egs_processed = list(nlp.pipe(egs, as_tuples=True, n_process=n_process))
+    assert [doc.text for doc, _ in egs_processed] == texts
+    assert [eg for _, eg in egs_processed] == list(range(len(texts)))
+
 
 def test_llm_pipe_empty(nlp):
     """Test call .pipe() with empty batch."""