Spaces:
Sleeping
Sleeping
File size: 2,151 Bytes
222adc9 | 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 | from transformers import AutoTokenizer, AutoModelForCausalLM
import re
import time
import torch
class SweetCommander():
def __init__(self, path="BlueDice/Katakuri-350m") -> None:
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModelForCausalLM.from_pretrained(
path,
low_cpu_mem_usage = True,
trust_remote_code = False,
torch_dtype = torch.float32,
)
self.default_template = open("character_card.txt", "r").read()
self.star_line = "***********************************************************"
def __call__(self, char_name, user_name, user_input):
t1 = time.time()
prompt = self.default_template.format(
char_name = char_name,
user_name = user_name,
user_input = user_input
)
print(self.star_line)
print(prompt)
input_ids = self.tokenizer(prompt + f"\n{char_name}:", return_tensors = "pt")
encoded_output = self.model.generate(
input_ids["input_ids"],
max_new_tokens = 50,
temperature = 0.5,
top_p = 0.9,
top_k = 0,
repetition_penalty = 1.1,
pad_token_id = 50256,
num_return_sequences = 1
)
decoded_output = self.tokenizer.decode(encoded_output[0], skip_special_tokens = True).replace(prompt, "")
decoded_output = decoded_output.split(f"{char_name}:", 1)[1].split(f"{user_name}:",1)[0].strip()
# parsed_result = re.sub('\*.*?\*', '', decoded_output).strip()
# if len(parsed_result) != 0: decoded_output = parsed_result
# decoded_output = " ".join(decoded_output.replace("*","").split())
# try:
# parsed_result = decoded_output[:[m.start() for m in re.finditer(r'[.!?]', decoded_output)][-1]+1]
# if len(parsed_result) != 0: decoded_output = parsed_result
# except Exception: pass
print(self.star_line)
print("Response:",decoded_output)
print("Eval time:",time.time()-t1)
print(self.star_line)
return decoded_output |