languagemodels

  1import requests
  2import datetime
  3import json
  4import re
  5from typing import overload
  6
  7from languagemodels.config import config
  8from languagemodels.inference import (
  9    generate,
 10    rank_instruct,
 11    parse_chat,
 12    list_tokens,
 13)
 14from languagemodels import embeddings
 15
 16docs = embeddings.RetrievalContext()
 17
 18
 19def complete(prompt: str) -> str:
 20    """Provide one completion for a given open-ended prompt
 21
 22    :param prompt: Prompt to use as input to the model
 23    :return: Completion returned from the language model
 24
 25    Examples:
 26
 27    >>> complete("Luke thought that he") #doctest: +SKIP
 28    'was going to be a doctor.'
 29
 30    >>> complete("There are many mythical creatures who") #doctest: +SKIP
 31    'are able to fly'
 32
 33    >>> complete("She hid in her room until") #doctest: +SKIP
 34    'she was sure she was safe'
 35    """
 36
 37    result = generate(
 38        ["Write a sentence"],
 39        prefix=prompt,
 40        max_tokens=config["max_tokens"],
 41        temperature=0.7,
 42        topk=40,
 43    )[0]
 44
 45    if result.startswith(prompt):
 46        prefix_length = len(prompt)
 47        return result[prefix_length:]
 48    else:
 49        return result
 50
 51
 52@overload
 53def do(prompt: list) -> list:
 54    ...
 55
 56
 57@overload
 58def do(prompt: str) -> str:
 59    ...
 60
 61
 62def do(prompt, choices=None):
 63    """Follow a single-turn instructional prompt
 64
 65    :param prompt: Instructional prompt(s) to follow
 66    :param choices: If provided, outputs are restricted to values in choices
 67    :return: Completion returned from the language model
 68
 69    Note that this function is overloaded to return a list of results if
 70    a list if of prompts is provided and a single string if a single
 71    prompt is provided as a string
 72
 73    Examples:
 74
 75    >>> do("Translate Spanish to English: Hola mundo!") #doctest: +SKIP
 76    'Hello world!'
 77
 78    >>> do("Pick the planet from the list: baseball, Texas, Saturn")
 79    '...Saturn...'
 80
 81    >>> do("Answer: What is the capital of England?")
 82    '...London...'
 83
 84    >>> do(["Pick the planet from the list: baseball, Texas, Saturn"] * 2)
 85    ['...Saturn...', '...Saturn...']
 86
 87    >>> do(["Say red", "Say blue"], choices=["red", "blue"])
 88    ['red', 'blue']
 89
 90    >>> do("Classify as positive or negative: LLMs are bad",
 91    ... choices=["Positive", "Negative"])
 92    'Negative'
 93
 94    >>> do("Classify as positive or negative: LLMs are great",
 95    ... choices=["Positive", "Negative"])
 96    'Positive'
 97    """
 98
 99    prompts = [prompt] if isinstance(prompt, str) else prompt
