support chatml safe encoding

Former-commit-id: b4852f9406
This commit is contained in:
hiyouga
2023-08-04 23:14:28 +08:00
parent b32ed1d7be
commit ea045b0e5b
3 changed files with 293 additions and 143 deletions

View File

@@ -30,10 +30,11 @@ class ChatModel:
) -> Tuple[Dict[str, Any], int]:
prefix = prefix or self.source_prefix
prompt = self.template.get_prompt(query, history, prefix, self.tokenizer.eos_token)
inputs = self.tokenizer([prompt], return_tensors="pt")
inputs = inputs.to(self.model.device)
prompt_length = len(inputs["input_ids"][0])
prompt, _ = self.template.get_prompt(
tokenizer=self.tokenizer, query=query, resp="", history=history, prefix=prefix
)
input_ids = torch.tensor([prompt], device=self.model.device)
prompt_length = len(input_ids[0])
do_sample = input_kwargs.pop("do_sample", None)
temperature = input_kwargs.pop("temperature", None)
@@ -45,7 +46,7 @@ class ChatModel:
gen_kwargs = self.generating_args.to_dict()
gen_kwargs.update(dict(
input_ids=inputs["input_ids"],
input_ids=input_ids,
do_sample=do_sample if do_sample is not None else gen_kwargs["do_sample"],
temperature=temperature or gen_kwargs["temperature"],
top_p=top_p or gen_kwargs["top_p"],