Caikit NLP

Caikit-NLP is a python library providing various Natural Language Processing (NLP) capabilities built on top of caikit framework.

Introduction

Caikit-NLP implements concept of "task" from caikit framework to define (and consume) interfaces for various NLP problems and implements various "modules" to provide functionalities for these "modules".

Capabilities provided by caikit-nlp:

Getting Started

Notebooks

To help you quickly get started with using Caikit, we have prepared a Jupyter notebook 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.

Installation

To install from git repo:

python -m venv .venv
source .venv/bin/activate
pip install git+https://github.com/caikit/caikit-nlp

Bootstrapping models

caikit_nlp can use Hugging Face models, allowing for direct download and bootstrapping.

For example, to use google/flan-t5-small:

import os
# The env var ALLOW_DOWNLOADS has to be set to allow model downloads before importing caikit_nlp
os.environ['ALLOW_DOWNLOADS'] = "1"

import caikit_nlp

model_name = "google/flan-t5-small"
model = caikit_nlp.text_generation.TextGeneration.bootstrap(model_name)
model.save(f"{model_name}-caikit") # optionally save the model

Serving models

To serve models, the following basic configuration can be used:

# config.yml
runtime:
  library: caikit_nlp
  local_models_dir: ./models

log:
  formatter: pretty # optional: log formatter is set to json by default

Start the server:

env CONFIG_FILES=./config.yml python -m caikit.runtime

The model can now be queried at localhost:8080 via http or at localhost:8085 via grpc.

For example, using the http server and using curl to send a POST request:

curl --json '{
    "model_id": "flan-t5-small-caikit",
    "inputs": "At what temperature does liquid Nitrogen boil?"
}' localhost:8080/api/v1/task/text-generation

We get the following response:

{
  "generated_text": "74 degrees F",
  "generated_tokens": 5,
  "finish_reason": "MAX_TOKENS",
  "producer_id": {
    "name": "Text Generation",
    "version": "0.1.0"
  },
  "input_token_count": 10,
  "seed": null
}

All the available API endpoints and protos can be dumped using scripts/dump_apis.sh.

Docker

To build the docker image:

python -m build --wheel
docker build -t caikit-nlp:latest .

A volume can be mounted at /caikit providing configuration and (optionally) models:

mkdir -p caikit
$EDITOR caikit/config.yml # edit as required
cp -r