Architecture Layers

import caikit.config
my_config_defaults = {"foo": {"bar": 1, "baz": "two", "biz": True}}
caikit.configure(config_dict=my_config_defaults)
print(caikit.get_config().foo)

from caikit.core.toolkit.factory import ImportableFactory, FactoryConstructible
import abc
import aconfig
import caikit
import os

###########
## Bases ##
###########

class DatastoreClient(abc.ABC):
    @abc.abstractmethod
    def put_data(self, key: str, data: str):
        """Abstract put method"""


class DatastoreConnectorBase(FactoryConstructible):
    @abc.abstractmethod
    def get_client(self) -> DatastoreClient:
        """All connectors must be able to get a client handle"""

DATASTORE_FACTORY = ImportableFactory("datastore")

###########
## My DB ##
###########

class MyDBClient(DatastoreClient):

    def __init__(self, hostname: str, port: int):
        # STUB
        self._client = {"hostname": hostname, "port": port}

    def put_data(self, key: str, data: str):
        # STUB
        self._client.write(key, data)


class MyDBConnector(DatastoreConnectorBase):

    name = "MYDB"

    def __init__(self, config: aconfig.Config, instance_name: str):
        self.instance_name = instance_name
        self._client = MyDBClient(config.hostname, config.port)

    def get_client(self) -> MyDBClient:
        return self._client


DATASTORE_FACTORY.register(MyDBConnector)

#############
## Your DB ##
#############

class YourDBClient(DatastoreClient):

    def __init__(self, root_path: str):
        self._root = root_path
        if not os.path.isdir(self._root):
            os.makedirs(self._root)

    def put_data(self, key: str, data: str):
        with open(os.path.join(self._root, key), "w") as handle:
            handle.write(data)


class YourDBConnector(DatastoreConnectorBase):

    name = "YOURDB"

    def __init__(self, config: aconfig.Config, instance_name: str):
        self.instance_name = instance_name
        self._client = YourDBClient(config.root)

    def get_client(self) -> MyDBClient:
        return self._client


DATASTORE_FACTORY.register(YourDBConnector)

##########
## Main ##
##########

if __name__ == "__main__":
    config_defaults = {
        "datastore": {
            "type": "YOURDB",
            "config": {
                "root": ".yourdb"
            }
        }
    }
    caikit.configure(config_dict=config_defaults)
    inst = DATASTORE_FACTORY.construct(caikit.get_config().datastore)
    print(inst.instance_name)

from caikit.core.toolkit.concurrency.destroyable_process import DestroyableProcess
import time


def work(message: str = "Hello World", max_duration: int = 20):
    start_time = time.time()
    while time.time() < start_time + max_duration:
        print(f"Message: {message}")
        time.sleep(1)


proc = DestroyableProcess("fork", work)

proc.start()
time.sleep(3)
proc.destroy()
proc.join()
print(f"Destroyed? {proc.destroyed}")
print(f"Canceled? {proc.canceled}")
print(f"Ran? {proc.ran}")

from caikit.core.exceptions import error_handler
import alog
import caikit

log = alog.use_channel("DEMO")
error = error_handler.get(log)


@alog.logged_function(log.debug)
def doit(data: str, param: int) -> str:
    error.type_check("<DMO00255362E>", str, data=data)
    error.type_check("<DMO58840196E>", int, param=param)
    error.value_check(
        "<DMO62491029E>", param > 0, "Invalid param value (%d). Must be > 0", param
    )
    return data * param


custom_config = {"data": "Hello world", "param": 1}
caikit.configure(config_dict=custom_config)
cfg = caikit.get_config()
alog.configure(
    default_level=cfg.log.level,
    filters=cfg.log.filters,
    formatter=cfg.log.formatter,
    thread_id=cfg.log.thread_id,
)
print(doit(cfg.data, cfg.param))

from caikit.core import DataObjectBase, dataobject
from enum import Enum
from typing import List, Iterable, Tuple, Union
import struct


@dataobject
class DataTypes(Enum):
    FLOAT = 1
    INT = 2


@dataobject
class DataBuffer(DataObjectBase):
    data: bytes
    data_type: DataTypes

    @classmethod
    def from_iterable(
        cls, vals: Union[Iterable[int], Iterable[float]],
    ) -> "DataBuffer":
        data_type = (
            DataTypes.INT
            if vals and isinstance(vals[0], int)
            else DataTypes.FLOAT
        )
        data = struct.pack("!I" + "d" * len(vals), len(vals), *vals)
        return cls(data, data_type)

    def to_type(self) -> Union[Tuple[float], Tuple[int]]:
        unpacked = struct.unpack_from(
            "!" + "d" * struct.unpack_from("!I", self.data)[0], self.data, 4
        )
        if self.data_type == DataTypes.INT:
            return tuple(int(element) for element in unpacked)
        return unpacked


