languagemodels.inference

  1from typing import List
  2import requests
  3import re
  4import os
  5import sys
  6from time import perf_counter
  7
  8from languagemodels.models import get_model, get_model_info
  9from languagemodels.config import config
 10
 11
 12class InferenceException(Exception):
 13    pass
 14
 15
 16def list_tokens(prompt):
 17    """Generates a list of tokens for a supplied prompt
 18
 19    >>> list_tokens("Hello, world!") # doctest: +SKIP
 20    [('▁Hello', 8774), (',', 6), ('▁world', 296), ('!', 55)]
 21
 22    >>> list_tokens("Hello, world!")
 23    [('...Hello', ...), ... ('...world', ...), ...]
 24    """
 25    tokenizer, _ = get_model("instruct")
 26
 27    output = tokenizer.encode(prompt, add_special_tokens=False)
 28    tokens = output.tokens
 29    ids = output.ids
 30
 31    return list(zip(tokens, ids))
 32
 33
 34def generate_ts(engine, prompt, max_tokens=200):
 35    """Generates a single text response for a prompt from a textsynth server
 36
 37    The server and API key are provided as environment variables:
 38
 39    LANGUAGEMODELS_TS_SERVER is the server such as http://localhost:8080
 40    LANGUAGEMODELS_TS_KEY is the API key
 41    """
 42    apikey = os.environ.get("LANGUAGEMODELS_TS_KEY") or ""
 43    server = os.environ.get("LANGUAGEMODELS_TS_SERVER") or "https://api.textsynth.com"
 44
 45    response = requests.post(
 46        f"{server}/v1/engines/{engine}/completions",
 47        headers={"Authorization": f"Bearer {apikey}"},
 48        json={"prompt": prompt, "max_tokens": max_tokens},
 49    )
 50    resp = response.json()
 51    if "text" in resp:
 52        return resp["text"]
 53    else:
 54        raise InferenceException(f"TextSynth error: {resp}")
 55
 56
 57def generate_oa(engine, prompt, max_tokens=200, temperature=0):
 58    """Generates a single text response for a prompt using OpenAI
 59
 60    The server and API key are provided as environment variables:
 61
 62    LANGUAGEMODELS_OA_KEY is the API key
 63    """
 64    apikey = os.environ.get("LANGUAGEMODELS_OA_KEY")
 65
 66    response = requests.post(
 67        "https://api.openai.com/v1/completions",
 68        headers={
 69            "Authorization": f"Bearer {apikey}",
 70            "Content-Type": "application/json",
 71        },
 72        json={
 73            "model": engine,
 74            "prompt": prompt,
 75            "max_tokens": max_tokens,
 76            "temperature": temperature,
 77        },
 78    )
 79    resp = response.json()
 80
 81    try:
 82        return resp["choices"][0]["text"]
 83    except KeyError:
 84        raise InferenceException(f"OpenAI error: {resp}")
 85
 86
 87def chat_oa(engine, prompt, max_tokens=200, temperature=0):
 88    """Generates a single text response for a prompt using OpenAI
 89
 90    The server and API key are provided as environment variables:
 91
 92    LANGUAGEMODELS_OA_KEY is the API key
 93    """
 94    apikey = os.environ.get("LANGUAGEMODELS_OA_KEY")
 95
 96    response = requests.post(
 97        "https://api.openai.com/v1/chat/completions",
 98        headers={
 99            "Authorization": f"Bearer {apikey}",
100            "Content-Type": "application/json",
101        },
102        json={
103            "model": engine,
104            "messages": [{"role": "user", "content": prompt}],
105            "max_tokens": max_tokens,
106            "temperature": temperature,
107        },
108    )
109    resp = response.json()
110
111    try:
112        return resp["choices"][0]["message"]["content"]
113    except KeyError:
114        raise InferenceException(f"OpenAI error: {resp}")
115
116
117def stream_results(results, tokenizer):
118    """Map a token iterator to a substring iterator"""
119    tokens = []
120    last_len = 0
121
122    for result in results:
123        tokens.append(result.token_id)
124        text = tokenizer.decode(tokens)
125        yield text[last_len:]
126        last_len = len(text)
127
128
129def echo_results(results, tokenizer):
130    """Output results to stderr as they are collected"""
131    tokens = []
132    last_len = 0
133
134    for result in results:
135        tokens.append(result.token_id)
136        text = tokenizer.decode(tokens)
137        sys.stderr.write(text[last_len:])
138        sys.stderr.flush()
139        last_len = len(text)
140
141    sys.stderr.write("\n\n")
142    sys.stderr.flush()
143    return tokens
144
145
146def generate(
147    instructions: List[str],
148    max_tokens: int = 200,
149    temperature: float = 0.1,
150    topk: int = 1,
151    repetition_penalty: float = 0.0,
152    prefix: str = "",
153    suppress: List[str] = [],
154    model: str = "instruct",
155    stream: bool = False,
156):
157    """Generates completions for a prompt
158
159    This may use a local model, or it may make an API call to an external
160    model if API keys are available.
161
162    >>> generate(["What is the capital of France?"])
163    ['...Paris...']
164
165    >>> list(generate(["What is the capital of France?"], stream=True))
166    ['...Paris...']
167    """
168    if os.environ.get("LANGUAGEMODELS_TS_KEY") or os.environ.get(
169        "LANGUAGEMODELS_TS_SERVER"
170    ):
171        return generate_ts("flan_t5_xxl_q4", instructions, max_tokens).strip()
172
173    if os.environ.get("LANGUAGEMODELS_OA_KEY"):
174        return chat_oa("gpt-3.5-turbo", instructions, max_tokens).strip()
175
176    tokenizer, model = get_model(model)
177
178    start_time = perf_counter()
179
180    suppress = [tokenizer.encode(s, add_special_tokens=False).tokens for s in suppress]
181
182    model_info = get_model_info("instruct")
183
184    fmt = model_info.get("prompt_fmt", "{instruction}")
185
186    if repetition_penalty == 0.0:
187        repetition_penalty = model_info.get("repetition_penalty", 1.3)
188
189    prompts = [fmt.replace("{instruction}", inst) for inst in instructions]
190    prompts_tok = [tokenizer.encode(p).tokens for p in prompts]
191
192    outputs_ids = []
193    if hasattr(model, "translate_batch"):
194        prefix = tokenizer.encode(prefix, add_special_tokens=False).tokens
195        if stream or (config["echo"] and len(prompts_tok) == 1):
196            results = model.generate_tokens(
197                prompts_tok[0],
198                target_prefix=prefix,
199                repetition_penalty=repetition_penalty,
200                max_decoding_length=max_tokens,
201                sampling_temperature=temperature,
202                sampling_topk=topk,
203                suppress_sequences=suppress,
204            )
205
206            if stream:
207                return stream_results(results, tokenizer)
208            else:
209                outputs_ids = [echo_results(results, tokenizer)]
210        else:
211            results = model.translate_batch(
212                prompts_tok,
213                target_prefix=[prefix] * len(prompts),
214                repetition_penalty=repetition_penalty,
215                max_decoding_length=max_tokens,
216                sampling_temperature=temperature,
217                sampling_topk=topk,
218                suppress_sequences=suppress,
219                beam_size=1,
220            )
221            outputs_tokens = [r.hypotheses[0] for r in results]
222            for output in outputs_tokens:
223                outputs_ids.append([tokenizer.token_to_id(t) for t in output])
224    else:
225        if stream or (config["echo"] and len(prompts_tok) == 1):
226            results = model.generate_tokens(
227                prompts_tok,
228                repetition_penalty=repetition_penalty,
229                max_length=max_tokens,
230                sampling_temperature=temperature,
231                sampling_topk=topk,
232                suppress_sequences=suppress,
233            )
234
235            if stream:
236                return stream_results(results, tokenizer)
237            else:
238                outputs_ids = [echo_results(results, tokenizer)]
239        else:
240            results = model.generate_batch(
241                prompts_tok,
242                repetition_penalty=repetition_penalty,
243                max_length=max_tokens,
244                sampling_temperature=temperature,
245                sampling_topk=topk,
246                suppress_sequences=suppress,
247                beam_size=1,
248                include_prompt_in_result=False,
249            )
250            outputs_ids = [r.sequences_ids[0] for r in results]
251
252    model_info["requests"] = model_info.get("requests", 0) + len(prompts)
253
254    in_toks = sum(len(p) for p in prompts_tok)
255    model_info["input_tokens"] = model_info.get("input_tokens", 0) + in_toks
256
257    out_toks = sum(len(o) for o in outputs_ids)
258    model_info["output_tokens"] = model_info.get("output_tokens", 0) + out_toks
259
260    elapsed_time = perf_counter() - start_time
261    model_info["runtime"] = model_info.get("runtime", 0) + elapsed_time
262
263    return [tokenizer.decode(i, skip_special_tokens=True).lstrip() for i in outputs_ids]
264
265
266def rank_instruct(inputs, targets):
267    """Sorts a list of targets by their probabilities
268
269    >>> rank_instruct(["Classify positive or negative: I love python. Classification:"],
270    ... ['positive', 'negative'])
271    [['positive', 'negative']]
272
273    >>> rank_instruct(["Classify fantasy or documentary: "
274    ... "The wizard raised their wand. Classification:"],
275    ... ['fantasy', 'documentary'])
276    [['fantasy', 'documentary']]
277
278    >>> rank_instruct(["Say six", "Say seven"], ["six", "seven"])
279    [['six', 'seven'], ['seven', 'six']]
280    """
281    tokenizer, model = get_model("instruct")
282
283    targ_tok = [tokenizer.encode(t).tokens for t in targets]
284    targ_tok *= len(inputs)
285
286    in_tok = []
287    for input in inputs:
288        toks = [tokenizer.encode(input).tokens]
289        in_tok += toks * len(targets)
290
291    if "Generator" in str(type(model)):
292        scores = model.score_batch([i + t for i, t in zip(in_tok, targ_tok)])
293    else:
294        scores = model.score_batch(in_tok, target=targ_tok)
295
296    ret = []
297    for i in range(0, len(inputs) * len(targets), len(targets)):
298        logprobs = [sum(r.log_probs) for r in scores[i : i + len(targets)]]
299        results = sorted(zip(targets, logprobs), key=lambda r: -r[1])
300        ret.append([r[0] for r in results])
301
302    return ret
303
304
305def parse_chat(prompt):
306    """Converts a chat prompt using special tokens to a plain-text prompt
307
308    This is useful for prompting generic models that have not been fine-tuned
309    for chat using specialized tokens.
310
311    >>> parse_chat('User: What time is it?')
312    Traceback (most recent call last):
313        ....
314    inference.InferenceException: Chat prompt must end with 'Assistant:'
315
316    >>> parse_chat('''User: What time is it?
317    ...
318    ...               Assistant:''')
319    [{'role': 'user', 'content': 'What time is it?'}]
320
321    >>> parse_chat('''
322    ...              A helpful assistant
323    ...
324    ...              User: What time is it?
325    ...
326    ...              Assistant:
327    ...              ''')
328    [{'role': 'system', 'content': 'A helpful assistant'},
329     {'role': 'user', 'content': 'What time is it?'}]
330
331    >>> parse_chat('''
332    ...              A helpful assistant
333    ...
334    ...              User: What time is it?
335    ...
336    ...              Assistant: The time is
337    ...              ''')
338    Traceback (most recent call last):
339        ....
340    inference.InferenceException: Final assistant message must be blank
341
342    >>> parse_chat('''
343    ...              A helpful assistant
344    ...
345    ...              User: First para
346    ...
347    ...              Second para
348    ...
349    ...              Assistant:
350    ...              ''')
351    [{'role': 'system', 'content': 'A helpful assistant'},
352     {'role': 'user', 'content': 'First para\\n\\nSecond para'}]
353
354    >>> parse_chat('''
355    ...              A helpful assistant
356    ...
357    ...              User: What time is it?
358    ...
359    ...              InvalidRole: Nothing
360    ...
361    ...              Assistant:
362    ...              ''')
363    Traceback (most recent call last):
364        ....
365    inference.InferenceException: Invalid chat role: invalidrole
366    """
367
368    if not re.match(r"^\s*\w+:", prompt):
369        prompt = "System: " + prompt
370
371    prompt = "\n\n" + prompt
372
373    chunks = re.split(r"[\r\n]\s*(\w+):", prompt, flags=re.M)
374    chunks = [m.strip() for m in chunks if m.strip()]
375
376    messages = []
377
378    for i in range(0, len(chunks), 2):
379        role = chunks[i].lower()
380
381        try:
382            content = chunks[i + 1]
383            content = re.sub(r"\s*\n\n\s*", "\n\n", content)
384        except IndexError:
385            content = ""
386        messages.append({"role": role, "content": content})
387
388    for message in messages:
389        if message["role"] not in ["system", "user", "assistant"]:
390            raise InferenceException(f"Invalid chat role: {message['role']}")
391
392    if messages[-1]["role"] != "assistant":
393        raise InferenceException("Chat prompt must end with 'Assistant:'")
394
395    if messages[-1]["content"] != "":
396        raise InferenceException("Final assistant message must be blank")
397
398    return messages[:-1]