<commit_before>import os
import pickle
import logging
from typing import Optional, List
from collections import defaultdict
import numpy as np
import baseline
from baseline.utils import (
    exporter,
    unzip_files,
    find_model_basename,
    find_files_with_prefix,
    import_user_module,
    read_json,
    is_sequence,
    Offsets,
    topk,
    to_numpy,
    revlut,
    load_vectorizers,
    load_vocabs,
    lookup_sentence,
    normalize_backend,
)
from baseline.model import load_model_for


logger = logging.getLogger('baseline')
__all__ = []
export = exporter(__all__)


class Service(object):

    def __init__(self, vocabs=None, vectorizers=None, model=None, preproc='client'):
        self.vectorizers = vectorizers
        self.model = model
        self.vocabs = vocabs
        self.preproc = preproc

    def get_vocab(self, vocab_type='word'):
        return self.vocabs.get(vocab_type)

    def get_labels(self):
        return self.model.get_labels()

    @classmethod
    def signature_name(cls):
        raise Exception("Undefined signature name")

    @classmethod
    def task_name(cls):
        raise Exception("Undefined task name")

    def predict(self, tokens, **kwargs):
        pass

    def format_output(self, predicted, **kwargs):
        """Turn the results of self.model.predict into our output format

        :param predicted: The results from self.model.predict
        :returns: Formatted output, different for each task
        """

    def batch_input(self, tokens):
        """Turn the input into a consistent format.
        :param tokens: tokens in format List[str] or List[List[str]]
        :return: List[List[str]]
        """
        # If the input is List[str] wrap it in list to make a batch of size one.
        return (tokens,) if isinstance(tokens[0], str) else tokens

    def prepare_vectorizers(self, tokens_batch):
        """Batch the input tokens, and call reset and count method on each vectorizers to set up their mxlen.
           This method is mainly for reducing repeated code blocks.

        :param tokens_batch: input tokens in format or List[List[str]]
        """
        for vectorizer in self.vectorizers.values():
            vectorizer.reset()
            for tokens in tokens_batch:
                _ = vectorizer.count(tokens)

    def vectorize(self, tokens_batch):
        """Turn the input into that batch dict for prediction.

        :param tokens_batch: `List[List[str]]`: The input text batch.

        :returns: dict[str] -> np.ndarray: The vectorized batch.
        """
        examples = defaultdict(list)
        for i, tokens in enumerate(tokens_batch):
            for k, vectorizer in self.vectorizers.items():
                vec, length = vectorizer.run(tokens, self.vocabs[k])
                examples[k].append(vec)
                if length is not None:
                    lengths_key = '{}_lengths'.format(k)
                    examples[lengths_key].append(length)

        for k in self.vectorizers.keys():
            examples[k] = np.stack(examples[k])
            lengths_key = '{}_lengths'.format(k)
            if lengths_key in examples:
                examples[lengths_key] = np.stack(examples[lengths_key])
        return examples

    @classmethod
    def load(cls, bundle, **kwargs):
        """Load a model from a bundle.

        This can be either a local model or a remote, exported model.

        :returns a Service implementation
        """
        # can delegate
        if os.path.isdir(bundle):
            directory = bundle
        else:
            directory = unzip_files(bundle)

        model_basename = find_model_basename(directory)
        vocabs = load_vocabs(directory)

        be = normalize_backend(kwargs.get('backend', 'tf'))

        remote = kwargs.get("remote", None)
        name = kwargs.get("name", None)
        if remote:
            logging.debug("loading remote model")
            beam = int(kwargs.get('beam', 30))
            model, preproc = Service._create_remote_model(
                directory, be, remote, name, cls.task_name(), cls.signature_name(), beam,
                preproc=kwargs.get('preproc', 'client'),
                version=kwargs.get('version'),
                remote_type=kwargs.get('remote_type'),
            )
            vectorizers = load_vectorizers(directory)
            return cls(vocabs, vectorizers, model, preproc)

        # Currently nothing to do here
        # labels = read_json(os.path.join(directory, model_basename) + '.labels')

        import_user_module('baseline.{}.embeddings'.format(be))
        try:
            import_user_module('baseline.{}.{}'.format(be, cls.task_name()))
        except:
            pass
        model = load_model_for(cls.task_name(), model_basename, **kwargs)
        vectorizers = load_vectorizers(directory)
        return cls(vocabs, vectorizers, model, 'client')

    @staticmethod
    def _create_remote_model(directory, backend, remote, name, task_name, signature_name, beam, **kwargs):
        """Reads the necessary information from the remote bundle to instatiate
        a client for a remote model.

        :directory the location of the exported model bundle
        :remote a url endpoint to hit
        :name the model name, as defined in tf-serving's model.config
        :signature_name  the signature to use.
        :beam used for s2s and found in the kwargs. We default this and pass it in.

        :returns a RemoteModel
        """
        from baseline.remote import create_remote
        assets = read_json(os.path.join(directory, 'model.assets'))
        model_name = assets['metadata']['exported_model']
        preproc = assets['metadata'].get('preproc', kwargs.get('preproc', 'client'))
        labels = read_json(os.path.join(directory, model_name) + '.labels')
        lengths_key = assets.get('lengths_key', None)
        inputs = assets.get('inputs', [])
        return_labels = bool(assets['metadata']['return_labels'])
        version = kwargs.get('version')

        if backend not in {'tf'}:
            raise ValueError("only Tensorflow is currently supported for remote Services")
        import_user_module('baseline.{}.remote'.format(backend))
        exp_type = kwargs.get('remote_type')
        if exp_type is None:
            exp_type = 'http' if remote.startswith('http') else 'grpc'
            exp_type = '{}-preproc'.format(exp_type) if preproc == 'server' else exp_type
            exp_type = f'{exp_type}-{task_name}'
        model = create_remote(
            exp_type,
            remote=remote, name=name,
            signature=signature_name,
            labels=labels,
            lengths_key=lengths_key,
            inputs=inputs,
            beam=beam,
            return_labels=return_labels,
            version=version,
        )
        return model, preproc


