textslinger 0.1.0__py3-none-any.whl

textslinger/causal.py

@@ -0,0 +1,437 @@
+import torch
+from typing import List, Tuple
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import itertools
+import heapq
+from textslinger.language_model import LanguageModel, DEFAULT_SYMBOL_SET
+from textslinger.exceptions import InvalidLanguageModelException
+from scipy.special import logsumexp
+from scipy.special import softmax
+import time
+from collections import defaultdict
+from typing import Final
+
+
+class CausalLanguageModel(LanguageModel):
+    """Character language model based on a pre-trained causal model, GPT-2 by default."""
+
+    def __init__(self,
+                 symbol_set: List[str],
+                 lang_model_name: str,
+                 lm_path: str = None,
+                 lm_device: str = "cpu",
+                 lm_left_context: str = "",
+                 beam_width: int = None,
+                 fp16: bool = False,
+                 mixed_case_context: bool = False,
+                 case_simple: bool = False,
+                 max_completed: int = None,
+                 ):
+        """
+        Initialize instance variables and load the language model with given path
+        Args:
+            symbol_set         - list of symbol strings
+            lang_model_name    - name of the Hugging Face casual language model to load
+            lm_path            - load fine-tuned model from specified directory
+            lm_device          - device to use for making predictions (cpu, mps, or cuda)
+            lm_left_context    - text to condition start of sentence on
+            beam_width         - how many hypotheses to keep during the search, None=off
+            fp16               - convert model to fp16 to save memory/compute on CUDA
+            mixed_case_context - use mixed case for language model left context
+            case_simple        - simple fixing of left context case
+            max_completed      - stop search once we reach this many completed hypotheses, None=don't prune
+        """
+        super().__init__(symbol_set=symbol_set)
+        self.model = None
+        self.tokenizer = None
+        self.vocab_size = 0
+        self.valid_vocab = []
+        self.vocab = defaultdict(list)
+        # Since subword token ids are integers, use a list instead of a dictionary
+        self.index_to_word = []
+        self.index_to_word_lower = []
+        self.symbol_set_lower = None
+        self.device = lm_device
+        self.left_context = lm_left_context
+        self.fp16 = fp16
+        self.mixed_case_context = mixed_case_context
+        self.case_simple = case_simple
+        self.max_completed = max_completed
+
+        if not max_completed and not beam_width:
+            print(f"WARNING: using causal language model without any pruning, this can be slow!")
+        else:
+            print(f"Causal language model, beam_width {beam_width}, max_completed {max_completed}")
+
+        # We optionally load the model from a local directory, but if this is not
+        # specified, we load a Hugging Face model
+        self.model_name = lang_model_name
+        self.model_dir = lm_path if lm_path else self.model_name
+
+        # Parameters for the search
+        self.beam_width = beam_width
+
+        # Simple heuristic to correct case in the LM context
+        self.simple_upper_words = {"i": "I",
+                                    "i'll": "I'll",
+                                    "i've": "I've",
+                                    "i'd": "I'd",
+                                    "i'm": "I'm"}
+
+        # Track how much time spent in different parts of the predict function
+        self.predict_total_ns = 0
+        self.predict_inference_ns = 0
+
+
+        # Are we a model that automatically inserts a start token that we need to get rid of
+        self.drop_first_token = False
+
+        self.load()
+
+    def _build_vocab(self) -> None:
+        """
+        Build a vocabulary table mapping token index to word strings
+        """
+
+        # Loop over all the subword tokens in the LLM
+        for i in range(self.vocab_size):
+            # Create a map from the subword token integer ID to the mixed and lowercase string versions
+            word = self.tokenizer.decode([i])
+            word_lower = word.lower()
+            self.index_to_word += word,
+            self.index_to_word_lower += word_lower,
+
+            # Check if all the characters in the subword token are in our valid symbol set
+            valid = True
+            for ch in word_lower:
+                if ch == ' ':
+                    continue
+                elif ch not in self.symbol_set_lower:
+                    valid = False
+                    break
+
+            # If the subword token symbols are all valid, then add it to the list of valid token IDs
+            if valid:
+                self.valid_vocab += i,
+                # Add this token ID to all lists for its valid text prefixes
+                for j in range(len(word)):
+                    key = word_lower[0:j + 1]
+                    self.vocab[key] += i,
+
+        # When done, self.vocab can be used to map to possible following subword tokens given some text, e.g.:
+        # self.vocab["cyclo"] = [47495, 49484]
+        # self.index_to_word[self.vocab["cyclo"][0]] = cyclop
+        # self.index_to_word[self.vocab["cyclo"][1]] = cyclopedia
+
+        (self.model_name.startswith("facebook/opt") or self.model_name.startswith("figmtu/opt") or "Llama-3.1" in self.model_name)
+
+        # Get the index we use for the start or end pseudo-word
+        if self.left_context == "":
+            if "gpt2" in self.model_name:
+                self.left_context = "<|endoftext|>"
+            elif "Llama-3.1" in self.model_name:
+                self.left_context = "<|begin_of_text|>"
+            # Seems to have both sentence start and end tokens: https://docs.mistral.ai/guides/tokenization/
+            elif "Mistral" in self.model_name:
+                self.left_context = "<s>"
+            else:
+                self.left_context = "</s>"
+
+        # OPT, Llama and Mistral all insert start token
+        self.drop_first_token = (self.model_name.startswith("facebook/opt") or
+                                 self.model_name.startswith("figmtu/opt") or
+                                 "Llama-3.1" in self.model_name or
+                                 "Mistral" in self.model_name)
+
+        # Get token id(s) for the left context we condition all sentences on
+        self.left_context_tokens = self._encode(self.left_context)
+        print(f"Causal: left_context = '{self.left_context}', left_context_tokens = {self.left_context_tokens}")
+
+    def _encode(self, text: str) -> List[int]:
+        tokens = self.tokenizer.encode(text)
+        # Both OPT and Llama automatically insert a start token which we want to control ourselves
+        if len(tokens) > 1 and self.drop_first_token:
+            tokens = tokens[1:]
+
+        return tokens
+
+    def _sequence_string(self, sequence: List[int]) -> str:
+        """
+        Convert a sequence of subword token IDs into a string with each token in ()'s
+        :param sequence: List of subword token IDs
+        :return: String
+        """
+        return "".join([f"({self.index_to_word[x]})" for x in sequence])
+
+    def get_all_tokens_text(self):
+        """
+        Return an array with the text of all subword tokens.
+        The array is in order by the integer index into the vocabulary.
+        This is mostly just for exploring the tokens in different LLMs.
+        :return: Array of subword token text strings.
+        """
+        result = []
+        for i in range(self.vocab_size):
+            result.append(self.tokenizer.decode([i]))
+        return result
+
+    def predict(self, evidence: List[str]) -> List[Tuple]:
+        """
+        Given an evidence of typed string, predict the probability distribution of
+        the next symbol
+        Args:
+            evidence - a list of characters (typed by the user)
+        Response:
+            A list of symbols with probability
+        """
+
+        assert self.model is not None, "language model does not exist!"
+        start_ns = time.time_ns()
+
+        context = "".join(evidence)
+
+        # If using the simple case feature, we need to go through the actual
+        # left context and capitalize the first letter in the sentence as
+        # well as any word in our list of words that should be capitalized.
+        if self.case_simple and len(context) > 0:
+            cased_context = ""
+            words = context.split()
+            for i, word in enumerate(words):
+                if i == 0 and word[0] >= 'a' and word[0] <= 'z':
+                    word = word[0].upper() + word[1:]
+                if i > 0:
+                    if word in self.simple_upper_words:
+                        word = self.simple_upper_words[word]
+                    cased_context += " "
+                cased_context += word
+            # Handle ending space in the context
+            if context[-1] == ' ':
+                cased_context += " "
+            context = cased_context
+
+        context_lower = context.lower()
+
+        # Index in the hypothesis string that is the next character after our context
+        target_pos = len(context_lower)
+
+        # For stats purposes track length of the prefix we are extending from space to match
+        # prefix_len = target_pos
+
+        # Look for the last space in the context, or -1 if no begin_text in context yet
+        pos = context_lower.rfind(" ")
+        tokens = []
+        tokens.extend(self.left_context_tokens)
+        if pos >= 0:
+            # Optionally, we condition on upper and lower case left context
+            if self.mixed_case_context:
+                truncated_context = context[0:pos]
+            else:
+                truncated_context = context_lower[0:pos]
+            tokens.extend(self._encode(truncated_context))
+            #prefix_len -= pos
+
+        #print(f"DEBUG, {context_lower} pos {pos}, prefix_len {prefix_len}")
+
+        # Constant indexes for use with the hypotheses tuples
+        LOGP: Final[int] = 0
+        SEQ: Final[int] = 1
+        LEN: Final[int] = 2
+
+        # Our starting hypothesis that we'll be extending.
+        # Format is (log likelihood, token id sequence, text length).
+        # Note: we only include tokens after any in left context.
+        start_length = 0
+        for x in tokens[len(self.left_context_tokens):]:
+            start_length += len(self.index_to_word_lower[x])
+        current_hypos = [(0.0, tokens, start_length)]
+
+        # We use a priority queue to track the top hypotheses during the beam search.
+        # For a beam of 8, empirical testing showed this was about the same amount
+        # of time as a simpler list that used a linear search to replace when full.
+        heapq.heapify(current_hypos)
+
+        # Create a hash mapping each valid following character to a list of log probabilities
+        char_to_log_probs = defaultdict(list)
+
+        # Add new extended hypotheses to this heap
+        next_hypos = []
+
+        # Tracks count of completed hypotheses
+        completed = 0
+
+        # Used to signal to while loop to stop the search
+        done = False
+
+        # Start a beam search forward from the backed off token sequence.
+        # Each iteration of this while loop extends hypotheses by all valid tokens.
+        # We only keep at most self.beam_width hypotheses in the valid heap.
+        # Stop extending search once we reach our max completed target.
+        while len(current_hypos) > 0 and not done:
+            # We'll explore hypothesis in order from most probable to least.
+            # This has little impact on how long it takes since this is only sorting a small number of things.
+            # But it is important with max_completed pruning since we want to bias for completing high probability things.
+            current_hypos.sort(reverse=True)
+
+            # Work on the hypotheses from the last round of extension.
+            # Create the torch tensor for the inference with a row for each hypothesis.
+            tokens_tensor = torch.tensor([x[SEQ] for x in current_hypos]).reshape(len(current_hypos), -1).to(self.device)
+
+            before_inference_ns = time.time_ns()
+            # Ask the LLM to predict tokens that come after our current set of hypotheses
+            with torch.no_grad():
+                # Compute the probabilities from the logits
+                log_probs = torch.log_softmax(self.model(tokens_tensor).logits[:, -1, :], dim=1)
+
+                # Create a big 2D tensor where each row is that hypothesis' current likelihood.
+                # First create a list of just the hypotheses' likelihoods.
+                # Then reshape to be a column vector.
+                # Then duplicate the column based on the number of subword tokens in the LLM.
+                add_tensor = torch.tensor([x[LOGP] for x in current_hypos]).reshape((log_probs.size()[0], 1)).repeat(1, log_probs.size()[1]).to(self.device)
+
+                # Add the current likelihoods with each subtoken's probability.
+                # Move it back to the CPU and convert to numpy since this makes it a lot faster to access for some reason.
+                new_log_probs = torch.add(log_probs, add_tensor).detach().cpu().numpy()
+            self.predict_inference_ns += time.time_ns() - before_inference_ns
+
+            for current_index, current in enumerate(current_hypos):
+                vocab = []
+                extra_vocab = []
+                # Extending this hypothesis must match the remaining text
+                remaining_context = context_lower[current[LEN]:]
+                if len(remaining_context) == 0:
+                    # There is no remaining context thus all subword tokens that are valid under our symbol set
+                    # should be considered when computing the probability of the next character.
+                    vocab = self.valid_vocab
+                else:
+                    if remaining_context in self.vocab:
+                        # We have a list of subword tokens that match the remaining text.
+                        # They could be the same length as the remaining text or longer and have the remaining text as a prefix.
+                        vocab = self.vocab[remaining_context]
+
+                    # We may need to use a subword token that doesn't completely consume the remaining text.
+                    # Find these by tokenizing all possible lengths of text starting from the current position.
+                    for i in range(1, len(remaining_context)):
+                        tokenization = self._encode(context_lower[current[LEN]:current[LEN] + i])
+                        # Ignore tokenizations involving multiple tokens since they involve an ID we would have already added.
+                        if len(tokenization) == 1:
+                            extra_vocab += tokenization[0],
+
+                # The below code takes the most time, results from pprofile on 5 phrases on an 2080 GPU:
+                #    299|  22484582|      89.5763|   3.9839e-06| 14.24%|                for token_id in itertools.chain(vocab, extra_vocab):
+                #    300|         0|            0|            0|  0.00%|                    # For a hypothesis to finish it must extend beyond the existing typed context
+                #    301|  22483271|      93.7939|  4.17172e-06| 14.91%|                    subword_len = len(self.index_to_word_lower[token_id])
+                #    302|  22483271|      92.8608|  4.13022e-06| 14.76%|                    if (current[LEN] + subword_len) > len(context):
+                #    303|         0|            0|            0|  0.00%|                        # Add this likelihood to the list for the character at the prediction position.
+                #    304|         0|            0|            0|  0.00%|                        # Tracking the list and doing logsumpexp later was faster than doing it for each add.
+                #    305|  22480431|      106.353|  4.73094e-06| 16.90%|                        char_to_log_probs[self.index_to_word_lower[token_id][target_pos - current[LEN]]] += new_log_probs[current_index][token_id],
+                #    306|  22480431|       92.689|   4.1231e-06| 14.73%|                        completed += 1
+                #    307|      2840|    0.0124488|  4.38338e-06|  0.00%|                    elif not self.beam_width or len(next_hypos) <
+                #
+                # Tuning notes:
+                #  - With a beam of 8 and max completed of 32,000, getting around 5x speedup on written dev set.
+                #  - This results in a PPL increase of 0.0025 versus old results using only beam of >= 8.
+                #  - Pruning based on log probability difference and based on minimum number of hypotheses per symbol in alphabet did worse.
+                #  - Code for these other pruning methods was removed.
+                # Possible ways to make it faster:
+                #  - Stop part way through the below for loop over (vocab, extra_vocab). But this seems weird since the token IDs are in
+                #    no particular order, we'd be just stopping early on the last hypothesis being explored by the enclosing loop.
+                #  - Sort the rows in the log prob results on the GPU. Use these to limit which token IDs we explore in the below
+                #    for loop. Is it possible to do this without introducing too much extra work to limit to the high probability ones?
+
+                # Create a list of token indexes that are a prefix of the target text.
+                # We go over all the integer IDs in the vocab and extra_vocab lists.
+                for token_id in itertools.chain(vocab, extra_vocab):
+                    # For a hypothesis to finish it must extend beyond the existing typed context
+                    subword_len = len(self.index_to_word_lower[token_id])
+                    if (current[LEN] + subword_len) > len(context):
+                        # Add this likelihood to the list for the character at the prediction position.
+                        # Tracking the list and doing logsumpexp later was faster than doing it for each add.
+                        char_to_log_probs[self.index_to_word_lower[token_id][target_pos - current[LEN]]] += new_log_probs[current_index][token_id],
+                        completed += 1
+                    elif not self.beam_width or len(next_hypos) < self.beam_width:
+                        # If we are under the beam limit then just add it
+                        heapq.heappush(next_hypos,
+                                       (new_log_probs[current_index][token_id],
+                                        current[SEQ] + [token_id],
+                                        current[LEN] + subword_len))
+                    elif new_log_probs[current_index][token_id] > next_hypos[0][LOGP]:
+                        # Or replace the worst hypotheses with the new one
+                        heapq.heappushpop(next_hypos,
+                                          (new_log_probs[current_index][token_id],
+                                           current[SEQ] + [token_id],
+                                           current[LEN] + subword_len))
+
+                # Break out of the for loop over hypotheses and while loop if we reach our max completed goal
+                if self.max_completed and completed >= self.max_completed:
+                    done = True
+                    break
+
+            # Swap in the extended set as the new current working set
+            current_hypos = next_hypos
+            next_hypos = []
+
+        # Parallel array to symbol_set for storing the marginals
+        char_probs = []
+        for ch in self.symbol_set_lower:
+            # Handle cases when symbols are never seen
+            if ch in char_to_log_probs:
+                char_probs += logsumexp(char_to_log_probs[ch]),
+            else:
+                char_probs += float("-inf"),
+
+        # Normalize to a distribution that sums to 1
+        char_probs = softmax(char_probs)
+
+        next_char_pred = {}
+        for i, ch in enumerate(self.symbol_set_lower):
+            # This works even if ch is a space
+            next_char_pred[ch.upper()] = char_probs[i]
+
+        end_ns = time.time_ns()
+        self.predict_total_ns += end_ns - start_ns
+
+        return list(sorted(next_char_pred.items(), key=lambda item: item[1], reverse=True))
+
+    def dump_predict_times(self) -> None:
+        """Print some stats about the prediction timing"""
+        if self.predict_total_ns > 0:
+            print(f"Predict %: "
+                  f"inference {self.predict_inference_ns / self.predict_total_ns * 100.0:.3f}")
+
+    def load(self) -> None:
+        """
+            Load the language model and tokenizer, initialize class variables
+        """
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=False)
+        except BaseException:
+            raise InvalidLanguageModelException(f"{self.model_name} is not a valid model identifier on HuggingFace.")
+        self.vocab_size = self.tokenizer.vocab_size
+        try:
+            self.model = AutoModelForCausalLM.from_pretrained(self.model_dir)
+            if self.fp16 and self.device == "cuda":
+                self.model = self.model.half()
+        except:
+            raise InvalidLanguageModelException(f"{self.model_dir} is not a valid local folder or model identifier on HuggingFace.")
+
+        self.model.eval()
+
+        self.model.to(self.device)
+
+        self.symbol_set_lower = []
+
+        for ch in self.symbol_set:
+            # This works even if ch is a space            
+            self.symbol_set_lower.append(ch.lower())
+
+        self._build_vocab()
+
+    def get_num_parameters(self) -> int:
+        """
+            Find out how many parameters the loaded model has
+        Args:
+        Response:
+            Integer number of parameters in the transformer model
+        """
+        return sum(p.numel() for p in self.model.parameters())
+    