@dataobject
class Embeddings(DataObjectBase):
    embeddings: List[DataBuffer]

    def to_types(self) -> List[Union[Tuple[float], Tuple[int]]]:
        return [element.to_type() for element in self.embeddings]


embedding_set = Embeddings([
    DataBuffer.from_iterable([1, 2, 3]),
    DataBuffer.from_iterable([4, 5, 6]),
    DataBuffer.from_iterable([0.7, 0.8, 0.9]),
])

print("## To Types")
print(embedding_set.to_types())
print()

print("## To Json")
print(embedding_set.to_json())
print()

print("## To Protobuf")
print(embedding_set.to_proto())
print()

from caikit.core.data_model import DataStream
import os

in_memory_stream = DataStream.from_iterable((1, 2, 3, 4))
print(list(in_memory_stream))
print(list(in_memory_stream))

io_bound_stream = DataStream(lambda: iter(os.listdir(".")))
print(list(io_bound_stream))
print(list(io_bound_stream))

from caikit.core import ModuleBase, ModuleConfig, ModuleLoader, ModuleSaver, module
from caikit.core.data_model import DataStream
from typing import List, Tuple, Union
import argparse
import re
import tempfile


@module(
    "0f148161-90fc-4275-b6b5-6cbdf9826af6",
    "NameFinder",
    "1.0.0",
)
class NameFinder(ModuleBase):

    def __init__(self, name_list: List[str]):
        self._name_list = set(name_list)

    def run(self, text: str) -> List[str]:
        return [
            name for name in self._name_list
            if name in self._preprocess(text)
        ]

    @classmethod
    def train(
        cls,
        training_data: DataStream[str],
        delimiters = (" ", ".", ",", ":", ";", "'", "\""),
        ngram_size: int = 2,
    ) -> "NameFinder":
        names = []
        expr = re.compile("[{}]".format("".join(delimiters)))
        for sample in training_data:
            sample = cls._preprocess(sample)
            token_boundaries = cls._get_token_boundaries(sample, expr)
            for ng_len in range(1, ngram_size + 1):
                for idx in range(len(token_boundaries)):
                    if idx + ng_len > len(token_boundaries):
                        break
                    tokens = token_boundaries[idx:idx + ng_len]
                    names.append(sample[tokens[0][0]: tokens[-1][1]])
        return cls(names)

    def save(self, model_path: str):
        with ModuleSaver(self, model_path=model_path) as module_saver:
            module_saver.update_config({
                "name_list": self._name_list,
            })

    @classmethod
    def load(cls, model_path: Union[str, ModuleConfig]) -> "NameFinder":
        loader = ModuleLoader(model_path)
        return cls(loader.config["name_list"])

    ## Impl ##

    @staticmethod
    def _preprocess(text: str) -> str:
        return text.lower()

    @staticmethod
    def _get_token_boundaries(
        sample: str, expr: re.Pattern,
    ) -> List[Tuple[int, int]]:
        token_boundaries = []
        start = 0
        for match in expr.finditer(sample):
            match_offsets = match.span()
            token_boundaries.append((start, match_offsets[0]))
            start = match_offsets[1]
        if start < len(sample):
            token_boundaries.append((start, len(sample)))
        return token_boundaries


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Caikit module example")
    parser.add_argument(
        "--training-sample", "-s", nargs="+", required=True, help="Names to train with",
    )
    parser.add_argument(
        "--test-sample", "-t", nargs="+", required=True, help="Sample text to run inference with",
    )
    parser.add_argument(
        "--ngram_size", "-n", type=int, default=None, help="N-Gram length for names",
    )
    parser.add_argument(
        "--delimiters", "-d", nargs="*", default=None, help="Tokenization delimiters for training samples",
    )
    args = parser.parse_args()

    # Train the model
    train_kwargs = {"training_data": DataStream.from_iterable(args.training_sample)}
    if args.delimiters is not None:
        train_kwargs["delimiters"] = args.delimiters
    if args.ngram_size is not None:
        train_kwargs["ngram_size"] = args.ngram_size
    trained_model = NameFinder.train(**train_kwargs)

    # Run the samples through the trained model
    print("## Trained Model ##")
    for sample in args.test_sample:
        res = trained_model.run(sample)
        print(f"Results for [{sample}]: {res}")
    print()

    # Save and reload the model
    with tempfile.TemporaryDirectory() as workdir:
        trained_model.save(workdir)
        loaded_model = NameFinder.load(workdir)

    # Rerun with the loaded model
    print("## Loaded Model ##")
    for sample in args.test_sample:
        res = trained_model.run(sample)
        print(f"Results for [{sample}]: {res}")
    print()

