1import numpy as np
  2from time import perf_counter
  3
  4from languagemodels.models import get_model, get_model_info
  5
  6
  7def embed(docs):
  8    """Compute embeddings for a batch of documents
  9
 10    >>> embed(["I love Python!"])[0].shape
 11    (384,)
 12
 13    >>> embed(["I love Python!"])[0][-3:]
 14    array([0.1..., 0.1..., 0.0...], dtype=float32)
 15
 16    >>> float(np.linalg.norm(embed(["I love Python!"])[0]))
 17    1.0
 18
 19    Embeddings are computed by running the first 512 tokens of each doc
 20    through a forward pass of the embedding model. The last hidden state
 21    of the model is mean pooled to produce a single vector
 22
 23    Documents will be processed in batches. The batch size is fixed at 64
 24    as this size was found to maximize throughput on a number of test
 25    systems while limiting memory usage.
 26    """
 27
 28    tokenizer, model = get_model("embedding")
 29    model_info = get_model_info("embedding")
 30
 31    start_time = perf_counter()
 32
 33    tokens = [tokenizer.encode(doc[:8192]).ids[:512] for doc in docs]
 34
 35    def mean_pool(last_hidden_state):
 36        embedding = np.mean(last_hidden_state, axis=0)
 37        embedding = embedding / np.linalg.norm(embedding)
 38        return embedding
 39
 40    bs = 64
 41    embeddings = []
 42    for i in range(0, len(docs), bs):
 43        outputs = model.forward_batch(tokens[i : i + bs])
 44        embeddings += [mean_pool(lhs) for lhs in np.array(outputs.last_hidden_state)]
 45
 46    model_info["requests"] = model_info.get("requests", 0) + len(tokens)
 47
 48    in_toks = sum(len(d) for d in tokens)
 49    model_info["input_tokens"] = model_info.get("input_tokens", 0) + in_toks
 50
 51    runtime = perf_counter() - start_time
 52    model_info["runtime"] = model_info.get("runtime", 0) + runtime
 53
 54    return embeddings
 55
 56
 57def search(query, docs, count=16):
 58    """Return `count` `docs` sorted by match against `query`
 59
 60    :param query: Input to match in search
 61    :param docs: List of docs to search against
 62    :param count: Number of document to return
 63    :return: List of (doc_num, score) tuples sorted by score descending
 64    """
 65
 66    prefix = get_model_info("embedding").get("query_prefix", "")
 67
 68    query_embedding = embed([f"{prefix}{query}"])[0]
 69
 70    scores = np.dot([d.embedding for d in docs], query_embedding)
 71
 72    return [(i, scores[i]) for i in reversed(np.argsort(scores)[-count:])]
 73
 74
 75def get_token_ids(doc):
 76    """Return list of token ids for a document
 77
 78    Note that the tokenzier used here is from the generative model.
 79
 80    This is used for token counting for the context, not for tokenization
 81    before embedding.
 82    """
 83
 84    generative_tokenizer, _ = get_model("instruct", tokenizer_only=True)
 85
 86    # We need to disable and re-enable truncation here
 87    # This allows us to tokenize very large documents
 88    # We won't be feeding the tokens themselves to a model, so this
 89    # shouldn't cause any problems.
 90    trunk = generative_tokenizer.truncation
 91    if trunk:
 92        generative_tokenizer.no_truncation()
 93    ids = generative_tokenizer.encode(doc, add_special_tokens=False).ids
 94    if trunk:
 95        generative_tokenizer.enable_truncation(
 96            trunk["max_length"], stride=trunk["stride"], strategy=trunk["strategy"]
 97        )
 98
 99    return ids
