Classification prompt tuning

CODEOWNERS

@@ -8,4 +8,4 @@
 #  https://help.github.com/en/articles/about-code-owners
 #
 
-*  @alex-jw-brooks @gkumbhat @evaline-ju
+*  @alex-jw-brooks @gkumbhat @evaline-ju @gabe-l-hart

README.md

@@ -1,11 +1,36 @@
 # Caikit NLP
 
+Welcome to the README.md page of `caikit-nlp`! This repository provides NLP domain capabilities running on [caikit](https://github.com/caikit/caikit) framework.
+
+## Introduction
+
 `caikit_nlp` is a [Caikit](https://github.com/caikit/caikit) library that currently provides [PEFT prompt tuning](https://github.com/huggingface/peft) and MPT (multi-task prompt tuning) functionalities.
 
-More information on MPT can be found at: https://arxiv.org/abs/2303.02861
+### Getting Started
+
+To help you quickly get started with using Caikit, we have prepared a [Jupyter notebook](examples/Caikit_Getting_Started.ipynb) that can be run in Google Colab. Caikit-nlp is a powerful library that leverages prompt tuning and fine-tuning to add NLP domain capabilities to caikit.
+
+
+### Contributing
+
+We welcome contributions from the community! If you would like to contribute to `caikit-nlp`, please read the guidelines in the main project's [CONTRIBUTING.md](main/CONTRIBUTING.md) file. It includes information on submitting bug reports, feature requests, and pull requests. Make sure to follow our coding standards, [code of conduct](code-of-conduct.md), [security standards](https://github.com/caikit/community/blob/main/SECURITY.md), and documentation guidelines to streamline the contribution process.
+
+### License
+
+This project is licensed under the [ASFv2 License](LICENSE).
+
+### Glossary
+
+A list of terms that either may be unfamiliar or that have nebulous definitions based on who and where you hear them, defined for how they are used/thought of in the `caikit`/`caikit-nlp` project:
+
+* Fine tuning - trains the base model onto new data etc; this changes the base model.
+* Prompt engineering - (usually) manually crafting texts that make models do a better job that's left appended to the input text. E.g., if you wanted to do something like sentiment on movie reviews, you might come up with a prompt like The movie was: _____ and replace the _____  with the movie review you're consider to try to get something like happy/sad out of it.
+* PEFT - library by Huggingface containing implementations of different tuning methods that scale well - things like prompt tuning, and MPT live there. So PEFT itself isn't an approach even though parameter efficient fine-tuning sounds like one.
+Prompt tuning - learning soft prompts. This is different from prompt engineering in that you're not trying to learn tokens. Instead, you're basically trying to learn new embedded representations (sometimes called virtual tokens) that can be concatenated onto your embedded input text to improve the performance. This can work well, but also can be sensitive to initialization.
+* Multitask prompt tuning (MPT) - Tries to fix some of the issues with prompt tuning by allowing you to effectively learn 'source prompts' across different tasks & leverage them to initialize your prompt tuning etc. More information on MPT can be found at: https://arxiv.org/abs/2303.02861
 
-Currently causal language models and sequence-to-sequence models are supported.
+The important difference between fine tuning and capabilities like prompt tuning/multi-taskprompt tuning is that the latter doesn't change the base model's weights at all. So when you run inference for prompt tuned models, you can have n prompts to 1 base model, and just inject the prompt tensors you need when they're requested instead of having _n_ separate fine-tuned models.
 
 #### Notes
 
-- The data model for text generative capabilities is baked into this repository itself at `caikit_nlp/data_model/generation.py`.
+- Currently causal language models and sequence-to-sequence models are supported.

caikit_nlp/__init__.py

@@ -24,7 +24,7 @@
 
 # Local
 # Import subpackages
-from . import config, data_model
+from . import config, data_model, model_management
 from .config import *
 from .data_model import *
 from .modules import *

caikit_nlp/data_model/__init__.py

@@ -15,7 +15,5 @@
 """
 
 # Local
-from . import classification, generation, text
-from .classification import *
+from . import generation
 from .generation import *
-from .text import *

caikit_nlp/modules/tokenization/tokenization_task.py → caikit_nlp/model_management/__init__.py

@@ -12,16 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# First Party
-from caikit.core import TaskBase, task
-
 # Local
-from ...data_model import TokenizationResult
-
-
-@task(
-    required_parameters={"text": str},
-    output_type=TokenizationResult,
-)
-class TokenizationTask(TaskBase):
-    pass
+from .tgis_auto_finder import TGISAutoFinder

caikit_nlp/model_management/tgis_auto_finder.py

@@ -0,0 +1,142 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+The TGISAutoFinder implements the ModelFinder interface to provide automatic
+discovery of text-generation models that can be auto-configured to run against
+a remote TGIS model.
+"""
+# Standard
+from typing import Optional
+
+# First Party
+from caikit.core import MODEL_MANAGER, error_handler
+from caikit.core.model_management import ModelFinderBase, model_finder_factory
+from caikit.core.modules import ModuleConfig
+from caikit_tgis_backend import TGISBackend
+import aconfig
+import alog
+
+# Local
+from ..modules.text_generation import TextGenerationTGIS
+
+log = alog.use_channel("TGIS_FND")
+error = error_handler.get(log)
+
+
+class TGISAutoFinder(ModelFinderBase):
+    __doc__ = __doc__
+
+    name = "TGIS-AUTO"
+
+    # Constants for the keys of the config blob
+    _LOCAL_INITIALIZER_NAME_KEY = "local_initializer_name"
+    _TGIS_BACKEND_PRIORITY_KEY = "tgis_backend_priority"
+
+    def __init__(self, config: aconfig.Config, instance_name: str = ""):
+        """Initialize from the model finder factory config
+
+        Config schema:
+
+        local_initializer_name:
+            type: string
+            default: "default"
+            description: The name within the initializers config for the LOCAL
+                initializer that will hold the tgis backend to use
+
+        tgis_backend_priority:
+            type: integer
+            description: Index within the backend_priority list for the TGIS
+                backend to use. If not set, the first TGIS backend found will be
+                used.
+
+        Args:
+            config (aconfig.Config): The configuration blob from caikit's
+                model_management factory construction
+            instance_name (str): The name of this finder instance
+        """
+        local_initializer_name = config.get(self._LOCAL_INITIALIZER_NAME_KEY, "default")
+        tgis_backend_priority = config.get(self._TGIS_BACKEND_PRIORITY_KEY)
+        error.type_check(
+            "<NLP97312902E>", str, local_initializer_name=local_initializer_name
+        )
+        error.type_check(
+            "<NLP97312903E>",
+            int,
+            tgis_backend_priority=tgis_backend_priority,
+            allow_none=True,
+        )
+
+        # Extract the TGIS backend instance
+        local_initializer = MODEL_MANAGER.get_initializer(local_initializer_name)
+        backends = local_initializer.backends
+        if tgis_backend_priority is not None:
+            error.value_check(
+                "<NLP87928813E>",
+                0 <= tgis_backend_priority < len(backends),
+                "Invalid {}: {}",
+                self._TGIS_BACKEND_PRIORITY_KEY,
+                tgis_backend_priority,
+            )
+            self._tgis_backend = backends[tgis_backend_priority]
+            error.value_check(
+                "<NLP77150201E>",
+                self._tgis_backend.backend_type == TGISBackend.backend_type,
+                "Index {} is not a TGIS backend",
+                tgis_backend_priority,
+            )
+        else:
+            tgis_backend = None
+            for backend in backends:
+                if backend.backend_type == TGISBackend.backend_type:
+                    tgis_backend = backend
+                    break
+            error.value_check(
+                "<NLP96294266E>",
+                tgis_backend is not None,
+                "No TGIS backend found!",
+            )
+            self._tgis_backend = tgis_backend
+
+    def find_model(
+        self,
+        model_path: str,
+        **kwargs,
+    ) -> Optional[ModuleConfig]:
+        """Find the model if"""
+
+        # Get a connection to this model in tgis
+        log.debug2("Attempting to setup TGIS client for %s", model_path)
+        if self._tgis_backend.get_connection(model_id=model_path) is None:
+            log.debug2("TGIS cannot connect to model %s", model_path)
+            return None
+
+        # If connection is ok, set up the module config to point to the remote
+        # TGIS text generation module
+        cfg = ModuleConfig(
+            {
+                "module_id": TextGenerationTGIS.MODULE_ID,
+                "module_class": TextGenerationTGIS.MODULE_CLASS,
+                "name": TextGenerationTGIS.MODULE_NAME,
+                "version": TextGenerationTGIS.MODULE_VERSION,
+                "model_name": model_path,
+            }
+        )
+        # Set a special indicator in the module config to use the backend that
+        # this finder found. This will override the backend found by the local
+        # initializer.
+        cfg.tgis_backend = self._tgis_backend
+        return cfg
+
+
+model_finder_factory.register(TGISAutoFinder)

caikit_nlp/modules/text_classification/__init__.py

@@ -14,4 +14,3 @@
 
 # Local
 from .sequence_classification import SequenceClassification
-from .text_classification_task import TextClassificationTask