from caikit.core import ModuleBase, TaskBase, module, task
from caikit.core.data_model import DataObjectBase, DataStream, dataobject
from caikit.core.exceptions import error_handler
from typing import Dict, Iterable
import alog

log = alog.use_channel("DEMO")
error = error_handler.get(log)

@dataobject
class Document(DataObjectBase):
    data: str


@dataobject
class TranslationRequestChunk(DataObjectBase):
    text: str
    target_language: str


@task(
    unary_parameters={"text": str, "target_language": str},
    unary_output_type=Document,
    streaming_parameters={"chunks": Iterable[TranslationRequestChunk]},
    streaming_output_type=Iterable[str],
)
class LanguageTranslationTask(TaskBase):
    """This task translates from one language to another target language"""


@module(
    "59c226a0-aabe-478d-841a-3bd6a030a897",
    "Sample LT Module",
    "1.0.0",
    task=LanguageTranslationTask,
)
class SampleLT(ModuleBase):

    def __init__(self, word_map: Dict[str, Dict[str, str]]):
        self._word_map = word_map

    @LanguageTranslationTask.taskmethod(input_streaming=True, output_streaming=True)
    def run_bidi_stream(
        self, chunks: DataStream[TranslationRequestChunk],
    ) -> DataStream[str]:

        def translation_generator():
            for chunk in chunks:
                text = chunk.text
                target_language = chunk.target_language
                error.value_check(
                    "<DMO15854362E>",
                    target_language in self._word_map,
                    "Unsupported target language: %s",
                    target_language,
                )
                target_word_map = self._word_map[target_language]
                for word in text.split():
                    yield target_word_map.get(word, f"UNK[{word}]")

        return DataStream(translation_generator)

    @LanguageTranslationTask.taskmethod()
    def run(self, text: str, target_language: str) -> Document:
        return Document(
            " ".join(
                self.run_bidi_stream([TranslationRequestChunk(text, target_language)])
            )
        )


if __name__ == "__main__":
    model = SampleLT({"foo": {"foo": "bar", "baz": "bat"}})
    print(model.run("foo is a bar", "foo"))

    input_stream = DataStream.from_iterable([
        TranslationRequestChunk("foo is", "foo"),
        TranslationRequestChunk("a bar", "foo"),
    ])
    for result in model.run_bidi_stream(input_stream):
        print(result)

import random
from caikit.core.augmentors import AugmentorBase
from caikit.core.data_model import DataStream


class DoublingAugmentor(AugmentorBase):

    augmentor_type = int

    def _augment(self, val: int) -> int:
        return val * 2


class RandomNoiseAugmentor(AugmentorBase):

    augmentor_type = float

    def __init__(self, random_seed: int = 42, range_size: float = 1.0):
        self._range_size = range_size
        super().__init__(random_seed, False)

    def _augment(self, val: float) -> float:
        delta = (random.random() - 0.5) * self._range_size
        return val + delta


def int_to_float(val: int) -> float:
    return float(val)


# Create a data stream that has been augmented with the doubled sequence and the
# randomly permuted sequence
stream = DataStream.from_iterable(range(10)).augment(
    DoublingAugmentor(42, False), 1, post_augment_func=int_to_float,
).augment(
    RandomNoiseAugmentor(), 1,
)
print(list(stream))

from caikit.core import DataObjectBase, ModuleBase, TaskBase, dataobject, module, task
from caikit.core.modules import ModuleLoader, ModuleSaver
from caikit.interfaces.runtime.data_model import TrainingInfoRequest, TrainingStatus
from caikit.runtime import grpc_server
import caikit.config
import os


@dataobject
class Greeting(DataObjectBase):
    greeting: str


@task(unary_parameters={"name": str}, unary_output_type=Greeting)
class GreetingTask(TaskBase):
    pass


@module("greeter", "Sample Greeter", "0.0.0", task=GreetingTask)
class GreeterModule(ModuleBase):

    def __init__(self, greeting_template: str = "Hello {}"):
        self._greeting_template = greeting_template

    def run(self, name: str) -> Greeting:
        return Greeting(self._greeting_template.format(name))

    @classmethod
    def train(cls, greeting_prefix: str) -> "GreeterModule":
        return cls(f"{greeting_prefix} ")

    def save(self, model_path: str):
        with ModuleSaver(module=self, model_path=model_path) as saver:
            saver.update_config({"greeting_template": self._greeting_template})

    @classmethod
    def load(cls, model_path: str) -> "GreeterModule":
        return cls(ModuleLoader(model_path).config.greeting_template)