@export
class ClassifierService(Service):
    def __init__(self, vocabs=None, vectorizers=None, model=None, preproc='client'):
        super(ClassifierService, self).__init__(vocabs, vectorizers, model, preproc)
        if hasattr(self.model, 'return_labels'):
            self.return_labels = self.model.return_labels
        else:
            self.return_labels = True  # keeping the default classifier behavior
        if not self.return_labels:
            self.label_vocab = {index: label for index, label in enumerate(self.get_labels())}

    @classmethod
    def task_name(cls):
        return 'classify'

    @classmethod
    def signature_name(cls):
        return 'predict_text'

    def predict(self, tokens, preproc=None, raw=False, dense=False):
        """Take tokens and apply the internal vocab and vectorizers.  The tokens should be either a batch of text
        single utterance of type ``list``
        """
        if preproc is not None:
            logger.warning("Warning: Passing `preproc` to `ClassifierService.predict` is deprecated.")
        if raw and not dense:
            logger.warning("Warning: `raw` parameter is deprecated pass `dense=True` to get back values as a single tensor")
            dense = True
        tokens_batch = self.batch_input(tokens)
        self.prepare_vectorizers(tokens_batch)
        if self.preproc == "client":
            examples = self.vectorize(tokens_batch)
        elif self.preproc == 'server':
            # TODO: here we allow vectorizers even for preproc=server to get `word_lengths`.
            # vectorizers should not be available when preproc=server.
            featurized_examples = self.vectorize(tokens_batch)
            examples = {
                        'tokens': np.array([" ".join(x) for x in tokens_batch]),
                        self.model.lengths_key: featurized_examples[self.model.lengths_key]
            }

        outcomes_list = self.model.predict(examples, dense=dense)
        return self.format_output(outcomes_list, dense=dense)

    def format_output(self, predicted, dense=False, **kwargs):
        if dense:
            return predicted
        results = []
        for outcomes in predicted:
            if self.return_labels:
                results += [list(map(lambda x: (x[0], x[1].item()), sorted(outcomes, key=lambda tup: tup[1], reverse=True)))]
            else:
                results += [list(map(lambda x: (self.label_vocab[x[0].item()], x[1].item()),
                                     sorted(outcomes, key=lambda tup: tup[1], reverse=True)))]
        return results

@export
class EmbeddingsService(Service):
    @classmethod
    def task_name(cls):
        return 'servable-embeddings'

    @classmethod
    def signature_name(cls):
        return 'embed_text'

    def predict(self, tokens, preproc=None):
        if preproc is not None:
            logger.warning("Warning: Passing `preproc` to `EmbeddingsService.predict` is deprecated.")
        tokens_batch = self.batch_input(tokens)
        self.prepare_vectorizers(tokens_batch)
        if self.preproc == 'client':
            examples = self.vectorize(tokens_batch)
        else:
            examples = {
                'tokens': np.array([" ".join(x) for x in tokens_batch]),
            }
        return self.format_output(self.model.predict(examples))

    def format_output(self, predicted):
        return predicted

    @classmethod
    def load(cls, bundle, **kwargs):
        import_user_module('create_servable_embeddings')
        return super().load(bundle, **kwargs)


@export
class TaggerService(Service):

    def __init__(self, vocabs=None, vectorizers=None, model=None, preproc='client'):<commit_msg>chore: load the servable embeddings from mead-hub<commit_after>import os