textslinger/classifier.py

@@ -0,0 +1,227 @@
+from collections import Counter
+import torch
+from typing import List, Tuple
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from textslinger.language_model import LanguageModel
+from textslinger.exceptions import InvalidLanguageModelException, WordPredictionsNotSupportedException
+from scipy.special import softmax
+
+
+class ClassifierLanguageModel(LanguageModel):
+    """Character language model based on a transformer with a classification head"""
+
+    def __init__(self,
+                 symbol_set: List[str],
+                 lang_model_name: str,
+                 lm_path: str = None,
+                 lm_device: str = "cpu",
+                 lm_left_context: str = "",
+                 beam_width: int = 8,
+                 batch_size: int = 8,
+                 fp16: bool = False,
+                 mixed_case_context = False,
+                 case_simple = False,
+                 ):
+        """
+        Initialize instance variables and load the language model with given path
+        Args:
+            response_type      - SYMBOL only
+            symbol_set         - list of symbol strings
+            lang_model_name    - name of the Hugging Face casual language model to load
+            lm_path            - load fine-tuned model from specified directory
+            lm_device          - device to use for making predictions (cpu, mps, or cuda)
+            lm_left_context    - text to condition start of sentence on
+            beam_width         - how many hypotheses to keep during the search
+            batch_size         - how many sequences to pass in at a time during inference
+            fp16               - convert model to fp16 to save memory/compute on CUDA
+            mixed_case_context - use mixed case for language model left context
+            case_simple        - simple fixing of left context case
+        """
+        super().__init__(symbol_set=symbol_set)
+        self.model = None
+        self.tokenizer = None
+        self.vocab_size = 0
+        self.valid_vocab = []
+        self.vocab = {}
+        self.longest_token = 0
+        self.index_to_word = {}
+        self.index_to_word_lower = {}
+        self.symbol_set_lower = None
+        self.device = lm_device
+        self.left_context = lm_left_context
+        self.fp16 = fp16
+        self.mixed_case_context = mixed_case_context
+        self.case_simple = case_simple
+
+        # We optionally load the model from a local directory, but if this is not
+        # specified, we load a Hugging Face model
+        self.model_name = lang_model_name
+        self.model_dir = lm_path if lm_path else self.model_name
+
+        # parameters for search
+        self.beam_width = beam_width
+        self.batch_size = batch_size
+
+        self.simple_upper_words = {"i": "I",
+                                    "i'll": "I'll",
+                                    "i've": "I've",
+                                    "i'd": "I'd",
+                                    "i'm": "I'm"}
+        self.load()
+
+    def _build_vocab(self) -> None:
+        """
+        Build a vocabulary table mapping token index to word strings
+        """
+
+        for i in range(self.vocab_size):
+            word = self.tokenizer.decode([i])
+            word_lower = word.lower()
+            self.index_to_word[i] = word
+            self.index_to_word_lower[i] = word_lower
+            valid = True
+            for ch in word_lower:
+                if ch not in self.symbol_set_lower:
+                    valid = False
+                    break
+            if valid:
+                self.valid_vocab += i,
+                length = len(word)
+                if length > self.longest_token:
+                    self.longest_token = length
+                for j in range(length):
+                    key = word_lower[0:j + 1]
+                    if key not in self.vocab:
+                        self.vocab[key] = []
+                    self.vocab[key] += i,
+
+        # Get the index we use for the start or end pseudo-word
+        if self.left_context == "":
+            if "gpt2" in self.model_name:
+                self.left_context = "<|endoftext|>"
+            else:
+                self.left_context = "</s>"
+        # Get token id(s) for the left context we condition all sentences on
+        self.left_context_tokens = self._encode(self.left_context)
+        # print(f"left_context_tokens = {self.left_context_tokens}")
+
+    def _encode(self, text: str) -> List[int]:
+        if text == "":
+            return []
+
+        tokens = self.tokenizer.encode(text)
+        if len(tokens) > 1 and (self.model_name.startswith("facebook/opt") or self.model_name.startswith("figmtu/opt")):
+            tokens = tokens[1:]
+
+        return tokens
+
+    def predict_character(self, evidence: List[str]) -> List[Tuple]:
+        """
+        Given an evidence of typed string, predict the probability distribution of
+        the next symbol
+        Args:
+            evidence - a list of characters (typed by the user)
+        Response:
+            A list of symbols with probability
+        """
+
+        assert self.model is not None, "language model does not exist!"
+
+        context = "".join(evidence)
+
+        if self.case_simple and len(context) > 0:
+            cased_context = ""
+            words = context.split()
+            for i, word in enumerate(words):
+                if i == 0 and word[0] >= 'a' and word[0] <= 'z':
+                    word = word[0].upper() + word[1:]
+                if i > 0:
+                    if word in self.simple_upper_words:
+                        word = self.simple_upper_words[word]
+                    cased_context += " "
+                cased_context += word
+            # Handle ending space in the context
+            if context[-1] == ' ':
+                cased_context += " "
+            #print(f"Simple casing of left context, from '{context}' to '{cased_context}'")
+            context = cased_context
+
+        else:
+            context = context.lower()
+
+        tokens = []
+        tokens.extend(self.left_context_tokens)
+
+        # Optionally, we condition on upper and lower case left context
+        if not self.mixed_case_context:
+            context = context.lower()
+        tokens.extend(self._encode(context))
+
+        tensor = torch.tensor([tokens]).to(self.device)
+
+        with torch.no_grad():
+            logits = self.model(tensor).logits
+            char_probs = torch.softmax(logits, dim=1).to("cpu").numpy()[0]
+
+        keys = [*"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz' "]
+        next_char_pred = Counter(dict(zip(keys, char_probs)))
+
+        for low in self.symbol_set_lower:
+            if low.isalpha():
+                next_char_pred[low.upper()] += next_char_pred[low]
+                del next_char_pred[low]
+
+        return list(sorted(next_char_pred.items(), key=lambda item: item[1], reverse=True))
+
+    def predict_word(self, 
+                     left_context: List[str], 
+                     right_context: List[str] = [" "],
+                     nbest: int = 3,
+                     ) -> List[Tuple]:
+        """
+        Using the provided data, compute log likelihoods over the next sequence of symbols
+        Args:
+            left_context - The text that precedes the desired prediction.
+            right_context - The text that will follow the desired prediction. For simple word
+                predictions, this should be a single space.
+            nbest - The number of top predictions to return
+
+        Response:
+            A list of tuples, (predicted text, log probability)
+        """
+        raise WordPredictionsNotSupportedException("Word predictions are not supported for this model.")
+
+    def load(self) -> None:
+        """
+            Load the language model and tokenizer, initialize class variables
+        """
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=False)
+        except BaseException:
+            raise InvalidLanguageModelException(f"{self.model_name} is not a valid model identifier on HuggingFace.")
+        self.vocab_size = self.tokenizer.vocab_size
+        try:
+            self.model = AutoModelForSequenceClassification.from_pretrained(self.model_dir)
+            if self.fp16 and self.device == "cuda":
+                self.model = self.model.half()
+        except:
+            raise InvalidLanguageModelException(f"{self.model_dir} is not a valid local folder or model identifier on HuggingFace.")
+
+        self.model.eval()
+
+        self.model.to(self.device)
+
+        self.symbol_set_lower = []
+        for ch in self.symbol_set:
+            self.symbol_set_lower.append(ch.lower())
+
+        self._build_vocab()
+
+    def get_num_parameters(self) -> int:
+        """
+            Find out how many parameters the loaded model has
+        Args:
+        Response:
+            Integer number of parameters in the transformer model
+        """
+        return sum(p.numel() for p in self.model.parameters())