100
101    if choices:
102        results = [r[0] for r in rank_instruct(prompts, choices)]
103    else:
104        results = generate(prompts, max_tokens=config["max_tokens"], topk=1)
105
106        for i, result in enumerate(results):
107            if len(result.split()) == 1:
108                results[i] = result.title()
109
110                if result[-1] not in (".", "!", "?"):
111                    results[i] = results[i] + "."
112
113    return results[0] if isinstance(prompt, str) else results
114
115
116@overload
117def embed(doc: list) -> list:
118    ...
119
120
121@overload
122def embed(doc: str) -> str:
123    ...
124
125
126def embed(doc):
127    """Create embedding for a document
128
129    :param doc: Document(s) to embed
130    :return: Embedding
131
132    Note that this function is overloaded to return a list of embeddings if
133    a list if of docs is provided and a single embedding if a single
134    doc is provided as a string
135
136    Examples:
137
138    >>> embed("Hello, world")
139    [-0.0...]
140
141    >>> embed(["Hello", "world"])
142    [[-0.0...]]
143    """
144
145    docs = [doc] if isinstance(doc, str) else doc
146
147    # Create embeddings and convert to lists of floats
148    emb = [[float(n) for n in e] for e in embeddings.embed(docs)]
149
150    return emb[0] if isinstance(doc, str) else emb
151
152
153def chat(prompt: str) -> str:
154    """Get new message from chat-optimized language model
155
156    The `prompt` for this model is provided as a series of messages as a single
157    plain-text string. Several special tokens are used to delineate chat
158    messages.
159
160    - `system:` - Indicates the start of a system message providing
161    instructions about how the assistant should behave.
162    - `user:` - Indicates the start of a prompter (typically user)
163    message.
164    - `assistant:` - Indicates the start of an assistant message.
165
166    A complete prompt may look something like this:
167
168    ```
169    Assistant is helpful and harmless
170
171    User: What is the capital of Germany?
172
173    Assistant: The capital of Germany is Berlin.
174
175    User: How many people live there?
176
177    Assistant:
178    ```
179
180    The completion from the language model is returned.
181
182    :param message: Prompt using formatting described above
183    :return: Completion returned from the language model
184
185    Examples:
186
187    >>> response = chat('''
188    ...      System: Respond as a helpful assistant. It is 5:00pm.
189    ...
190    ...      User: What time is it?
191    ...
192    ...      Assistant:
193    ...      ''') # doctest: +SKIP
194    "It's 5:00pm."
195    """
196
197    messages = parse_chat(prompt)
198
199    # Suppress starts of all assistant messages to avoid repeat generation
200    suppress = [
201        "Assistant: " + m["content"].split(" ")[0]
202        for m in messages
203        if m["role"] in ["assistant", "user"]
204    ]
205
206    # Suppress all user messages to avoid repeating them
207    suppress += [m["content"] for m in messages if m["role"] == "user"]
208
209    system_msgs = [m for m in messages if m["role"] == "system"]
210    assistant_msgs = [m for m in messages if m["role"] == "assistant"]
211    user_msgs = [m for m in messages if m["role"] == "user"]
212
213    # The current model is tuned on instructions and tends to get
214    # lost if it sees too many questions
215    # Use only the most recent user and assistant message for context
216    # Keep all system messages
217    messages = system_msgs + assistant_msgs[-1:] + user_msgs[-1:]
218
219    rolemap = {
220        "system": "System",
221        "user": "Question",
222        "assistant": "Assistant",
223    }
224
225    messages = [f"{rolemap[m['role']]}: {m['content']}" for m in messages]
226
227    prompt = "\n\n".join(messages) + "\n\n" + "Assistant:"
228
229    if prompt.startswith("System:"):
230        prompt = prompt[7:].strip()
231
232    response = generate(
233        [prompt],
234        max_tokens=config["max_tokens"],
235        temperature=0.3,
236        topk=40,
237        prefix="Assistant:",
238        suppress=suppress,
239    )[0]
240
241    # Remove duplicate assistant being generated
242    if response.startswith("Assistant:"):
243        response = response[10:]
244
245    return response.strip()
246
247
248def code(prompt: str) -> str:
249    """Complete a code prompt
250
251    This assumes that users are expecting Python completions. Default models
252    are fine-tuned on Python where applicable.
253
254    :param prompt: Code context to complete
255    :return: Completion returned from the language model
256
257    Examples:
258
259    >>> code("# Print Hello, world!\\n")
260    'print("Hello, world!")\\n'
261
262    >>> code("def return_4():")
263    '...return 4...'
264    """
265    return generate([prompt], max_tokens=config["max_tokens"], topk=1, model="code")[0]
266
267
268def extract_answer(question: str, context: str) -> str:
269    """Extract an answer to a `question` from a provided `context`
270
271    The returned answer will always be a substring extracted from `context`.
272    It may not always be a correct or meaningful answer, but it will never be
273    an arbitrary hallucination.
274
275    :param question: A question to answer using knowledge from context
276    :param context: Knowledge used to answer the question
277    :return: Answer to the question.
278
279    Examples:
280
281    >>> context = "There is a green ball and a red box"
282    >>> extract_answer("What color is the ball?", context).lower()
283    '...green...'
284
285    >>> extract_answer("Who created Python?", get_wiki('Python')) #doctest: +SKIP
286    '...Guido van Rossum...'
287    """
288
289    return generate([f"{context}\n\n{question}"])[0]
290
291
292def classify(doc: str, label1: str, label2: str) -> str:
293    """Performs binary classification on an input
294
295    :param doc: A plain text input document to classify
296    :param label1: The first label to classify against
297    :param label2: The second label to classify against
298    :return: The closest matching class. The return value will always be
299    `label1` or `label2`
300
301    Examples:
302
303    >>> classify("That book was good.","positive","negative")
304    'positive'
305    >>> classify("That movie was terrible.","positive","negative")
306    'negative'
307    """
308
309    return do(
310        f"Classify as {label1} or {label2}: {doc}\n\nClassification:",
311        choices=[label1, label2],
312    )
313
314
315def store_doc(doc: str, name: str = "") -> None:
316    """Store document for later retrieval
317
318    :param doc: A plain text document to store.
319    :param name: Optional name for the document. This is used as a chunk prefix.
320
321    Examples:
322
323    >>> store_doc("The sky is blue.")
324    """
325    docs.store(doc, name)
326
327
328def load_doc(query: str) -> str:
329    """Load a matching document
330
331    A single document that best matches `query` will be returned.
332
333    :param query: Query to compare to stored documents
334    :return: Content of the closest matching document
335
336    Examples:
337
338    >>> store_doc("Paris is in France.")
339    >>> store_doc("The sky is blue.")
340    >>> load_doc("Where is Paris?")
341    'Paris is in France.'
342    """
343    return docs.get_match(query)
344
345
346def get_doc_context(query: str) -> str:
347    """Loads context from documents
348
349    A string representing the most relevant content from all stored documents
350    will be returned. This may be a blend of chunks from multiple documents.
351
352    :param query: Query to compare to stored documents
353    :return: Up to 128 tokens of context
354
355    Examples:
356
357    >>> store_doc("Paris is in France.")
358    >>> store_doc("Paris is nice.")
359    >>> store_doc("The sky is blue.")
360    >>> get_doc_context("Where is Paris?")
361    'Paris is in France.\\n\\nParis is nice.'
362    """
363    return docs.get_context(query)
364
365
366def get_wiki(topic: str) -> str:
367    """
368    Return Wikipedia summary for a topic
369
370    This function ignores the complexity of disambiguation pages and simply
371    returns the first result that is not a disambiguation page
372
373    :param topic: Topic to search for on Wikipedia
374    :return: Text content of the lead section of the most popular matching article
375
376    Examples:
377
378    >>> get_wiki('Python language')
379    'Python is a high-level...'
380
381    >>> get_wiki('Chemistry')
382    'Chemistry is the scientific study...'
383    """
384
385    url = "https://api.wikimedia.org/core/v1/wikipedia/en/search/title"
386    response = requests.get(url, params={"q": topic, "limit": 5})
387    response = json.loads(response.text)
388
389    for page in response["pages"]:
390        wiki_result = requests.get(
391            f"https://en.wikipedia.org/w/api.php?action=query&prop=extracts|pageprops&"
392            f"exintro&redirects=1&titles={page['title']}&format=json"
393        ).json()
394
395        first = wiki_result["query"]["pages"].popitem()[1]
396        if "disambiguation" in first["pageprops"]:
397            continue
398
399        summary = first["extract"]
400
401        cutoffs = [
402            "See_also",
403            "Notes",
404            "References",
405            "Further_reading",
406            "External_links",
407        ]
408
409        for cutoff in cutoffs:
410            summary = summary.split(f'<span id="{cutoff}">', 1)[0]
411
412        summary = re.sub(r"<p>", "\n\n", summary, flags=re.I)
413        summary = re.sub(r"<!\-\-.*?\-\->", "", summary, flags=re.I | re.DOTALL)
414        summary = re.sub(r"<.*?>", "", summary, flags=re.I)
415        summary = re.sub(r"\s*[\n\r]+\s*[\r\n]+[\s\r\n]*", "\n\n", summary, flags=re.I)
416        summary = summary.strip()
417        return summary
418    else:
419        return "No matching wiki page found."
420
421
422def get_weather(latitude, longitude):
423    """Fetch the current weather for a supplied longitude and latitude
424
425    Weather is provided by the US government and this function only supports
426    locations in the United States.
427
428    :param latitude: Latitude value representing this location
429    :param longitude: Longitude value representing this location
430    :return: Plain text description of the current weather forecast
431
432    Examples:
433
434    >>> get_weather(41.8, -87.6) # doctest: +SKIP
435    'Scattered showers and thunderstorms before 1pm with a high of 73.'
436    """
437
438    res = requests.get(f"https://api.weather.gov/points/{latitude},{longitude}")
439    points = json.loads(res.text)
440    forecast_url = points["properties"]["forecast"]
441
442    res = requests.get(forecast_url)
443    forecast = json.loads(res.text)
444    current = forecast["properties"]["periods"][0]
445
446    return current["detailedForecast"]
447
448
449def get_date() -> str:
450    """Returns the current date and time in natural language
451
452    >>> get_date() # doctest: +SKIP
453    'Friday, May 12, 2023 at 09:27AM'
454    """
455
456    now = datetime.datetime.now()
457
458    return now.strftime("%A, %B %d, %Y at %I:%M%p")
459
460
461def print_tokens(prompt: str) -> None:
462    """Prints a list of tokens in a prompt
463
464    :param prompt: Prompt to use as input to tokenizer
465    :return: Nothing
466
467    Examples:
468
469    >>> print_tokens("Hello world")
470    ' Hello' (token 8774)
471    ' world' (token 296)
472
473    >>> print_tokens("Hola mundo")
474    ' Hol' (token 5838)
475    'a' (token 9)
476    ' mun' (token 13844)
477    'd' (token 26)
478    'o' (token 32)
479    """
480
481    tokens = list_tokens(prompt)
482
483    for token in tokens:
484        print(f"'{token[0].replace('▁',' ')}' (token {token[1]})")
485
486
487def count_tokens(prompt: str) -> None:
488    """Counts tokens in a prompt
489
490    :param prompt: Prompt to use as input to tokenizer
491    :return: Nothing
492
493    Examples:
494
495    >>> count_tokens("Hello world")
496    2
497
498    >>> count_tokens("Hola mundo")
499    5
500    """
501
502    return len(list_tokens(prompt))
503
504
505def set_max_ram(value):
506    """Sets max allowed RAM
507
508    This value takes priority over environment variables
509
510    Returns the numeric value set in GB
511
512    >>> set_max_ram(16)
513    16.0
514
515    >>> set_max_ram('512mb')
516    0.5
517    """
518
519    config["max_ram"] = value
520
521    return config["max_ram"]
522
523
524def require_model_license(match_re):
525    """Require models to match supplied regex
526
527    This can be used to enforce certain licensing constraints when using this
528    package.
529    """
530    config["model_license"] = match_re