nltkor 1.2.14__cp311-cp311-macosx_13_0_x86_64.whl

nltkor/metrics/mauve.py

@@ -0,0 +1,273 @@
+# Author: Krishna Pillutla
+# License: GPLv3
+
+import math
+import numpy as np
+import time
+from types import SimpleNamespace
+from nltkor.make_requirement import make_requirement
+
+import faiss
+from sklearn.preprocessing import normalize
+from sklearn.decomposition import PCA
+from sklearn.metrics import auc as compute_area_under_curve
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+try:
+    import torch
+    FOUND_TORCH = True
+except (ImportError, ModuleNotFoundError):
+    FOUND_TORCH = False
+
+try:
+    import transformers
+    FOUND_TRANSFORMERS = True
+except (ImportError, ModuleNotFoundError):
+    FOUND_TRANSFORMERS = False
+
+if FOUND_TORCH and FOUND_TRANSFORMERS:
+    # only needed for tokenizing
+    from .mauve_utils import get_tokenizer, get_model, featurize_tokens_from_model, get_device_from_arg
+
+
+MODEL, TOKENIZER, MODEL_NAME = None, None, None
+
+class Mauve:
+    def __init__(self, model_name_or_path='skt/kobert-base-v1'):
+        self.featurize_model_name=model_name_or_path
+
+    def compute(self,
+            p_features=None, q_features=None,
+            p_tokens=None, q_tokens=None,
+            p_text=None, q_text=None,
+            num_buckets='auto', pca_max_data=-1, kmeans_explained_var=0.9,
+            kmeans_num_redo=5, kmeans_max_iter=500,
+            device_id=-1, max_text_length=1024,
+            divergence_curve_discretization_size=25, mauve_scaling_factor=5,
+            verbose=False, seed=25, batch_size=1, use_float64=False,
+    ):
+
+        """
+        Compute the MAUVE score between two text generations P and Q.
+
+        P is either specified as ``p_features``, ``p_tokens``, or ``p_text``. Same with Q.
+
+        :param ``p_features``: ``numpy.ndarray`` of shape (n, d), where n is the number of generations.
+        :param ``q_features``: ``numpy.ndarray`` of shape (n, d), where n is the number of generations.
+        :param ``p_tokens``: list of length n, each entry is torch.LongTensor of shape (1, length).
+        :param ``q_tokens``: list of length n, each entry is torch.LongTensor of shape (1, length).
+        :param ``p_text``: list of length n, each entry is a string.
+        :param ``q_text``: list of length n, each entry is a string.
+        :param ``num_buckets``: the size of the histogram to quantize P and Q. Options: ``'auto'`` (default, which is n/10) or an integer.
+        :param ``pca_max_data``: the number data points to use for PCA. If `-1`, use all the data. Default -1.
+        :param ``kmeans_explained_var``: amount of variance of the data to keep in dimensionality reduction by PCA. Default 0.9.
+        :param ``kmeans_num_redo``: number of times to redo k-means clustering (the best objective is kept). Default 5.
+            Try reducing this to 1 in order to reduce running time.
+        :param ``kmeans_max_iter``: maximum number of k-means iterations. Default 500.
+            Try reducing this to 100 in order to reduce running time.
+        :param ``featurize_model_name``: name of the model from which features are obtained. Default 'gpt2-large'.
+            We support all models which can be loaded from ``transformers.AutoModel.from_pretrained(featurize_model_name)``.
+        :param ``device_id``: Device for featurization. Supply gpu_id (e.g. 0 or 3) to use GPU or -1 to use CPU.
+        :param ``max_text_length``: maximum number of tokens to consider. Default 1024.
+        :param ``divergence_curve_discretization_size``: Number of points to consider on the divergence curve. Default 25.
+            Larger values do not offer much of a difference.
+        :param ``mauve_scaling_factor``: The constant``c`` from the paper. Default 5.
+            See `Best Practices <index.html#best-practices-for-mauve>`_ for details.
+        :param ``verbose``: If True, print running time updates.
+        :param ``seed``: random seed to initialize k-means cluster assignments.
+        :param ``batch_size``: Batch size for feature extraction.
+            A larger batch size speeds up computation.
+            You might have to experiment to find the largest batch size that fits in your GPU memory.
+            See `here <https://github.com/krishnap25/mauve/issues/8#issuecomment-1082075240>`_ for details.
+
+        :return: an object with fields p_hist, q_hist, divergence_curve and mauve.
+
+        * ``out.mauve`` is a number between 0 and 1, the MAUVE score. Higher values means P is closer to Q.
+        * ``out.frontier_integral``, a number between 0 and 1. Lower values mean that P is closer to Q.
+        * ``out.p_hist`` is the obtained histogram for P. Same for ``out.q_hist``.
+        * ``out.divergence_curve`` contains the points in the divergence curve. It is of shape (m, 2), where m is ``divergence_curve_discretization_size``
+
+        """
+
+        if p_features is None and p_tokens is None and p_text is None:
+            raise ValueError('Supply at least one of p_features, p_tokens, p_text')
+        if q_features is None and q_tokens is None and q_text is None:
+            raise ValueError('Supply at least one of q_features, q_tokens, q_text')
+        p_features = self.get_features_from_input(
+            p_features, p_tokens, p_text, self.featurize_model_name, max_text_length,
+            device_id, name="p", verbose=verbose, batch_size=batch_size, use_float64=use_float64,
+        )
+        q_features = self.get_features_from_input(
+            q_features, q_tokens, q_text, self.featurize_model_name, max_text_length,
+            device_id, name="q", verbose=verbose, batch_size=batch_size, use_float64=use_float64,
+        )
+        if num_buckets == 'auto':
+            # heuristic: use num_clusters = num_generations / 10
+            num_buckets = max(2, int(round(min(p_features.shape[0], q_features.shape[0]) / 10)))
+        elif not isinstance(num_buckets, int):
+            raise ValueError('num_buckets is expected to be an integer or "auto"')
+
+        # Acutal binning
+        t1 = time.time()
+        p, q = self.cluster_feats(p_features, q_features,
+                            num_clusters=num_buckets,
+                            norm='l2', whiten=False,
+                            pca_max_data=pca_max_data,
+                            explained_variance=kmeans_explained_var,
+                            num_redo=kmeans_num_redo,
+                            max_iter=kmeans_max_iter,
+                            seed=seed, verbose=verbose)
+        t2 = time.time()
+        if verbose:
+            print('total discretization time:', round(t2-t1, 2), 'seconds')
+
+        # Divergence curve and mauve
+        mixture_weights = np.linspace(1e-6, 1-1e-6, divergence_curve_discretization_size)
+        divergence_curve = self.get_divergence_curve_for_multinomials(p, q, mixture_weights, mauve_scaling_factor)
+        x, y = divergence_curve.T
+        idxs1 = np.argsort(x)
+        idxs2 = np.argsort(y)
+        mauve_score = 0.5 * (
+            compute_area_under_curve(x[idxs1], y[idxs1]) +
+            compute_area_under_curve(y[idxs2], x[idxs2])
+        )
+        fi_score = self.get_fronter_integral(p, q)
+        to_return = SimpleNamespace(
+            p_hist=p, q_hist=q, divergence_curve=divergence_curve,
+            mauve=mauve_score,
+            frontier_integral=fi_score,
+            num_buckets=num_buckets,
+        )
+        return to_return
+
+    def get_features_from_input(self, features, tokenized_texts, texts,
+                                featurize_model_name, max_len, device_id, name, batch_size,
+                                verbose=False, use_float64=False):
+        global MODEL, TOKENIZER, MODEL_NAME
+        if features is None:
+            # Featurizing is necessary. Make sure the required packages are available
+            if not FOUND_TORCH:
+                raise ModuleNotFoundError(
+                    """PyTorch not found. Please install PyTorch if you would like to use the featurization.
+                        For details, see `https://github.com/krishnap25/mauve`
+                        and `https://pytorch.org/get-started/locally/`.
+                    """)
+            if not FOUND_TRANSFORMERS:
+                raise ModuleNotFoundError(
+                    """Transformers not found. Please install Transformers if you would like to use the featurization.
+                        For details, see `https://github.com/krishnap25/mauve`
+                        and `https://huggingface.co/transformers/installation.html`.
+                    """)
+
+            if tokenized_texts is None:
+                # tokenize texts
+                if TOKENIZER is None or MODEL_NAME != featurize_model_name:
+                    if verbose: print('Loading tokenizer')
+                    TOKENIZER = get_tokenizer(featurize_model_name)
+                if verbose: print('Tokenizing text...')
+                tokenized_texts = [
+                    TOKENIZER.encode(sen, return_tensors='pt', truncation=True, max_length=max_len)
+                    for sen in texts
+                ]
+            # use tokenized_texts to featurize
+            if TOKENIZER is None or MODEL_NAME != featurize_model_name:
+                if verbose: print('Loading tokenizer')
+                TOKENIZER = get_tokenizer(featurize_model_name)
+            if MODEL is None or MODEL_NAME != featurize_model_name:
+                if verbose: print('Loading model')
+                MODEL = get_model(featurize_model_name, TOKENIZER, device_id)
+                MODEL_NAME = featurize_model_name
+            else:
+                MODEL = MODEL.to(get_device_from_arg(device_id))
+            if use_float64:
+                MODEL = MODEL.double()
+            if verbose: print('Featurizing tokens')
+            features = featurize_tokens_from_model(MODEL, tokenized_texts, batch_size, name).detach().cpu().numpy()
+        else:
+            features = np.asarray(features)
+        return features
+
+    def cluster_feats(self, p, q, num_clusters,
+                    norm='none', whiten=True,
+                    pca_max_data=-1,
+                    explained_variance=0.9,
+                    num_redo=5, max_iter=500,
+                    seed=0, verbose=False):
+        assert 0 < explained_variance < 1
+        if verbose:
+            print(f'seed = {seed}')
+        assert norm in ['none', 'l2', 'l1', None]
+        data1 = np.vstack([q, p])
+        if norm in ['l2', 'l1']:
+            data1 = normalize(data1, norm=norm, axis=1)
+        pca = PCA(n_components=None, whiten=whiten, random_state=seed+1)
+        if pca_max_data < 0 or pca_max_data >= data1.shape[0]:
+            pca.fit(data1)
+        elif 0 < pca_max_data < data1.shape[0]:
+            rng = np.random.RandomState(seed+5)
+            idxs = rng.choice(data1.shape[0], size=pca_max_data, replace=False)
+            pca.fit(data1[idxs])
+        else:
+            raise ValueError(f'Invalid argument pca_max_data={pca_max_data} with {data1.shape[0]} datapoints')
+        s = np.cumsum(pca.explained_variance_ratio_)
+        idx = np.argmax(s >= explained_variance)  # last index to consider
+        if verbose:
+            print(f'performing clustering in lower dimension = {idx}')
+        data1 = pca.transform(data1)[:, :idx+1]
+        # Cluster
+        data1 = data1.astype(np.float32)
+        t1 = time.time()
+        kmeans = faiss.Kmeans(data1.shape[1], num_clusters, niter=max_iter,
+                            verbose=verbose, nredo=num_redo, update_index=True,
+                            seed=seed+2)
+        kmeans.train(data1)
+        _, labels = kmeans.index.search(data1, 1)
+        labels = labels.reshape(-1)
+        t2 = time.time()
+        if verbose:
+            print('kmeans time:', round(t2-t1, 2), 's')
+
+        q_labels = labels[:len(q)]
+        p_labels = labels[len(q):]
+
+        q_bins = np.histogram(q_labels, bins=num_clusters,
+                            range=[0, num_clusters], density=True)[0]
+        p_bins = np.histogram(p_labels, bins=num_clusters,
+                            range=[0, num_clusters], density=True)[0]
+        return p_bins / p_bins.sum(), q_bins / q_bins.sum()
+
+
+    def kl_multinomial(self, p, q):
+        assert p.shape == q.shape
+        if np.logical_and(p != 0, q == 0).any():
+            return np.inf
+        else:
+            idxs = np.logical_and(p != 0, q != 0)
+            return np.sum(p[idxs] * np.log(p[idxs] / q[idxs]))
+
+
+    def get_divergence_curve_for_multinomials(self, p, q, mixture_weights, scaling_factor):
+        # TODO: check if extreme points are needed
+        divergence_curve = [[0, np.inf]] # extreme point
+        for w in np.sort(mixture_weights):
+            r = w * p + (1 - w) * q
+            divergence_curve.append([self.kl_multinomial(q, r), self.kl_multinomial(p, r)])
+        divergence_curve.append([np.inf, 0]) # other extreme point
+        return np.exp(-scaling_factor * np.asarray(divergence_curve))
+
+    def get_fronter_integral(self, p, q, scaling_factor=2):
+        total = 0.0
+        for p1, q1 in zip(p, q):
+            if p1 == 0 and q1 == 0:
+                pass
+            elif p1 == 0:
+                total += q1 / 4
+            elif q1 == 0:
+                total += p1 / 4
+            elif abs(p1 - q1) > 1e-8:
+                t1 = p1 + q1
+                t2 = p1 * q1 * (math.log(p1) - math.log(q1)) / (p1 - q1)
+                total += 0.25 * t1 - 0.5 * t2
+            # else: contribution is 0
+        return total * scaling_factor