textslinger/exceptions.py

@@ -0,0 +1,25 @@
+
+class TextPredictException(Exception):
+    """textpredict core exception.
+
+    Thrown when an error occurs specific to textpredict core concepts.
+    """
+
+    def __init__(self, message, errors=None):
+        super().__init__(message)
+        self.message = message
+        self.errors = errors
+
+
+class InvalidLanguageModelException(TextPredictException):
+    """Invalid Language Model Exception.
+
+    Thrown when attempting to load a language model from an invalid path"""
+    ...
+
+
+class KenLMInstallationException(TextPredictException):
+    """KenLM Installation Exception.
+
+    Thrown when attempting to import kenlm without installing the module"""
+    ...

textslinger/language_model.py

@@ -0,0 +1,52 @@
+"""Defines the language model base class."""
+from abc import ABC, abstractmethod
+from typing import List, Tuple
+from string import ascii_uppercase
+
+# Returns all uppercase English letters
+def alphabet():
+    """Alphabet.
+
+    Function used to standardize the symbols we use as alphabet.
+
+    Returns
+    -------
+        array of letters.
+    """
+    return list(ascii_uppercase)
+
+DEFAULT_SYMBOL_SET = alphabet() + [' ']
+
+class LanguageModel(ABC):
+    """Parent class for Language Models."""
+
+    symbol_set: List[str] = None
+
+    def __init__(self,
+                 symbol_set: List[str] = DEFAULT_SYMBOL_SET):
+        self.symbol_set = symbol_set
+
+    @classmethod
+    def name(cls) -> str:
+        """Model name used for configuration"""
+        suffix = 'LanguageModel'
+        if cls.__name__.endswith(suffix):
+            return cls.__name__[0:-len(suffix)].upper()
+        return cls.__name__.upper()
+
+    @abstractmethod
+    def predict(self, evidence: List[str]) -> List[Tuple]:
+        """
+        Using the provided data, compute log likelihoods over the entire symbol set.
+        Args:
+            evidence - ['A', 'B']
+
+        Response:
+            probability - dependent on response type, a list of words or symbols with probability
+        """
+        ...
+
+    @abstractmethod
+    def load(self) -> None:
+        """Load model from the provided assets/path"""
+        ...