caikit.configure(
    config_dict={
        "runtime": {
            # Import modules from this script
            "library": "__main__",
            # Auto-load models found in the local "models" directory
            "local_models_dir": "models",
            "lazy_load_local_models": True,
            "training": {
                # Save trained models in the local "models" directory
                "output_dir": "models",
                # Don't save with the model ID for ease of auto-loading
                "save_with_id": False,
            }
        }
    }
)
os.makedirs("models", exist_ok=True)

with grpc_server.RuntimeGRPCServer() as server:
    # Set up service clients
    chan = server.make_local_channel()
    train_client = server.training_service.stub_class(chan)
    train_status_client = server.training_management_service.stub_class(chan)
    inference_client = server.inference_service.stub_class(chan)

    # Launch a training
    training_request = server.training_service.messages.GreetingTaskGreeterModuleTrainRequest(
        model_name="greeter",
        parameters=server.training_service.messages.GreetingTaskGreeterModuleTrainParameters(
            greeting_prefix="Greetings",
        )
    )
    training_handle = train_client.GreetingTaskGreeterModuleTrain(training_request)
    print(f"Started Training {training_handle.training_id} for model {training_handle.model_name}")

    # Wait until the training completes
    while True:
        training_status = train_status_client.GetTrainingStatus(
            TrainingInfoRequest(training_handle.training_id).to_proto()
        )
        if training_status.state == TrainingStatus.COMPLETED.value:
            print(f"Finished training {training_handle.training_id}")
            break

    # Make an inference request
    inference_request = server.inference_service.messages.GreetingTaskRequest(name="Gabe")
    greeting = inference_client.GreetingTaskPredict(
        inference_request, metadata=list({"mm-model-id": training_handle.model_name}.items())
    )
    print(f"Got greeting: {greeting.greeting}")

from caikit.core import DataObjectBase, ModuleBase, TaskBase, dataobject, module, task
from caikit.core.modules import ModuleLoader, ModuleSaver
from caikit.runtime import http_server
import caikit.config
import os
import requests


@dataobject
class Greeting(DataObjectBase):
    greeting: str


@task(unary_parameters={"name": str}, unary_output_type=Greeting)
class GreetingTask(TaskBase):
    pass


@module("greeter", "Sample Greeter", "0.0.0", task=GreetingTask)
class GreeterModule(ModuleBase):

    def __init__(self, greeting_template: str = "Hello {}"):
        self._greeting_template = greeting_template

    def run(self, name: str) -> Greeting:
        return Greeting(self._greeting_template.format(name))

    @classmethod
    def train(cls, greeting_prefix: str) -> "GreeterModule":
        return cls(f"{greeting_prefix} ")

    def save(self, model_path: str):
        with ModuleSaver(module=self, model_path=model_path) as saver:
            saver.update_config({"greeting_template": self._greeting_template})

    @classmethod
    def load(cls, model_path: str) -> "GreeterModule":
        return cls(ModuleLoader(model_path).config.greeting_template)


caikit.configure(
    config_dict={
        "runtime": {
            # Import modules from this script
            "library": "__main__",
            # Auto-load models found in the local "models" directory
            "local_models_dir": "models",
            "lazy_load_local_models": True,
            "lazy_load_poll_period_seconds": 1,
            "training": {
                # Save trained models in the local "models" directory
                "output_dir": "models",
                # Don't save with the model ID for ease of auto-loading
                "save_with_id": False,
            }
        }
    }
)
os.makedirs("models", exist_ok=True)

with http_server.RuntimeHTTPServer() as server:
    # Set up service clients
    base_url = f"http://localhost:{server.port}"
    train_url = f"{base_url}/api/v1/GreetingTaskGreeterModuleTrain"
    model_info_url = f"{base_url}/info/models"
    inference_url = f"{base_url}/api/v1/task/greeting"

    # Launch a training
    model_name = "greeter"
    training_handle = requests.post(
        train_url, json={"model_name": model_name, "parameters": {"greeting_prefix": "Heyo"}}
    ).json()
    training_id = training_handle["training_id"]
    print(f"Started Training {training_id} for model {model_name}")

    # Wait until the training completes
    # NOTE: Training management not yet available in HTTP server
    while True:
        model_info = requests.get(model_info_url).json()
        if any(model["name"] == model_name for model in model_info["models"]):
            print(f"Finished training {training_id}")
            break

    # Make an inference request
    greeting = requests.post(
        inference_url, json={"inputs": "Gabe", "model_id": model_name},
    ).json()["greeting"]
    print(f"Got greeting: {greeting}")