|
import time |
|
import gradio as gr |
|
import torch |
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer |
|
|
|
def run_LLM (model, tokenizer, streamer, prompt): |
|
|
|
token_ids = tokenizer.encode(prompt, return_tensors="pt") |
|
output_ids = model.generate( |
|
input_ids=token_ids.to(model.device), |
|
|
|
max_new_tokens=3000000, |
|
do_sample=True, |
|
temperature=0.8, |
|
) |
|
|
|
n_tokens = len(output_ids[0]) |
|
output_text = tokenizer.decode(output_ids[0]) |
|
|
|
return (output_text, n_tokens) |
|
|
|
def display_message(): |
|
model = AutoModelForCausalLM.from_pretrained("cyberagent/calm2-7b-chat", |
|
device_map="cuda", |
|
torch_dtype="auto") |
|
tokenizer = AutoTokenizer.from_pretrained("cyberagent/calm2-7b-chat") |
|
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) |
|
|
|
prompt = """わが国の経済について今後の予想を教えてください。 |
|
ASSISTANT: """ |
|
|
|
|
|
t_start = time.perf_counter() |
|
t_prev = t_start |
|
t_sum = 0.0 |
|
total_tokens = 0 |
|
|
|
log = '' |
|
|
|
for i in range(10): |
|
(result, n_tokens) = run_LLM(model, tokenizer, streamer, prompt) |
|
|
|
total_tokens = total_tokens + n_tokens |
|
t_curr = time.perf_counter() |
|
t_lap = t_curr - t_prev |
|
t_prev = t_curr |
|
t_sum = t_sum + t_lap |
|
t_avg = t_sum/(i+1.0) |
|
|
|
speed_now = n_tokens/t_lap |
|
speed_avg = total_tokens/t_sum |
|
|
|
|
|
row = [] |
|
row.append(i+1) |
|
row.append(t_lap) |
|
row.append(t_avg) |
|
row.append(speed_now) |
|
row.append(speed_avg) |
|
row.append(n_tokens) |
|
row.append(total_tokens) |
|
row.append(t_sum) |
|
|
|
|
|
line = "%d %f %f %f %f %d %d %f" % (i+1, t_lap, t_avg, speed_now, speed_avg, n_tokens, total_tokens, t_sum) |
|
log = log + line + "¥n" |
|
|
|
return log |
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
iface = gr.Interface(fn=display_message, inputs=None, outputs="text") |
|
iface.launch() |
|
|