textslinger/mixture.py

@@ -0,0 +1,126 @@
+from collections import Counter
+from typing import Optional, Dict, List, Tuple
+from math import isclose
+
+from textslinger.language_model import LanguageModel
+from textslinger.exceptions import InvalidLanguageModelException
+
+# pylint: disable=unused-import
+# flake8: noqa
+"""All supported models must be imported"""
+from textslinger.causal import CausalLanguageModel
+from textslinger.ngram import NGramLanguageModel
+
+class MixtureLanguageModel(LanguageModel):
+    """
+        Character language model that mixes any combination of other models
+    """
+
+    supported_lm_types = ["CAUSAL", "NGRAM", "UNIFORM"]
+
+    @staticmethod
+    def language_models_by_name() -> Dict[str, LanguageModel]:
+        """Returns available language models indexed by name."""
+        return {lm.name(): lm for lm in LanguageModel.__subclasses__()}
+
+    @staticmethod
+    def validate_parameters(types: List[str], weights: List[float], params: List[Dict[str, str]]):
+        if params is not None:
+            if (types is None) or (len(types) != len(params)):
+                raise InvalidLanguageModelException("Length of parameters does not match length of types")
+
+        if weights is not None:
+            if (types is None) or (len(types) != len(weights)):
+                raise InvalidLanguageModelException("Length of weights does not match length of types")
+            if not isclose(sum(weights), 1.0, abs_tol=1e-05):
+                raise InvalidLanguageModelException("Weights do not sum to 1")
+
+        if types is not None:
+            if weights is None:
+                raise InvalidLanguageModelException("Model weights not provided")
+            if params is None:
+                raise InvalidLanguageModelException("Model parameters not provided")
+            if not all(x in MixtureLanguageModel.supported_lm_types for x in types):
+                raise InvalidLanguageModelException(f"Supported model types: {MixtureLanguageModel.supported_lm_types}")
+
+    def __init__(self,
+                 symbol_set: List[str],
+                 lm_types: List[str],
+                 lm_weights: List[float],
+                 lm_params: List[Dict[str, str]] = None):
+        """
+        Initialize instance variables and load the language model with given path
+        Args:
+            response_type - SYMBOL only
+            symbol_set - list of symbol strings
+            lm_types - list of types of models to mix
+            lm_weights - list of weights to use when mixing the models
+            lm_params - list of dictionaries to pass as parameters for each model's instantiation
+        """
+
+        MixtureLanguageModel.validate_parameters(lm_types, lm_weights, lm_params)
+
+        super().__init__(symbol_set=symbol_set)
+        self.models = list()
+        self.symbol_set = symbol_set
+
+        self.lm_types = lm_types
+        self.lm_weights = lm_weights
+        self.lm_params = lm_params
+
+        MixtureLanguageModel.validate_parameters(self.lm_types, self.lm_weights, self.lm_params)
+
+        self.load()
+
+    @staticmethod
+    def interpolate_language_models(lms: List[Dict[str, float]], coeffs: List[float]) -> List[Tuple]:
+        """
+        interpolate two or more language models
+        Args:
+            lms - output from the language models (a list of dicts with char as keys and prob as values)
+            coeffs - list of rescale coefficients, lms[0] will be scaled by coeffs[0] and so on
+        Response:
+            a list of (char, prob) tuples representing an interpolated language model
+        """
+        combined_lm = Counter()
+
+        for i, lm in enumerate(lms):
+            for char in lm:
+                combined_lm[char] += lm[char] * coeffs[i]
+
+        return list(sorted(combined_lm.items(), key=lambda item: item[1], reverse=True))
+
+    def predict(self, evidence: List[str]) -> List[Tuple]:
+        """
+        Given an evidence of typed string, predict the probability distribution of
+        the next symbol
+        Args:
+            evidence - a list of characters (typed by the user)
+        Response:
+            A list of symbols with probability
+        """
+
+        pred_list = list()
+
+        # Generate predictions from each component language model
+        pred_list = [dict(model.predict(evidence)) for model in self.models]
+
+        # Mix the component models
+        next_char_pred = MixtureLanguageModel.interpolate_language_models(pred_list, self.lm_weights)
+
+        return next_char_pred
+
+    def load(self) -> None:
+        """
+            Load the language models to be mixed
+        """
+        language_models = MixtureLanguageModel.language_models_by_name()
+        for lm_type, params in zip(self.lm_types, self.lm_params):
+            model = language_models[lm_type]
+            lm = None
+            try:
+                lm = model(self.symbol_set, **params)
+            except InvalidLanguageModelException as e:
+                raise InvalidLanguageModelException(f"Error in creation of model type {lm_type}: {e.message}")
+
+            self.models.append(lm)