nltkor/metrics/mauve_utils.py

@@ -0,0 +1,131 @@
+# Author: Krishna Pillutla
+# License: GPLv3
+import json
+import os
+import time
+from tqdm.auto import tqdm as tqdm_original
+
+import torch
+from transformers import AutoModel, AutoTokenizer, XLNetTokenizer
+# from nltkor.search.kobert_tokenizer import KoBERTTokenizer
+
+CPU_DEVICE = torch.device('cpu')
+tqdm = lambda *args, **kwargs: tqdm_original(
+    *args, **kwargs, disable=os.environ.get("DISABLE_TQDM", False))
+
+
+def get_device_from_arg(device_id):
+    if (device_id is not None and
+            torch.cuda.is_available() and
+            0 <= device_id < torch.cuda.device_count()):
+        return torch.device(f'cuda:{device_id}')
+    else:
+        return CPU_DEVICE
+
+def get_model(model_name, tokenizer, device_id):
+    device = get_device_from_arg(device_id)
+    if 'gpt2' in model_name or "bert" in model_name:
+        model = AutoModel.from_pretrained(model_name, pad_token_id=tokenizer.eos_token_id).to(device)
+        model = model.eval()
+    else:
+        raise ValueError(f'Unknown model: {model_name}')
+    return model
+
+def get_tokenizer(model_name='skt/kobert-base-v1'):
+    if 'gpt2' in model_name or "bert" in model_name:
+        if model_name == 'skt/kobert-base-v1':
+            # tokenizer = KoBERTTokenizer.from_pretrained(model_name)
+            tokenizer = XLNetTokenizer.from_pretrained(model_name)
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+    else:
+        raise ValueError(f'Unknown model: {model_name}')
+    return tokenizer
+
+def load_json_dataset(data_path, max_num_data):
+    texts = []
+    for i, line in enumerate(open(data_path)):
+        if i >= max_num_data:
+            break
+        texts.append(json.loads(line)['text'])
+    return texts
+
+def load_and_tokenize_json_data(tokenizer, data_path, max_len=1024, max_num_data=float('inf')):
+    """ Load and tokenize the data in a jsonl format
+
+    :param tokenizer:  HF tokenizer object
+    :param data_path: jsonl file to read. Read the "text" field of each line
+    :param max_len: maximum length of tokenized data
+    :param max_num_data: maximum number of lines to load
+    :return: list of `torch.LongTensor`s of shape (1, num_tokens), one for each input line
+    """
+    assert max_len <= 1024 and max_num_data >= 2000, f"max_len={max_len}, max_num_data={max_num_data} are insufficent"
+    t1 = time.time()
+    texts = load_json_dataset(data_path, max_num_data=max_num_data)
+    t2 = time.time()
+    print(f'dataset load time: {round(t2-t1, 2)} sec')
+    t1 = time.time()
+    tokenized_texts = [tokenizer.encode(sen, return_tensors='pt', truncation=True, max_length=max_len)
+                      for sen in texts]
+    t2 = time.time()
+    print(f'tokenizing time: {round(t2-t1, 2)} sec')
+    return tokenized_texts
+
+def decode_samples_from_lst(tokenizer, tokenized_texts):
+    """ Decode from tokens to string
+
+    :param tokenizer: HF tokenizer
+    :param tokenized_texts: list of list of tokens
+    :return: decoded output as a list of strings of the same length as tokenized_text_list
+    """
+    t1 = time.time()
+    output = []
+    for l in tokenized_texts:
+        o = tokenizer.decode(torch.LongTensor(l), skip_special_tokens=True)
+        output.append(o)
+    t2 = time.time()
+    print(f'de-tokenizing time: {round(t2-t1, 2)}')
+    return output
+
+@torch.no_grad()
+def featurize_tokens_from_model(model, tokenized_texts, batch_size, name="", verbose=False):
+    """Featurize tokenized texts using models, support batchify
+    :param model: HF Transformers model
+    :param batch_size: Batch size used during forward pass
+    :param tokenized_texts: list of torch.LongTensor of shape (1, length)
+    :param verbose: If True, print status and time
+    :return:
+    """
+    device = next(model.parameters()).device
+    t1 = time.time()
+    feats, chunks, chunk_sent_lengths = [], [], []
+    chunk_idx = 0
+
+    while chunk_idx * batch_size < len(tokenized_texts):
+        _chunk = [_t.view(-1) for _t in tokenized_texts[chunk_idx * batch_size: (chunk_idx + 1) * batch_size]]
+        chunks.append(_chunk)
+        chunk_sent_lengths.append([len(_c) for _c in _chunk])
+        chunk_idx += 1
+
+    for chunk, chunk_sent_length in tqdm(list(zip(chunks, chunk_sent_lengths)), desc=f"Featurizing {name}"):
+        padded_chunk = torch.nn.utils.rnn.pad_sequence(chunk,
+                                                       batch_first=True,
+                                                       padding_value=0).to(device)
+        attention_mask = torch.nn.utils.rnn.pad_sequence(
+            [torch.ones(sent_length).long() for sent_length in chunk_sent_length],
+            batch_first=True,
+            padding_value=0).to(device)
+        outs = model(input_ids=padded_chunk,
+                     attention_mask=attention_mask,
+                     past_key_values=None,
+                     output_hidden_states=True,
+                     return_dict=True)
+        h = []
+        for hidden_state, sent_length in zip(outs.hidden_states[-1], chunk_sent_length):
+            h.append(hidden_state[sent_length - 1])
+        h = torch.stack(h, dim=0)
+        feats.append(h.cpu())
+    t2 = time.time()
+    if verbose:
+        print(f'Featurize time: {round(t2-t1, 2)}')
+    return torch.cat(feats)