100
101
102def chunk_doc(doc, name="", chunk_size=64, chunk_overlap=8):
103    """Break a document into chunks
104
105    :param doc: Document to chunk
106    :param name: Optional document name
107    :param chunk_size: Length of individual chunks in tokens
108    :param chunk_overlap: Number of tokens to overlap when breaking chunks
109    :return: List of strings representing the chunks
110
111    The simple chunking approach used here consist of the following:
112
113    1. Attempt to chunk the remainder of the document.
114    2. If we can't fit all tokens in chunk_size, backtrack to look for a
115    meaningful cut point.
116    3. If a cut point is found, use that as the chunk boundary. There will
117    be no overlap between this chunk and the next in this case.
118    4. If a cut point is not found, use chunk_size a the boundary. There
119    will be chunk_overlap overlapping tokens starting the next chunk.
120    5. Repeat until entire document has been split into chunks.
121
122    >>> chunk_doc("")
123    []
124
125    >>> chunk_doc(
126    ... "It was the best of times, it was the worst of times, it was the age "
127    ... "of wisdom, it was the age of foolishness, it was the epoch of belief, "
128    ... "it was the epoch of incredulity, it was the season of Light, it was "
129    ... "the season of Darkness, it was the spring of hope, it was the winter "
130    ... "of despair, we had everything before us, we had nothing before us, we "
131    ... "were all going direct to Heaven, we were all going direct the other "
132    ... "way—in short, the period was so far like the present period, that "
133    ... "some of its noisiest authorities insisted on its being received, for "
134    ... "good or for evil, in the superlative degree of comparison only.")
135    ['It was the best of times...']
136
137    >>> chunk_doc(
138    ... "One morning, when Gregor Samsa woke from troubled dreams, he found "
139    ... "himself transformed in his bed into a horrible vermin. He lay on his "
140    ... "armour-like back, and if he lifted his head a little he could see "
141    ... "his brown belly, slightly domed and divided by arches into stiff "
142    ... "sections. The bedding was hardly able to cover it and seemed ready "
143    ... "to slide off any moment. His many legs, pitifully thin compared with "
144    ... "the size of the rest of him, waved about helplessly as he looked.")
145    ['One morning, ...']
146
147    >>> chunk_doc("Hello")
148    ['Hello']
149
150    >>> chunk_doc("Hello " * 65)
151    ['Hello Hello...', 'Hello...']
152
153    >>> chunk_doc("Hello world. " * 24)[0]
154    'Hello world. ...Hello world.'
155
156    >>> len(chunk_doc("Hello world. " * 20))
157    1
158
159    >>> len(chunk_doc("Hello world. " * 24))
160    2
161
162    # Check to make sure sentences aren't broken on decimal points
163    >>> chunk_doc(('z. ' + ' 37.468 ' * 5) * 3)[0]
164    'z. 37.468 ...z.'
165    """
166    generative_tokenizer, _ = get_model("instruct", tokenizer_only=True)
167
168    tokens = get_token_ids(doc)
169
170    separator_tokens = [".", "!", "?", ").", "\n\n", "\n", '."']
171
172    separators = [get_token_ids(t)[-1] for t in separator_tokens]
173
174    name_tokens = []
175
176    label = f"From {name} document:" if name else ""
177
178    if name:
179        name_tokens = get_token_ids(label)
180
181    i = 0
182    chunks = []
183    chunk = name_tokens.copy()
184    while i < len(tokens):
185        token = tokens[i]
186        chunk.append(token)
187        i += 1
188
189        # Save the last chunk if we're done
190        if i == len(tokens):
191            chunks.append(generative_tokenizer.decode(chunk))
192            break
193
194        if len(chunk) == chunk_size:
195            # Backtrack to find a reasonable cut point
196            for j in range(1, chunk_size // 2):
197                if chunk[chunk_size - j] in separators:
198                    ctx = generative_tokenizer.decode(
199                        chunk[chunk_size - j : chunk_size - j + 2]
200                    )
201                    if " " in ctx or "\n" in ctx:
202                        # Found a good separator
203                        text = generative_tokenizer.decode(chunk[: chunk_size - j + 1])
204                        chunks.append(text)
205                        chunk = name_tokens + chunk[chunk_size - j + 1 :]
206                        break
207            else:
208                # No semantically meaningful cutpoint found
209                # Default to a hard cut
210                text = generative_tokenizer.decode(chunk)
211                chunks.append(text)
212                # Share some overlap with next chunk
213                overlap = max(
214                    chunk_overlap, chunk_size - len(name_tokens) - (len(tokens) - i)
215                )
216                chunk = name_tokens + chunk[-overlap:]
217
218    return chunks
219
220
221class Document:
222    """
223    A document used for semantic search
224
225    Documents have content and an embedding that is used to match the content
226    against other semantically similar documents.
227    """
228
229    def __init__(self, content, name="", embedding=None):
230        self.content = content
231        self.embedding = embedding if embedding is not None else embed([content])[0]
232        self.name = name
233
234
235class RetrievalContext:
236    """
237    Provides a context for document retrieval
238
239    Documents are embedded and cached for later search.
240
241    Example usage:
242
243    >>> rc = RetrievalContext()
244    >>> rc.store("Paris is in France.")
245    >>> rc.store("The sky is blue.")
246    >>> rc.store("Mars is a planet.")
247    >>> rc.get_match("Paris is in France.")
248    'Paris is in France.'
249
250    >>> rc.get_match("Where is Paris?")
251    'Paris is in France.'
252
253    >>> rc.clear()
254    >>> rc.get_match("Where is Paris?")
255
256    >>> rc.clear()
257    >>> rc.store(' '.join(['Python'] * 4096))
258    >>> len(rc.chunks)
259    73
260
261    >>> rc.clear()
262    >>> rc.store(' '.join(['Python'] * 232))
263    >>> len(rc.chunks)
264    4
265
266    >>> rc.get_context("What is Python?")
267    'Python Python Python...'
268
269    >>> [len(c.content.split()) for c in rc.chunks]
270    [64, 64, 64, 64]
271
272    >>> len(rc.get_context("What is Python?").split())
273    128
274    """
275
276    def __init__(self, chunk_size=64, chunk_overlap=8):
277        self.chunk_size = chunk_size
278        self.chunk_overlap = chunk_overlap
279        self.clear()
280
281    def clear(self):
282        self.docs = []
283        self.chunks = []
284
285    def store(self, doc, name=""):
286        """Stores a document along with embeddings
287
288        This stores both the document as well as document chunks
289
290        >>> rc = RetrievalContext()
291        >>> rc.clear()
292        >>> rc.store(' '.join(['Python'] * 233))
293        >>> len(rc.chunks)
294        5
295
296        >>> rc.clear()
297        >>> rc.store(' '.join(['Python'] * 232))
298        >>> len(rc.chunks)
299        4
300
301        >>> rc.clear()
302        >>> rc.store('Python')
303        >>> len(rc.chunks)
304        1
305
306        >>> rc.clear()
307        >>> rc.store('It is a language.', 'Python')
308        >>> len(rc.chunks)
309        1
310        >>> [c.content for c in rc.chunks]
311        ['From Python document: It is a language.']
312
313        >>> rc = RetrievalContext()
314        >>> rc.clear()
315        >>> rc.store(' '.join(['details'] * 217), 'Python')
316        >>> len(rc.chunks)
317        5
318
319        >>> rc.clear()
320        >>> rc.store(' '.join(['details'] * 216), 'Python')
321        >>> len(rc.chunks)
322        4
323        >>> [c.content for c in rc.chunks]
324        ['From Python document: details details details...']
325        """
326
327        if doc not in self.docs:
328            self.docs.append(Document(doc))
329            self.store_chunks(doc, name)
330
331    def store_chunks(self, doc, name=""):
332        chunks = chunk_doc(doc, name, self.chunk_size, self.chunk_overlap)
333
334        embeddings = embed(chunks)
335
336        for embedding, chunk in zip(embeddings, chunks):
337            self.chunks.append(Document(chunk, embedding=embedding))
338
339    def get_context(self, query, max_tokens=128):
340        """Gets context matching a query
341
342        Context is capped by token length and is retrieved from stored
343        document chunks
344        """
345
346        if len(self.chunks) == 0:
347            return None
348
349        results = search(query, self.chunks)
350
351        chunks = []
352        tokens = 0
353
354        for chunk_id, score in results:
355            chunk = self.chunks[chunk_id].content
356            chunk_tokens = len(get_token_ids(chunk))
357            if tokens + chunk_tokens <= max_tokens and score > 0.1:
358                chunks.append(chunk)
359                tokens += chunk_tokens
360
361        context = "\n\n".join(chunks)
362
363        return context
364
365    def get_match(self, query):
366        if len(self.docs) == 0:
367            return None
368
369        return self.docs[search(query, self.docs)[0][0]].content