textslinger/ngram.py

@@ -0,0 +1,110 @@
+from collections import Counter
+from typing import Optional, List, Tuple
+from textslinger.language_model import LanguageModel
+from textslinger.exceptions import InvalidLanguageModelException
+import kenlm
+import numpy as np
+
+
+class NGramLanguageModel(LanguageModel):
+    """Character n-gram language model using the KenLM library for querying"""
+
+    def __init__(self,
+                 symbol_set: List[str],
+                 lm_path: str,
+                 skip_symbol_norm: Optional[bool] = False):
+
+        super().__init__(symbol_set=symbol_set)
+        print(f"Creating n-gram language model, lm_path = {lm_path}")
+        self.model = None
+        self.lm_path = lm_path
+        self.skip_symbol_norm = skip_symbol_norm
+        self.load()
+
+    def predict(self, evidence: List[str]) -> List[Tuple]:
+        """
+        Given an evidence of typed string, predict the probability distribution of
+        the next symbol
+        Args:
+            evidence - a list of characters (typed by the user)
+        Response:
+            A list of symbols with probability
+        """
+
+        # Do not modify the original parameter, could affect mixture model
+        context = evidence.copy()
+
+        if len(context) > 11:
+            context = context[-11:]
+
+        evidence_str = ''.join(context).lower()
+
+        for i, ch in enumerate(context):
+            if ch == ' ':
+                context[i] = "<sp>"
+
+        self.model.BeginSentenceWrite(self.state)
+
+        # Update the state one token at a time based on evidence, alternate states
+        for i, token in enumerate(context):
+            if i % 2 == 0:
+                self.model.BaseScore(self.state, token.lower(), self.state2)
+            else:
+                self.model.BaseScore(self.state2, token.lower(), self.state)
+
+        next_char_pred = None
+
+        # Generate the probability distribution based on the final state
+        if len(context) % 2 == 0:
+            next_char_pred = self.prob_dist(self.state)
+        else:
+            next_char_pred = self.prob_dist(self.state2)
+
+        return next_char_pred
+
+    def load(self) -> None:
+        """
+            Load the language model, initialize state variables
+        Args:
+            path: language model file path
+        """
+
+        try:
+            self.model = kenlm.LanguageModel(self.lm_path)
+        except BaseException:
+            raise InvalidLanguageModelException(
+                f"A valid model path must be provided for the KenLMLanguageModel.\nPath{self.lm_path} is not valid.")
+
+        self.state = kenlm.State()
+        self.state2 = kenlm.State()
+
+    def prob_dist(self, state: kenlm.State) -> List[Tuple]:
+        """
+            Take in a state and generate the probability distribution of next character
+        Args:
+            state - the kenlm state updated with the evidence
+        Response:
+            A list of symbols with probability
+        """
+        next_char_pred = Counter()
+
+        temp_state = kenlm.State()
+
+        for char in self.symbol_set:
+            # Replace the space character with KenLM's <sp> token
+            if char == ' ':
+                score = self.model.BaseScore(state, '<sp>', temp_state)
+            else:
+                score = self.model.BaseScore(state, char.lower(), temp_state)
+
+            # BaseScore returns log probs, convert by putting 10 to its power
+            next_char_pred[char] = pow(10, score)
+
+        # We can optionally disable normalization over our symbol set
+        # This is useful if we want to compare against SRILM with a LM with a larger vocab
+        if not self.skip_symbol_norm:
+            sum = np.sum(list(next_char_pred.values()))
+            for char in self.symbol_set:
+                next_char_pred[char] /= sum
+
+        return list(sorted(next_char_pred.items(), key=lambda item: item[1], reverse=True))