nltkor/misc/__init__.py

@@ -0,0 +1,11 @@
+# Natural Language Toolkit: Miscellaneous modules
+#
+# Copyright (C) 2001-2020 NLTK Project
+# Author: Steven Bird <stevenbird1@gmail.com>
+# URL: <http://nltk.org/>
+# For license information, see LICENSE.TXT
+
+from nltk.misc.chomsky import generate_chomsky
+from nltk.misc.wordfinder import word_finder
+from nltk.misc.minimalset import MinimalSet
+from nltk.misc.babelfish import babelize_shell

nltkor/misc/string2string_basic_functions.py

@@ -0,0 +1,59 @@
+"""
+string2string code
+src = https://github.com/stanfordnlp/string2string
+
+
+MIT License
+
+Copyright (c) 2023 Mirac Suzgun
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+
+"""
+
+from typing import List, Union
+
+# Take the Cartesian product of two lists of strings (or lists of lists of strings)
+def cartesian_product(
+    lst1: Union[List[str], List[List[str]]],
+    lst2: Union[List[str], List[List[str]]],
+    boolList: bool = False,
+    list_of_list_separator: str = " ## ",
+) -> Union[List[str], List[List[str]]]:
+    """
+    This function returns the Cartesian product of two lists of strings (or lists of lists of strings).
+
+    Arguments:
+        lst1: The first list of strings (or lists of lists of strings).
+        lst2: The second list of strings (or lists of lists of strings).
+        boolList: A boolean flag indicating whether the output should be a list of strings (or lists of lists of strings) (default: False).
+
+    Returns:
+        The Cartesian product of the two lists of strings (or lists of lists of strings).
+    """
+    if lst1 == []:
+        return lst2
+    elif lst2 == []:
+        return lst1
+    return [
+        s1 + ("" if not (boolList) else list_of_list_separator) + s2
+        for s1 in lst1
+        for s2 in lst2
+    ]

