# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02_analyzers.ipynb (unless otherwise specified).

__all__ = ['list2chunks', 'ZeroShotClassifier', 'TextEncoder', 'TopicModel', 'DEFAULT_TOKEN_PATTERN']

# Cell

import math
import warnings
import numpy as np
import pandas as pd

def list2chunks(a, n):
    k, m = divmod(len(a), n)
    return (a[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n))

# Cell

class ZeroShotClassifier():
    """
    Interface to Zero Shot Topic Classifier
    """

    def __init__(self, model_name='facebook/bart-large-mnli', device=None):
        """
        ZeroShotClassifier constructor

        **Args:**
          - model_name(str): name of a BART NLI model
          - device(str): device to use (e.g., 'cuda', 'cpu')
        """
        if 'mnli' not in model_name and 'xnli' not in model_name:
            raise ValueError('ZeroShotClasifier requires an MNLI or XNLI model')
        try:
            import torch
        except ImportError:
            raise Exception('ZeroShotClassifier requires PyTorch to be installed.')
        self.torch_device = device
        if self.torch_device is None: self.torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
        from transformers import AutoModelForSequenceClassification, AutoTokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name).to(self.torch_device)


    def predict(self, docs, labels=[], include_labels=False, multilabel=True,
               max_length=512, batch_size=8, nli_template='This text is about {}.',  topic_strings=[]):
        """
        This method performs zero-shot text classification using Natural Language Inference (NLI).


        **Parameters**:
          - docs(list|str): text of document or list of texts
          - labels(list): a list of strings representing topics of your choice
                          Example:
                           labels=['political science', 'sports', 'science']
          - include_labels(bool): If True, will return topic labels along with topic probabilities
          - multilabel(bool): If True, labels are considered independent and multiple labels can predicted true for document and be close to 1.
                            If False, scores are normalized such that probabilities sum to 1.
          - max_length(int): truncate long documents to this many tokens
          - batch_size(int): batch_size to use. default:8
                           Increase this value to speed up predictions - especially
                           if len(topic_strings) is large.
          - nli_template(str): labels are inserted into this template for use as hypotheses in natural language inference
          - topic_strings(list): alias for labels parameter for backwards compatibility

        **Returns:**


          inferred probabilities or list of inferred probabilities if doc is list
        """

        # error checks
        is_str_input = False
        if not isinstance(docs, (list, np.ndarray)):
            docs = [docs]
            is_str_input = True
        if not isinstance(docs[0], str): raise ValueError('docs must be string or a list of strings representing document(s)')
        if len(labels) > 0 and len(topic_strings) > 0: raise ValueError('labels and topic_strings are mutually exclusive')
        if not labels and not topic_strings: raise ValueError('labels must be a list of strings')
        if topic_strings:
            labels = topic_strings


        # convert to sequences
        sequence_pairs = []
        for premise in docs:
            sequence_pairs.extend([[premise, nli_template.format(label)] for label in labels])
        if batch_size  > len(sequence_pairs): batch_size = len(sequence_pairs)
        if len(sequence_pairs) >= 100 and batch_size==8:
            warnings.warn('TIP: Try increasing batch_size to speedup ZeroShotClassifier predictions')
        num_chunks = math.ceil(len(sequence_pairs)/batch_size)
        sequence_chunks = list2chunks(sequence_pairs, n=num_chunks)

        # inference
        import torch
        with torch.no_grad():
            outputs = []
            for sequences in sequence_chunks:
                batch = self.tokenizer.batch_encode_plus(sequences, return_tensors='pt', max_length=max_length, truncation='only_first', padding=True).to(self.torch_device)
                logits = self.model(batch['input_ids'], attention_mask=batch['attention_mask'], return_dict=False)[0]
                outputs.extend(logits.cpu().detach().numpy())
        outputs = np.array(outputs)
        outputs = outputs.reshape((len(docs), len(labels), -1))

        # process outputs
        if multilabel:
            # softmax over the entailment vs. contradiction dim for each label independently
            entail_contr_logits = outputs[..., [0, -1]]
            scores = np.exp(entail_contr_logits) / np.exp(entail_contr_logits).sum(-1, keepdims=True)
            scores = scores[..., 1]
        else:
            # softmax the "entailment" logits over all candidate labels
            entail_logits = outputs[..., -1]
            scores = np.exp(entail_logits) / np.exp(entail_logits).sum(-1, keepdims=True)
        scores = scores.tolist()
        if include_labels:
            scores = [list(zip(labels, s)) for s in scores]
        if is_str_input: scores = scores[0]
        return scores