textslinger/uniform.py

@@ -0,0 +1,65 @@
+"""Uniform language model"""
+from typing import Dict, List, Tuple, Union, Optional
+import numpy as np
+from textslinger.language_model import LanguageModel, DEFAULT_SYMBOL_SET
+
+
+class UniformLanguageModel(LanguageModel):
+    """Language model in which probabilities for symbols are uniformly
+    distributed.
+
+    Parameters
+    ----------
+        response_type - SYMBOL only
+        symbol_set - optional specify the symbol set, otherwise uses DEFAULT_SYMBOL_SET
+    """
+
+    def __init__(self,
+                 symbol_set: Optional[List[str]] = DEFAULT_SYMBOL_SET):
+        super().__init__(symbol_set=symbol_set)
+
+    def predict(self, evidence: Union[str, List[str]]) -> List[Tuple]:
+        """
+        Using the provided data, compute probabilities over the entire symbol.
+        set.
+
+        Parameters
+        ----------
+            evidence  - list of previously typed symbols
+
+        Returns
+        -------
+            list of (symbol, probability) tuples
+        """
+        probs = equally_probable(self.symbol_set)
+        return list(zip(self.symbol_set, probs))
+
+    def load(self) -> None:
+        """Restore model state from the provided checkpoint"""
+
+def equally_probable(alphabet: List[str],
+                     specified: Optional[Dict[str, float]] = None) -> List[float]:
+    """Returns a list of probabilities which correspond to the provided
+    alphabet. Unless overridden by the specified values, all items will
+    have the same probability. All probabilities sum to 1.0.
+
+    Parameters:
+    ----------
+        alphabet - list of symbols; a probability will be generated for each.
+        specified - dict of symbol => probability values for which we want to
+            override the default probability.
+    Returns:
+    --------
+        list of probabilities (floats)
+    """
+    n_letters = len(alphabet)
+    if not specified:
+        return np.full(n_letters, 1 / n_letters)
+
+    # copy specified dict ignoring non-alphabet items
+    overrides = {k: specified[k] for k in alphabet if k in specified}
+    assert sum(overrides.values()) < 1
+
+    prob = (1 - sum(overrides.values())) / (n_letters - len(overrides))
+    # override specified values
+    return [overrides[sym] if sym in overrides else prob for sym in alphabet]