import pickle
import logging
from typing import Optional, List
from collections import defaultdict
import numpy as np
import baseline
from baseline.utils import (
    exporter,
    unzip_files,
    find_model_basename,
    find_files_with_prefix,
    import_user_module,
    read_json,
    is_sequence,
    Offsets,
    topk,
    to_numpy,
    revlut,
    load_vectorizers,
    load_vocabs,
    lookup_sentence,
    normalize_backend,
)
from baseline.model import load_model_for


logger = logging.getLogger('baseline')
__all__ = []
export = exporter(__all__)


class Service(object):

    def __init__(self, vocabs=None, vectorizers=None, model=None, preproc='client'):
        self.vectorizers = vectorizers
        self.model = model
        self.vocabs = vocabs
        self.preproc = preproc

    def get_vocab(self, vocab_type='word'):
        return self.vocabs.get(vocab_type)

    def get_labels(self):
        return self.model.get_labels()

    @classmethod
    def signature_name(cls):
        raise Exception("Undefined signature name")

    @classmethod
    def task_name(cls):
        raise Exception("Undefined task name")

    def predict(self, tokens, **kwargs):
        pass

    def format_output(self, predicted, **kwargs):
        """Turn the results of self.model.predict into our output format

        :param predicted: The results from self.model.predict
        :returns: Formatted output, different for each task
        """

    def batch_input(self, tokens):
        """Turn the input into a consistent format.
        :param tokens: tokens in format List[str] or List[List[str]]
        :return: List[List[str]]
        """
        # If the input is List[str] wrap it in list to make a batch of size one.
        return (tokens,) if isinstance(tokens[0], str) else tokens

    def prepare_vectorizers(self, tokens_batch):
        """Batch the input tokens, and call reset and count method on each vectorizers to set up their mxlen.
           This method is mainly for reducing repeated code blocks.

        :param tokens_batch: input tokens in format or List[List[str]]
        """
        for vectorizer in self.vectorizers.values():
            vectorizer.reset()
            for tokens in tokens_batch:
                _ = vectorizer.count(tokens)

    def vectorize(self, tokens_batch):
        """Turn the input into that batch dict for prediction.

        :param tokens_batch: `List[List[str]]`: The input text batch.

        :returns: dict[str] -> np.ndarray: The vectorized batch.
        """
        examples = defaultdict(list)
        for i, tokens in enumerate(tokens_batch):
            for k, vectorizer in self.vectorizers.items():
                vec, length = vectorizer.run(tokens, self.vocabs[k])
                examples[k].append(vec)
                if length is not None:
                    lengths_key = '{}_lengths'.format(k)
                    examples[lengths_key].append(length)

        for k in self.vectorizers.keys():
            examples[k] = np.stack(examples[k])
            lengths_key = '{}_lengths'.format(k)
            if lengths_key in examples:
                examples[lengths_key] = np.stack(examples[lengths_key])
        return examples

    @classmethod
    def load(cls, bundle, **kwargs):
        """Load a model from a bundle.

        This can be either a local model or a remote, exported model.

        :returns a Service implementation
        """
        # can delegate
        if os.path.isdir(bundle):
            directory = bundle
        else:
            directory = unzip_files(bundle)

        model_basename = find_model_basename(directory)
        vocabs = load_vocabs(directory)

        be = normalize_backend(kwargs.get('backend', 'tf'))

        remote = kwargs.get("remote", None)
        name = kwargs.get("name", None)
        if remote:
            logging.debug("loading remote model")
            beam = int(kwargs.get('beam', 30))
            model, preproc = Service._create_remote_model(
                directory, be, remote, name, cls.task_name(), cls.signature_name(), beam,
                preproc=kwargs.get('preproc', 'client'),
                version=kwargs.get('version'),
                remote_type=kwargs.get('remote_type'),
            )
            vectorizers = load_vectorizers(directory)
            return cls(vocabs, vectorizers, model, preproc)

        # Currently nothing to do here
        # labels = read_json(os.path.join(directory, model_basename) + '.labels')

        import_user_module('baseline.{}.embeddings'.format(be))
        try:
            import_user_module('baseline.{}.{}'.format(be, cls.task_name()))
        except:
            pass
        model = load_model_for(cls.task_name(), model_basename, **kwargs)
        vectorizers = load_vectorizers(directory)
        return cls(vocabs, vectorizers, model, 'client')

    @staticmethod
    def _create_remote_model(directory, backend, remote, name, task_name, signature_name, beam, **kwargs):
        """Reads the necessary information from the remote bundle to instatiate
        a client for a remote model.

        :directory the location of the exported model bundle
        :remote a url endpoint to hit
        :name the model name, as defined in tf-serving's model.config
        :signature_name  the signature to use.
        :beam used for s2s and found in the kwargs. We default this and pass it in.

        :returns a RemoteModel
        """
        from baseline.remote import create_remote
        assets = read_json(os.path.join(directory, 'model.assets'))
        model_name = assets['metadata']['exported_model']
        preproc = assets['metadata'].get('preproc', kwargs.get('preproc', 'client'))
        labels = read_json(os.path.join(directory, model_name) + '.labels')
        lengths_key = assets.get('lengths_key', None)
        inputs = assets.get('inputs', [])
        return_labels = bool(assets['metadata']['return_labels'])
        version = kwargs.get('version')

        if backend not in {'tf'}:
            raise ValueError("only Tensorflow is currently supported for remote Services")
        import_user_module('baseline.{}.remote'.format(backend))
        exp_type = kwargs.get('remote_type')
        if exp_type is None:
            exp_type = 'http' if remote.startswith('http') else 'grpc'
            exp_type = '{}-preproc'.format(exp_type) if preproc == 'server' else exp_type
            exp_type = f'{exp_type}-{task_name}'
        model = create_remote(
            exp_type,
            remote=remote, name=name,
            signature=signature_name,
            labels=labels,
            lengths_key=lengths_key,
            inputs=inputs,
            beam=beam,
            return_labels=return_labels,
            version=version,
        )
        return model, preproc