# Cell

#from sentence_transformers import SentenceTransformer, util

class TextEncoder():
    """
    Tiny wrapper to sentence-transformers
    """

    def __init__(self, model_name='stsb-roberta-large', device=None):
        """
        TextEmbedder constructor.

        **Args:**
          - model_name(str): name of fine-tuned model for embeddings
          - device(str): device to use (e.g., 'cuda', 'cpu')
        """
        try:
            from sentence_transformers import SentenceTransformer, util
        except ImportError:
            raise Exception('TextEncoder requires: pip install sentence-transformers')
        try:
            import torch
        except ImportError:
            raise Exception('PyTorch must be installed.')

        self.torch_device = device
        if self.torch_device is None: self.torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.model = SentenceTransformer(model_name)

    def encode(self, texts, batch_size=32, normalize=False, show_progress_bar=False):
        """Generate embedding for supplied text"""
        if isinstance(texts, str): texts = [texts]
        return self.model.encode(texts, batch_size=batch_size,
                                 show_progress_bar=show_progress_bar, normalize_embeddings=normalize,
                                 convert_to_tensor=False, device=self.torch_device )

# Cell

from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.decomposition import NMF, LatentDirichletAllocation
import math
import numpy as np
DEFAULT_TOKEN_PATTERN = (r"\b[a-zA-Z][a-zA-Z0-9]*(?:[_/&-][a-zA-Z0-9]+)+\b|"
                         r"\b\d*[a-zA-Z][a-zA-Z0-9][a-zA-Z0-9]+\b")