nltkor/misc/string2string_default_tokenizer.py

@@ -0,0 +1,83 @@
+"""
+string2string code
+src = https://github.com/stanfordnlp/string2string
+
+
+MIT License
+
+Copyright (c) 2023 Mirac Suzgun
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+
+"""
+
+"""
+This file contains the default tokenizer.
+"""
+
+from typing import List
+
+# Tokenizer class
+class Tokenizer:
+    """
+    This class contains the tokenizer.
+    """
+
+    def __init__(self,
+        word_delimiter: str = " ",
+    ):
+        """
+        Initializes the Tokenizer class.
+
+        Arguments:
+            word_delimiter (str): The word delimiter. Default is " ".
+        """
+        # Set the word delimiter
+        self.word_delimiter = word_delimiter
+
+    # Tokenize
+    def tokenize(self,
+        text: str,
+    ) -> List[str]:
+        """
+        Returns the tokens from a string.
+
+        Arguments:
+            text (str): The text to tokenize.
+
+        Returns:
+            List[str]: The tokens.
+        """
+        return text.split(self.word_delimiter)
+    
+    # Detokenize
+    def detokenize(self,
+        tokens: List[str],
+    ) -> str:
+        """
+        Returns the string from a list of tokens.
+
+        Arguments:
+            tokens (List[str]): The tokens.
+
+        Returns:
+            str: The string.
+        """
+        return self.word_delimiter.join(tokens)