textslinger-0.1.0.dist-info/METADATA

@@ -0,0 +1,84 @@
+Metadata-Version: 2.4
+Name: textslinger
+Version: 0.1.0
+Summary: TextSlinger text prediction toolkit.
+Author-email: Dylan Gaines <dcgaines@mtu.edu>, Keith Vertanen <vertanen@mtu.edu>, CAMBI <cambi_support@googlegroups.com>
+License: MIT License
+        
+        Copyright (c) 2025 Keith Vertanen, Dylan Gaines
+        
+        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.
+Project-URL: Homepage, https://www.cambi.tech/
+Project-URL: Source, https://github.com/kdv123/textpredict
+Platform: Linux
+Platform: Windows
+Platform: Mac OS-X
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Topic :: Scientific/Engineering :: Human Machine Interfaces
+Classifier: Natural Language :: English
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: End Users/Desktop
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Requires-Python: <3.11,>=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: torch==2.2.0
+Requires-Dist: torchvision==0.17.0
+Requires-Dist: torchaudio==2.2.0
+Requires-Dist: datasets==2.0.0
+Requires-Dist: bitsandbytes==0.42.0
+Requires-Dist: requests==2.32.3
+Requires-Dist: kenlm==0.1
+Requires-Dist: nlpaug==1.1.11
+Requires-Dist: psutil==5.7.2
+Requires-Dist: ipywidgets==8.1.3
+Requires-Dist: sentencepiece==0.2.0
+Requires-Dist: protobuf==4.25.3
+Requires-Dist: evaluate==0.4.0
+Requires-Dist: scikit-learn==1.2.2
+Requires-Dist: accelerate==0.33.0
+Requires-Dist: transformers==4.47.0
+Requires-Dist: numpy==1.24.4
+Requires-Dist: tqdm==4.62.2
+Provides-Extra: release
+Requires-Dist: twine==5.0.0; extra == "release"
+Requires-Dist: build==1.2.2.post1; extra == "release"
+Requires-Dist: wheel==0.43.0; extra == "release"
+Dynamic: license-file
+
+Python library for making text predictions using a language model.
+
+**** Setting up a Python environment ****
+These instructions are the same as those in the classify-aac which actually does training of models.
+But the same environment should also work for evaluation using the models (though you may only need a subset of the packages).
+This was tested on cheetah on 9/14/24.
+
+If you don't have anaconda installed in your user account you'll first need to do that.
+See: https://docs.anaconda.com/anaconda/install/linux/
+
+% conda create -n aacllm python=3.10
+% conda activate aacllm
+% conda install pytorch torchvision torchaudio pytorch-cuda cuda mpi4py -c pytorch -c nvidia
+% pip install 'git+https://github.com/potamides/uniformers.git#egg=uniformers'
+% pip install --upgrade transformers
+% pip install kenlm==0.1 --global-option="--max_order=12"
+% pip install rbloom bitsandbytes requests nlpaug ipywidgets psutil datasets sentencepiece protobuf evaluate scikit-learn deepspeed accelerate