class TopicModel():


    def __init__(self,texts=None, n_topics=None, n_features=10000,
                 min_df=5, max_df=0.5,  stop_words='english',
                 model_type='lda',
                 lda_max_iter=5, lda_mode='online',
                 token_pattern=None, verbose=1,
                 hyperparam_kwargs=None
    ):
        """
        Fits a topic model to documents in <texts>.
        Example:
            tm = ktrain.text.get_topic_model(docs, n_topics=20,
                                            n_features=1000, min_df=2, max_df=0.95)
        Args:
            texts (list of str): list of texts
            n_topics (int): number of topics.
                            If None, n_topics = min{400, sqrt[# documents/2]})
            n_features (int):  maximum words to consider
            max_df (float): words in more than max_df proportion of docs discarded
            stop_words (str or list): either 'english' for built-in stop words or
                                      a list of stop words to ignore
            model_type(str): type of topic model to fit. One of {'lda', 'nmf'}.  Default:'lda'
            lda_max_iter (int): maximum iterations for 'lda'.  5 is default if using lda_mode='online'.
                                If lda_mode='batch', this should be increased (e.g., 1500).
                                Ignored if model_type != 'lda'
            lda_mode (str):  one of {'online', 'batch'}. Ignored if model_type !='lda'
            token_pattern(str): regex pattern to use to tokenize documents.
            verbose(bool): verbosity
        """
        self.verbose=verbose

        # estimate n_topics
        if n_topics is None:
            if texts is None:
                raise ValueError('If n_topics is None, texts must be supplied')
            estimated = max(1, int(math.floor(math.sqrt(len(texts) / 2))))
            n_topics = min(400, estimated)
            print('n_topics automatically set to %s' % (n_topics))

        # train model
        if texts is not None:
            (model, vectorizer) = self.train(texts, model_type=model_type,
                                             n_topics=n_topics, n_features=n_features,
                                             min_df = min_df, max_df = max_df,
                                             stop_words=stop_words,
                                             lda_max_iter=lda_max_iter, lda_mode=lda_mode,
                                             token_pattern=token_pattern,
                                             hyperparam_kwargs=hyperparam_kwargs)
        else:
            vectorizer = None
            model = None



        # save model and vectorizer and hyperparameter settings
        self.vectorizer = vectorizer
        self.model = model
        self.n_topics = n_topics
        self.n_features = n_features
        if verbose: print('done.')

        # these variables are set by self.build():
        self.topic_dict = None
        self.doc_topics = None
        self.bool_array = None

        self.scorer = None       # set by self.train_scorer()
        self.recommender = None  # set by self.train_recommender()
        return


    def train(self,texts, model_type='lda', n_topics=None, n_features=10000,
              min_df=5, max_df=0.5,  stop_words='english',
              lda_max_iter=5, lda_mode='online',
              token_pattern=None, hyperparam_kwargs=None):
        """
        Fits a topic model to documents in <texts>.
        """
        if hyperparam_kwargs is None:
            hyperparam_kwargs = {}
        alpha = hyperparam_kwargs.get('alpha', 5.0 / n_topics)
        beta = hyperparam_kwargs.get('beta', 0.01)
        nmf_alpha = hyperparam_kwargs.get('nmf_alpha', 0)
        l1_ratio = hyperparam_kwargs.get('l1_ratio', 0)
        ngram_range = hyperparam_kwargs.get('ngram_range', (1,1))

         # preprocess texts
        if self.verbose: print('preprocessing texts...')
        if token_pattern is None: token_pattern = DEFAULT_TOKEN_PATTERN
        #if token_pattern is None: token_pattern = r'(?u)\b\w\w+\b'
        vectorizer = CountVectorizer(max_df=max_df, min_df=min_df,
                                 max_features=n_features, stop_words=stop_words,
                                 token_pattern=token_pattern, ngram_range=ngram_range)


        x_train = vectorizer.fit_transform(texts)

        # fit model

        if self.verbose: print('fitting model...')
        if model_type == 'lda':
            model = LatentDirichletAllocation(n_components=n_topics, max_iter=lda_max_iter,
                                              learning_method=lda_mode, learning_offset=50.,
                                              doc_topic_prior=alpha,
                                              topic_word_prior=beta,
                                              verbose=self.verbose, random_state=0)
        elif model_type == 'nmf':
            model = NMF(
                n_components=n_topics,
                max_iter=lda_max_iter,
                verbose=self.verbose,
                alpha=nmf_alpha,
                l1_ratio=l1_ratio,
                random_state=0)
        else:
            raise ValueError("unknown model type:", str(model_type))
        model.fit(x_train)

        # save model and vectorizer and hyperparameter settings
        return (model, vectorizer)


    @property
    def topics(self):
        """
        convenience method/property
        """
        return self.get_topics()


    def get_document_topic_distribution(self):
        """
        Gets the document-topic distribution.
        """
        self._check_build()
        return self.doc_topics


    def get_sorted_docs(self, topic_id):
        """
        Returns all docs sorted by relevance to <topic_id>.
        """
        docs = self.get_docs()
        d = {}
        for doc in docs: d[doc['doc_id']] = doc
        m = self.get_document_topic_distribution()
        doc_ids = (-m[:,topic_id]).argsort()
        return [d[doc_id] for doc_id in doc_ids]



    def get_word_weights(self, topic_id, n_words=100):
        """
        Returns a list tuples of the form: (word, weight) for given topic_id.
        """
        self._check_model()
        if topic_id+1 > len(self.model.components_):
            raise ValueError('topic_id must be less than %s' % (len(self.model.components_)))
        feature_names = self.vectorizer.get_feature_names()
        word_probs = self.model.components_[topic_id]
        word_ids = [i for i in word_probs.argsort()[:-n_words - 1:-1]]
        words = [feature_names[i] for i in word_ids]
        probs = [word_probs[i] for i in word_ids]
        return list( zip(words, probs) )


    def get_topics(self, n_words=10, as_string=True):
        """
        Returns a list of discovered topics
        """
        self._check_model()
        feature_names = self.vectorizer.get_feature_names()
        topic_summaries = []
        for topic_idx, topic in enumerate(self.model.components_):
            summary = [feature_names[i] for i in topic.argsort()[:-n_words - 1:-1]]
            if as_string: summary = " ".join(summary)
            topic_summaries.append(summary)
        return topic_summaries


    def print_topics(self, n_words=10, show_counts=False):
        """
        print topics
        """
        topics = self.get_topics(n_words=n_words, as_string=True)
        if show_counts:
            self._check_build()
            topic_counts = sorted([ (k, topics[k], len(v)) for k,v in self.topic_dict.items()],
                                    key=lambda kv:kv[-1], reverse=True)
            for (idx, topic, count) in topic_counts:
                print("topic:%s | count:%s | %s" %(idx, count, topic))
        else:
            for i, t in enumerate(topics):
                print('topic %s | %s' % (i, t))
        return


    def build(self, texts):
        """
        Builds the document-topic distribution showing the topic probability distirbution
        """
        doc_topics = self.predict(texts)
        self.doc_topics = doc_topics

        self.topic_dict = self._rank_documents(texts, doc_topics=doc_topics)
        return



    def get_docs(self, topic_ids=[], doc_ids=[], rank=False):
        """
        Returns document entries for supplied topic_ids.
        """
        self._check_build()
        if not topic_ids:
            topic_ids = list(range(self.n_topics))
        result_texts = []
        for topic_id in topic_ids:
            if topic_id not in self.topic_dict: continue
            texts = [{'text':tup[0], 'doc_id':tup[1], 'topic_proba':tup[2], 'topic_id':topic_id} for tup in self.topic_dict[topic_id]
                                                                                                     if not doc_ids or tup[1] in doc_ids]
            result_texts.extend(texts)
        if not rank:
            result_texts = sorted(result_texts, key=lambda x:x['doc_id'])
        return result_texts


    def get_doctopics(self,  topic_ids=[], doc_ids=[]):
        """
        Returns a topic probability distribution for documents
        """
        docs = self.get_docs(topic_ids=topic_ids, doc_ids=doc_ids)
        return np.array([self.doc_topics[idx] for idx in [x['doc_id'] for x in docs]])




    def predict(self, texts, threshold=None, harden=False):
        """
        predict topics for doucments
        """
        self._check_model()
        transformed_texts = self.vectorizer.transform(texts)
        X_topics = self.model.transform(transformed_texts)
        return X_topics

    def _rank_documents(self,
                       texts,
                       doc_topics=None):
        """
        Rank documents by topic score.
        """
        if doc_topics is not None:
            X_topics = doc_topics
        else:
            if self.verbose: print('transforming texts to topic space...')
            X_topics = self.predict(texts)
        topics = np.argmax(X_topics, axis=1)
        scores = np.amax(X_topics, axis=1)
        doc_ids = np.array([i for i, x in enumerate(texts)])
        result = list(zip(texts, doc_ids, topics, scores))
        if self.verbose: print('done.')
        result = sorted(result, key=lambda x: x[-1], reverse=True)
        result_dict = {}
        for r in result:
            text = r[0]
            doc_id = r[1]
            topic_id = r[2]
            score = r[3]
            lst = result_dict.get(topic_id, [])
            lst.append((text, doc_id, score))
            result_dict[topic_id] = lst
        return result_dict


    def _check_build(self):
        self._check_model()
        if self.topic_dict is None:
            raise Exception('Must call build() method.')


    def _check_model(self):
        if self.model is None or self.vectorizer is None:
            raise Exception('Must call train()')
