Infini-gram and its Uses in our Ontology Topology Project

Introduction

Infini-gram is an n-gram Language Model developed at the University of Washington. Infini-gram utilizes the largest possible “n” context rather than a fixed context like traditional n-gram models. The unbounded context allows infini-gram to provide more accurate n-gram predictions than traditional models. Infini-gram’s efficient use of a suffix array allows it to have fast expression lookups within large text corpora without using large amounts of space. In our paper, “Ontologies in Language Model Unembedding Matrices”, we used infini-gram to calculate frequencies of terms in comparison to language models’ ontological accuracy of these terms. Infini-gram’s term lookup performance made it a helpful tool in collecting this data but it has many uses outside of expression counting including:

• Counting the number of occurrences of terms of any length within text corpuses or training data.

• Finding the probability of the last token of a given term

• Finding probable next word predictions for a given term (using the greatest possible token context)

• Find samples of all documents containing a given term

How Infini-gram Works

Infini-gram’s largest possible “n” is defined as the longest context string for which there are more than 0 results inside of the corpus.

For example, given the prompt “...conducts research at the Paul G. Allen School of Computer Science and Engineering, University of”, a traditional 5-gram LM will use “Engineering, University of” (5 tokens) as context while the infini-gram model will find the count of the largest possible context string available in the corpus: “research at the Paul G. Allen School of Computer Science and Engineering, University of” for which the only possible continuation of the string is “_Washington”. (As provided by infini-gram’s site here.)

Infini-gram is able to utilize this strategy through its use of a suffix array. In order to create this array, every possible suffix that makes up the targeted corpus is given an integer index based off of their position in the corpus. The indices are lexicographically sorted based on the lexicographical order of the suffixes that they represent and are stored in the suffix array. The lexicographically sorted nature of the suffix arrays allows infini-gram to use binary searches to find and count lexicographically equal occurrences of a given term. Infini-gram achieves O(n) space and (when searching for a term T) \(O(logn* |T|)\) searching rather than traditional O(n*|T|).

{
  ‘Index’: ‘v4_pileval_llama’,
  ‘query_type’: ‘count’,
  ‘query’: ‘programmer’
}

curl -X POST -H "Content-Type: application/json" -d '{"index": "v4_rpj_llama_s4", "query_type": "count", "query": "programmer"}' https://api.infini-gram.io/

import requests
                      
payload = {
  'index': 'v4_rpj_llama_s4',
  'query_type': 'count',
  'query': 'programmer',
}
result = requests.post('https://api.infini-gram.io/', json=payload).json()
print(result)

from infini_gram.engine import InfiniGramEngine
                        from transformers import AutoTokenizer
                        
                        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", add_bos_token=False, add_eos_token=False)
                        engine = InfiniGramEngine(index_dir='index/v4_pileval_llama', eos_token_id=tokenizer.eos_token_id)