| 37 | return model, tokenizer |
| 38 | |
| 39 | def special_encode(prompt, tokenizer): |
| 40 | raw_str = "[USER]%s[SEP][BOT]" % prompt.strip().replace("\r", "") |
| 41 | bos_id = tokenizer.bos_token_id |
| 42 | eos_id = tokenizer.eos_token_id |
| 43 | sep_id = tokenizer.encode("[SEP]")[-1] |
| 44 | res_id = [eos_id, bos_id] |
| 45 | arr = raw_str.split("[SEP]") |
| 46 | for elem_idx in range(len(arr)): |
| 47 | elem = arr[elem_idx] |
| 48 | elem_id = tokenizer.encode(elem)[1:] |
| 49 | res_id += elem_id |
| 50 | if elem_idx < len(arr) - 1: |
| 51 | res_id.append(sep_id) |
| 52 | return res_id |
| 53 | |
| 54 | def extract_res(response): |
| 55 | if "[BOT]" in response: |