@export
class ClassifierService(Service):
    def __init__(self, vocabs=None, vectorizers=None, model=None, preproc='client'):
        super(ClassifierService, self).__init__(vocabs, vectorizers, model, preproc)
        if hasattr(self.model, 'return_labels'):
            self.return_labels = self.model.return_labels
        else:
            self.return_labels = True  # keeping the default classifier behavior
        if not self.return_labels:
            self.label_vocab = {index: label for index, label in enumerate(self.get_labels())}

    @classmethod
    def task_name(cls):
        return 'classify'

    @classmethod
    def signature_name(cls):
        return 'predict_text'

    def predict(self, tokens, preproc=None, raw=False, dense=False):
        """Take tokens and apply the internal vocab and vectorizers.  The tokens should be either a batch of text
        single utterance of type ``list``
        """
        if preproc is not None:
            logger.warning("Warning: Passing `preproc` to `ClassifierService.predict` is deprecated.")
        if raw and not dense:
            logger.warning("Warning: `raw` parameter is deprecated pass `dense=True` to get back values as a single tensor")
            dense = True
        tokens_batch = self.batch_input(tokens)
        self.prepare_vectorizers(tokens_batch)
        if self.preproc == "client":
            examples = self.vectorize(tokens_batch)
        elif self.preproc == 'server':
            # TODO: here we allow vectorizers even for preproc=server to get `word_lengths`.
            # vectorizers should not be available when preproc=server.
            featurized_examples = self.vectorize(tokens_batch)
            examples = {
                        'tokens': np.array([" ".join(x) for x in tokens_batch]),
                        self.model.lengths_key: featurized_examples[self.model.lengths_key]
            }

        outcomes_list = self.model.predict(examples, dense=dense)
        return self.format_output(outcomes_list, dense=dense)

    def format_output(self, predicted, dense=False, **kwargs):
        if dense:
            return predicted
        results = []
        for outcomes in predicted:
            if self.return_labels:
                results += [list(map(lambda x: (x[0], x[1].item()), sorted(outcomes, key=lambda tup: tup[1], reverse=True)))]
            else:
                results += [list(map(lambda x: (self.label_vocab[x[0].item()], x[1].item()),
                                     sorted(outcomes, key=lambda tup: tup[1], reverse=True)))]
        return results

@export
class EmbeddingsService(Service):
    @classmethod
    def task_name(cls):
        return 'servable-embeddings'

    @classmethod
    def signature_name(cls):
        return 'embed_text'

    def predict(self, tokens, preproc=None):
        if preproc is not None:
            logger.warning("Warning: Passing `preproc` to `EmbeddingsService.predict` is deprecated.")
        tokens_batch = self.batch_input(tokens)
        self.prepare_vectorizers(tokens_batch)
        if self.preproc == 'client':
            examples = self.vectorize(tokens_batch)
        else:
            examples = {
                'tokens': np.array([" ".join(x) for x in tokens_batch]),
            }
        return self.format_output(self.model.predict(examples))

    def format_output(self, predicted):
        return predicted

    @classmethod
    def load(cls, bundle, **kwargs):
        import_user_module('hub:v1:addons:create_servable_embeddings')
        return super().load(bundle, **kwargs)


@export
class TaggerService(Service):

    def __init__(self, vocabs=None, vectorizers=None, model=None, preproc='client'):
{'fraction_non_alphanumeric': 0.1008883848407523, 'fraction_numerical': 0.0009923447689254325, 'mean_word_length': 2.4773250082155767, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3057499', 'n_tokens_mistral': 5908, 'n_tokens_neox': 5766, 'n_words': 1599}