languagemodels

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