|
import dataclasses |
|
import pprint |
|
from functools import partial |
|
import re |
|
import os |
|
from threading import Lock |
|
import urllib |
|
import time |
|
from typing import List, Optional, Union |
|
|
|
from pydantic import BaseModel |
|
import absl.logging |
|
from tqdm import tqdm, trange |
|
import numpy as np |
|
import mlxu |
|
from ml_collections import ConfigDict |
|
import uvicorn |
|
from fastapi import FastAPI |
|
import gradio as gr |
|
import requests |
|
from requests.exceptions import Timeout, ConnectionError |
|
|
|
|
|
class InferenceRequest(BaseModel): |
|
prefix_text: Optional[List[str]] = None |
|
text: Optional[List[str]] = None |
|
until: Optional[Union[List[str], List[List[str]]]] = None |
|
temperature: Optional[float] = None |
|
|
|
|
|
class ChatRequest(BaseModel): |
|
prompt: str |
|
context: str = '' |
|
temperature: Optional[float] = None |
|
|
|
|
|
class LMServer(object): |
|
""" HTTP server for serving langauge models. """ |
|
|
|
@staticmethod |
|
def get_default_config(updates=None): |
|
config = ConfigDict() |
|
config.host = '0.0.0.0' |
|
config.port = 5007 |
|
config.batch_size = 1 |
|
config.logging = False |
|
config.pre_compile = 'loglikelihood' |
|
config.default_temperature = 1.0 |
|
config.greedy_until_max_length = 5000 |
|
config.prepend_to_prefix = '' |
|
config.append_to_prefix = '' |
|
config.prepend_to_text = '' |
|
config.append_to_text = '' |
|
config.chat_prepend_text = '' |
|
config.chat_user_prefix = '' |
|
config.chat_user_suffix = '' |
|
config.chat_lm_prefix = '' |
|
config.chat_lm_suffix = '' |
|
config.notes = '' |
|
|
|
if updates is not None: |
|
config.update(ConfigDict(updates).copy_and_resolve_references()) |
|
return config |
|
|
|
def __init__(self, config): |
|
self.config = self.get_default_config(config) |
|
self.lock = Lock() |
|
self.app = FastAPI() |
|
self.app.post('/loglikelihood')(self.serve_loglikelihood) |
|
self.app.post('/loglikelihood-rolling')(self.serve_loglikelihood_rolling) |
|
self.app.post('/generate')(self.serve_generate) |
|
self.app.post('/greedy-until')(self.serve_greedy_until) |
|
self.app.post('/chat')(self.serve_chat) |
|
self.app.get('/ready')(self.serve_ready) |
|
self.app = gr.mount_gradio_app(self.app, self.create_chat_app(), '/') |
|
|
|
@staticmethod |
|
def loglikelihood(prefix_text, text): |
|
raise NotImplementedError() |
|
|
|
@staticmethod |
|
def loglikelihood_rolling(text): |
|
raise NotImplementedError() |
|
|
|
@staticmethod |
|
def generate(text, temperature): |
|
raise NotImplementedError() |
|
|
|
@staticmethod |
|
def greedy_until(prefix_text, until, max_length): |
|
raise NotImplementedError() |
|
|
|
@staticmethod |
|
def to_list(x): |
|
if isinstance(x, np.ndarray): |
|
return x.tolist() |
|
return x |
|
|
|
def serve_ready(self): |
|
return 'Ready!\n' |
|
|
|
def serve_loglikelihood(self, data: InferenceRequest): |
|
with self.lock: |
|
if self.config.logging: |
|
absl.logging.info( |
|
'\n========= Serving Log Likelihood Request ========= \n' |
|
+ pprint.pformat(data) + '\n' |
|
) |
|
|
|
if data.prefix_text is None: |
|
data.prefix_text = ['' for _ in data.text] |
|
|
|
prefix_text = [ |
|
self.config.prepend_to_prefix + p + self.config.append_to_prefix |
|
for p in data.prefix_text |
|
] |
|
text = [ |
|
self.config.prepend_to_text + t + self.config.append_to_text |
|
for t in data.text |
|
] |
|
|
|
log_likelihood = [] |
|
is_greedy = [] |
|
for i in trange(0, len(text), self.config.batch_size, ncols=0): |
|
batch_prefix_text = prefix_text[i:i + self.config.batch_size] |
|
batch_text = text[i:i + self.config.batch_size] |
|
batch_size = len(batch_text) |
|
|
|
if batch_size < self.config.batch_size: |
|
extra = self.config.batch_size - batch_size |
|
batch_prefix_text.extend(['a' for _ in range(extra)]) |
|
batch_text.extend(['a' for _ in range(extra)]) |
|
|
|
batch_log_likelihood, batch_is_greedy = self.loglikelihood( |
|
batch_prefix_text, batch_text |
|
) |
|
batch_log_likelihood = self.to_list(batch_log_likelihood) |
|
batch_is_greedy = self.to_list(batch_is_greedy) |
|
log_likelihood.extend(batch_log_likelihood[:batch_size]) |
|
is_greedy.extend(batch_is_greedy[:batch_size]) |
|
|
|
output = { |
|
'prefix_text': data.prefix_text, |
|
'text': data.text, |
|
'log_likelihood': log_likelihood, |
|
'is_greedy': is_greedy, |
|
} |
|
if self.config.logging: |
|
absl.logging.info( |
|
'\n========= Output ========= \n' |
|
+ pprint.pformat(output) + '\n' |
|
) |
|
|
|
return output |
|
|
|
def serve_loglikelihood_rolling(self, data: InferenceRequest): |
|
with self.lock: |
|
if self.config.logging: |
|
absl.logging.info( |
|
'\n========= Serving Log Likelihood Request ========= \n' |
|
+ pprint.pformat(data) + '\n' |
|
) |
|
|
|
text = [ |
|
self.config.prepend_to_text + t + self.config.append_to_text |
|
for t in data.text |
|
] |
|
log_likelihood = [] |
|
is_greedy = [] |
|
for i in trange(0, len(text), self.config.batch_size, ncols=0): |
|
batch_text = text[i:i + self.config.batch_size] |
|
batch_size = len(batch_text) |
|
|
|
if batch_size < self.config.batch_size: |
|
extra = self.config.batch_size - batch_size |
|
batch_text.extend(['a' for _ in range(extra)]) |
|
|
|
batch_log_likelihood, batch_is_greedy = self.loglikelihood_rolling( |
|
batch_text |
|
) |
|
batch_log_likelihood = self.to_list(batch_log_likelihood) |
|
batch_is_greedy = self.to_list(batch_is_greedy) |
|
log_likelihood.extend(batch_log_likelihood[:batch_size]) |
|
is_greedy.extend(batch_is_greedy[:batch_size]) |
|
|
|
output = { |
|
'text': data.text, |
|
'log_likelihood': log_likelihood, |
|
'is_greedy': is_greedy, |
|
} |
|
if self.config.logging: |
|
absl.logging.info( |
|
'\n========= Output ========= \n' |
|
+ pprint.pformat(output) + '\n' |
|
) |
|
|
|
return output |
|
|
|
def serve_generate(self, data: InferenceRequest): |
|
with self.lock: |
|
if self.config.logging: |
|
absl.logging.info( |
|
'\n========= Serving Generate Request ========= \n' |
|
+ pprint.pformat(data) + '\n' |
|
) |
|
prefix_text = [ |
|
self.config.prepend_to_prefix + p + self.config.append_to_prefix |
|
for p in data.prefix_text |
|
] |
|
|
|
if data.temperature is None: |
|
data.temperature = self.config.default_temperature |
|
|
|
output_text = [] |
|
for i in trange(0, len(prefix_text), self.config.batch_size, ncols=0): |
|
batch_prefix_text = prefix_text[i:i + self.config.batch_size] |
|
batch_size = len(batch_prefix_text) |
|
|
|
if batch_size < self.config.batch_size: |
|
extra = self.config.batch_size - batch_size |
|
batch_prefix_text.extend(['a' for _ in range(extra)]) |
|
|
|
batch_output_text = self.generate( |
|
batch_prefix_text, |
|
temperature=data.temperature, |
|
) |
|
output_text.extend(self.to_list(batch_output_text)[:batch_size]) |
|
|
|
output = { |
|
'prefix_text': data.prefix_text, |
|
'output_text': output_text, |
|
'temperature': data.temperature, |
|
} |
|
if self.config.logging: |
|
absl.logging.info( |
|
'\n========= Output ========= \n' |
|
+ pprint.pformat(output) + '\n' |
|
) |
|
return output |
|
|
|
def serve_greedy_until(self, data: InferenceRequest): |
|
with self.lock: |
|
if self.config.logging: |
|
absl.logging.info( |
|
'\n========= Serving Greedy Until Request ========= \n' |
|
+ pprint.pformat(data) + '\n' |
|
) |
|
prefix_text = [ |
|
self.config.prepend_to_prefix + p + self.config.append_to_prefix |
|
for p in data.prefix_text |
|
] |
|
until = data.until |
|
max_length = self.config.greedy_until_max_length |
|
|
|
output_text = [] |
|
for i in range(0, len(prefix_text), self.config.batch_size): |
|
batch_prefix_text = prefix_text[i:i + self.config.batch_size] |
|
batch_until = until[i:i + self.config.batch_size] |
|
batch_size = len(batch_prefix_text) |
|
|
|
batch_output_text = self.greedy_until(batch_prefix_text, batch_until, max_length) |
|
output_text.extend(self.to_list(batch_output_text)[:batch_size]) |
|
|
|
output = { |
|
'prefix_text': data.prefix_text, |
|
'until': data.until, |
|
'max_length': max_length, |
|
'output_text': output_text, |
|
} |
|
if self.config.logging: |
|
absl.logging.info( |
|
'\n========= Output ========= \n' |
|
+ pprint.pformat(output) + '\n' |
|
) |
|
return output |
|
|
|
def process_chat(self, prompt, context, temperature): |
|
context = ( |
|
context + self.config.chat_user_prefix |
|
+ prompt + self.config.chat_user_suffix |
|
+ self.config.chat_lm_prefix |
|
) |
|
response = self.generate( |
|
[self.config.chat_prepend_text + context], |
|
temperature=float(temperature), |
|
)[0] |
|
context = context + response + self.config.chat_lm_suffix |
|
return response, context |
|
|
|
def serve_chat(self, data: ChatRequest): |
|
if data.temperature is None: |
|
data.temperature = self.config.default_temperature |
|
response, context = self.process_chat( |
|
data.prompt, data.context, |
|
temperature=data.temperature, |
|
) |
|
return { |
|
'response': response, |
|
'context': context, |
|
'temperature': data.temperature, |
|
} |
|
|
|
def create_chat_app(self): |
|
with gr.Blocks(analytics_enabled=False, title='EasyLM Chat') as gradio_chatbot: |
|
gr.Markdown('# Chatbot Powered by [EasyLM](https://github.com/young-geng/EasyLM)') |
|
gr.Markdown(self.config.notes) |
|
chatbot = gr.Chatbot(label='Chat history') |
|
msg = gr.Textbox( |
|
placeholder='Type your message here...', |
|
show_label=False |
|
) |
|
with gr.Row(): |
|
send = gr.Button('Send') |
|
regenerate = gr.Button('Regenerate', interactive=False) |
|
clear = gr.Button('Reset') |
|
|
|
temp_slider = gr.Slider( |
|
label='Temperature', minimum=0, maximum=2.0, |
|
value=self.config.default_temperature |
|
) |
|
|
|
context_state = gr.State(['', '']) |
|
|
|
def user_fn(user_message, history, context): |
|
return { |
|
msg: gr.update(value='', interactive=False), |
|
clear: gr.update(interactive=False), |
|
send: gr.update(interactive=False), |
|
regenerate: gr.update(interactive=False), |
|
chatbot: history + [[user_message, None]], |
|
context_state: [context[1], context[1]], |
|
} |
|
|
|
def model_fn(history, context, temperature): |
|
history[-1][1], new_context = self.process_chat( |
|
history[-1][0], context[0], temperature |
|
) |
|
return { |
|
msg: gr.update(value='', interactive=True), |
|
clear: gr.update(interactive=True), |
|
send: gr.update(interactive=True), |
|
chatbot: history, |
|
context_state: [context[0], new_context], |
|
regenerate: gr.update(interactive=True), |
|
} |
|
|
|
def regenerate_fn(): |
|
return { |
|
msg: gr.update(value='', interactive=False), |
|
clear: gr.update(interactive=False), |
|
send: gr.update(interactive=False), |
|
regenerate: gr.update(interactive=False), |
|
} |
|
|
|
def clear_fn(): |
|
return { |
|
chatbot: None, |
|
msg: '', |
|
context_state: ['', ''], |
|
regenerate: gr.update(interactive=False), |
|
} |
|
|
|
msg.submit( |
|
user_fn, |
|
inputs=[msg, chatbot, context_state], |
|
outputs=[msg, clear, send, chatbot, context_state, regenerate], |
|
queue=False |
|
).then( |
|
model_fn, |
|
inputs=[chatbot, context_state, temp_slider], |
|
outputs=[msg, clear, send, chatbot, context_state, regenerate], |
|
queue=True |
|
) |
|
send.click( |
|
user_fn, |
|
inputs=[msg, chatbot, context_state], |
|
outputs=[msg, clear, send, chatbot, context_state, regenerate], |
|
queue=False |
|
).then( |
|
model_fn, |
|
inputs=[chatbot, context_state, temp_slider], |
|
outputs=[msg, clear, send, chatbot, context_state, regenerate], |
|
queue=True |
|
) |
|
regenerate.click( |
|
regenerate_fn, |
|
inputs=None, |
|
outputs=[msg, clear, send, regenerate], |
|
queue=False |
|
).then( |
|
model_fn, |
|
inputs=[chatbot, context_state, temp_slider], |
|
outputs=[msg, clear, send, chatbot, context_state, regenerate], |
|
queue=True |
|
) |
|
clear.click( |
|
clear_fn, |
|
inputs=None, |
|
outputs=[chatbot, msg, context_state, regenerate], |
|
queue=False |
|
) |
|
|
|
gradio_chatbot.queue(concurrency_count=1) |
|
return gradio_chatbot |
|
|
|
def run(self): |
|
if self.config.pre_compile != '': |
|
if self.config.pre_compile == 'all': |
|
pre_compile = ['loglikelihood', 'generate', 'greedy_until', 'chat'] |
|
else: |
|
pre_compile = self.config.pre_compile.split(',') |
|
|
|
pre_compile_data = ['a' for _ in range(self.config.batch_size)] |
|
for task in pre_compile: |
|
if task == 'loglikelihood': |
|
self.loglikelihood(pre_compile_data, pre_compile_data) |
|
self.loglikelihood_rolling(pre_compile_data) |
|
elif task == 'generate': |
|
self.generate(pre_compile_data, 1.0) |
|
elif task == 'greedy_until': |
|
self.greedy_until( |
|
pre_compile_data, pre_compile_data, |
|
self.config.greedy_until_max_length |
|
) |
|
elif task == 'chat': |
|
self.process_chat('a', 'a', 1.0) |
|
else: |
|
raise ValueError(f'Invalid precompile task: {task}!') |
|
|
|
uvicorn.run(self.app, host=self.config.host, port=self.config.port) |
|
|
|
|
|
class LMClient(object): |
|
""" A simple client for the LM server. """ |
|
|
|
@staticmethod |
|
def get_default_config(updates=None): |
|
config = ConfigDict() |
|
config.url = 'http://localhost:5007' |
|
config.batch_size = 1 |
|
config.wait_for_ready = True |
|
config.dummy = False |
|
|
|
if updates is not None: |
|
config.update(ConfigDict(updates).copy_and_resolve_references()) |
|
return config |
|
|
|
def __init__(self, config=None): |
|
self.config = self.get_default_config(config) |
|
if self.config.wait_for_ready: |
|
self.wait_for_ready() |
|
|
|
def wait_for_ready(self): |
|
if self.config.dummy: |
|
return |
|
while True: |
|
try: |
|
requests.get(urllib.parse.urljoin(self.config.url, 'ready')) |
|
return |
|
except (Timeout, ConnectionError) as e: |
|
time.sleep(10) |
|
|
|
@staticmethod |
|
def batched(iterator, batch_size): |
|
batch = [] |
|
for example in iterator: |
|
batch.append(example) |
|
if len(batch) == batch_size: |
|
yield batch |
|
batch = [] |
|
if len(batch) > 0: |
|
yield batch |
|
|
|
def loglikelihood(self, prefix, text): |
|
prefix, text = list(prefix), list(text) |
|
if self.config.dummy: |
|
return [-1.0 for _ in text], [False for _ in text] |
|
|
|
log_likelihood = [] |
|
is_greedy = [] |
|
|
|
batched_iterator = list(zip( |
|
self.batched(prefix, self.config.batch_size), |
|
self.batched(text, self.config.batch_size) |
|
)) |
|
for batch_prefix, batch_text in tqdm(batched_iterator, ncols=0): |
|
response = requests.post( |
|
urllib.parse.urljoin(self.config.url, 'loglikelihood'), |
|
json={'prefix_text': batch_prefix, 'text': batch_text} |
|
).json() |
|
log_likelihood.extend(response['log_likelihood']) |
|
is_greedy.extend(response['is_greedy']) |
|
|
|
return log_likelihood, is_greedy |
|
|
|
def loglikelihood_rolling(self, text): |
|
text = list(text) |
|
if self.config.dummy: |
|
return [-1.0 for _ in text], [False for _ in text] |
|
|
|
log_likelihood = [] |
|
is_greedy = [] |
|
batched_iterator = list(self.batched(text, self.config.batch_size)) |
|
for batch_text in tqdm(batched_iterator, ncols=0): |
|
response = requests.post( |
|
urllib.parse.urljoin(self.config.url, 'loglikelihood-rolling'), |
|
json={'text': batch_text} |
|
).json() |
|
log_likelihood.extend(response['log_likelihood']) |
|
is_greedy.extend(response['is_greedy']) |
|
return log_likelihood, is_greedy |
|
|
|
def greedy_until(self, prefix, until): |
|
prefix, until = list(prefix), list(until) |
|
if self.config.dummy: |
|
results = [] |
|
for u in until: |
|
if isinstance(u, str): |
|
results.append('dummy text ' + u) |
|
else: |
|
results.append('dummy text ' + u[0]) |
|
return results |
|
|
|
batched_iterator = list(zip( |
|
self.batched(prefix, self.config.batch_size), |
|
self.batched(until, self.config.batch_size), |
|
)) |
|
output_text = [] |
|
for batch_prefix, batch_until in tqdm(batched_iterator, ncols=0): |
|
response = requests.post( |
|
urllib.parse.urljoin(self.config.url, 'greedy-until'), |
|
json={'prefix_text': batch_prefix, 'until': batch_until} |
|
).json() |
|
output_text.extend(response['output_text']) |
|
return output_text |
|
|
|
def generate(self, prefix, temperature=None): |
|
prefix = list(prefix) |
|
if self.config.dummy: |
|
return ['' for _ in prefix] |
|
|
|
output_text = [] |
|
batched_iterator = list(self.batched(prefix, self.config.batch_size)) |
|
for batch_prefix in tqdm(batched_iterator, ncols=0): |
|
response = requests.post( |
|
urllib.parse.urljoin(self.config.url, 'generate'), |
|
json={ |
|
'prefix_text': batch_prefix, |
|
'temperature': temperature, |
|
} |
|
).json() |
|
output_text.extend(response['output_text']) |
|
return output_text |
|
|
|
def chat(self, prompt, context, temperature=None): |
|
if self.config.dummy: |
|
return '' |
|
response = requests.post( |
|
urllib.parse.urljoin(self.config.url, 'chat'), |
|
json={ |
|
'prompt': prompt, |
|
'context': context, |
|
'temperature': temperature, |
|
} |
|
).json() |
|
return response['response'], response['context'] |
|
|
