Spaces:
Runtime error
Runtime error
File size: 3,908 Bytes
88f55d9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import torch
from transformers import StoppingCriteria, StoppingCriteriaList
import copy
import json
import global_vars
from chats import pre, post
from pingpong import PingPong
from gens.batch_gen import get_output_batch
from pingpong.context import CtxLastWindowStrategy
class StopOnTokens(StoppingCriteria):
def __init__(self, tokenizer):
super().__init__()
self.stop_token_ids = tokenizer.convert_tokens_to_ids(
["<|im_end|>", "<|endoftext|>"]
)
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
for stop_id in self.stop_token_ids:
if input_ids[0][-1] == stop_id:
return True
return False
def build_prompts(ppmanager, user_message, global_context, win_size=3):
dummy_ppm = copy.deepcopy(ppmanager)
dummy_ppm.ctx = global_context
for pingpong in dummy_ppm.pingpongs:
pong = pingpong.pong
first_sentence = pong.split("\n")[0]
if first_sentence != "" and \
pre.contains_image_markdown(first_sentence):
pong = ' '.join(pong.split("\n")[1:]).strip()
pingpong.pong = pong
lws = CtxLastWindowStrategy(win_size)
prompt = lws(dummy_ppm)
return prompt
def text_stream(ppmanager, streamer):
for new_text in streamer:
ppmanager.append_pong(new_text)
yield ppmanager, ppmanager.build_uis()
yield ppmanager, ppmanager.build_uis()
def summarize(
ppmanager, prompt_to_summarize, win_size,
temperature, top_p, top_k, repetition_penalty, max_new_tokens,
num_beams, use_cache, do_sample, eos_token_id, pad_token_id
):
ctx = ppmanager.ctx
last_pong = ppmanager.pingpongs[-1].pong
ppmanager.add_pingpong(PingPong(prompt_to_summarize, ""))
prompt = ppmanager.build_prompts(from_idx=-win_size)
_, gen_config_summarization = pre.build_gen_config(
temperature, top_p, top_k, repetition_penalty, max_new_tokens,
num_beams, use_cache, do_sample, eos_token_id, pad_token_id
)
summarize_output = get_output_batch(
global_vars.model, global_vars.tokenizer, [prompt], gen_config_summarization
)[0].strip()
ppmanager.ctx = summarize_output
ppmanager.pop_pingpong()
return ppmanager
def chat_stream(
idx, local_data, user_message, state, model_num,
global_context, ctx_num_lconv, ctx_sum_prompt,
res_temp, res_topp, res_topk, res_rpen, res_mnts, res_beams, res_cache, res_sample, res_eosid, res_padid,
):
res = [
state["ppmanager_type"].from_json(json.dumps(ppm))
for ppm in local_data
]
ppm = res[idx]
# add_ping returns a prompt structured in Alpaca form
ppm.add_pingpong(
PingPong(user_message, "")
)
prompt = build_prompts(ppm, user_message, global_context, ctx_num_lconv)
# prepare text generating streamer & start generating
gen_kwargs, streamer = pre.build(
prompt,
res_temp, res_topp, res_topk, res_rpen, res_mnts,
res_beams, res_cache, res_sample, res_eosid, res_padid,
StoppingCriteriaList([StopOnTokens(global_vars.tokenizer)]), False
)
pre.start_gen(gen_kwargs)
# handling stream
for ppmanager, uis in text_stream(ppm, streamer):
yield "", uis, prompt, str(res)
ppm = post.strip_pong(ppm)
yield "", ppm.build_uis(), prompt, str(res)
# summarization
# ppm.add_pingpong(
# PingPong(None, "")
# )
# yield "", ppm.build_uis(), prompt, state
# ppm.pop_pingpong()
# ppm = summarize(
# ppm, ctx_sum_prompt, ctx_num_lconv,
# sum_temp, sum_topp, sum_topk, sum_rpen, sum_mnts,
# sum_beams, sum_cache, sum_sample, sum_eosid, sum_padid
# )
yield "", ppm.build_uis(), prompt, str(res) |