textslinger-0.1.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+textslinger/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+textslinger/causal.py,sha256=QgDf6EvP1XuVmA9vevMwSyqkeTfpZhjN6ktn2s0Z8p0,21989
+textslinger/classifier.py,sha256=7p1zPocHlv93WaCtC4ddCpqxhwK4JgHKOy-t45Sjm70,8972
+textslinger/exceptions.py,sha256=UuetkKZj7YWTNDVxyRGXyGjLB70f1Padq3VHkhjcGG8,663
+textslinger/language_model.py,sha256=bVwjRei78MV_-PbH15b_icfLCa7Isusj6t-nUDrbGUo,1388
+textslinger/mixture.py,sha256=pa7WWCFDcSBpgKkK_FjfgupiiMMtc2hADRIl8_PlDaY,5026
+textslinger/ngram.py,sha256=2qU_vXF2C1WTw5V96vaFcg58YAQVoaOXaFGgBmNR8Xo,3848
+textslinger/uniform.py,sha256=-osO-sw96TBF26XYdo2SuoJP-cvVFCiooEmA3oy38Xc,2258
+textslinger-0.1.0.dist-info/licenses/LICENSE.md,sha256=9GrMOuuZ2WSNM16xWY2bWj2r0PgqJ98dzowfYcNtDjs,1084
+textslinger-0.1.0.dist-info/METADATA,sha256=3HIjEAwB1o37Z3nez0SFaFfn_yi41qSwlGXSR6bGjj4,3946
+textslinger-0.1.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+textslinger-0.1.0.dist-info/top_level.txt,sha256=cYLIX2fs3IGchHPWv9BvkQ59hycDDNy2fc68WTpTGpo,12
+textslinger-0.1.0.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+textslinger-0.1.0.dist-info/RECORD,,

textslinger-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (78.1.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

textslinger-0.1.0.dist-info/licenses/LICENSE.md

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Keith Vertanen, Dylan Gaines
+
+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.

textslinger-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+textslinger

textslinger-0.1.0.dist-info/zip-safe

@@ -0,0 +1 @@
+