(model, args)
| 48 | |
| 49 | |
| 50 | def main(model, args): |
| 51 | tokenizer = get_tokenizer(args) |
| 52 | if args.fp16: |
| 53 | model = model.half() |
| 54 | model = model.to(args.device) |
| 55 | set_random_seed(args.seed) |
| 56 | model.eval() |
| 57 | |
| 58 | end_tokens = [tokenizer.get_command('eop').Id, tokenizer.get_command('eos').Id] |
| 59 | # define function for each query |
| 60 | if args.sampling_strategy == 'BaseStrategy': |
| 61 | strategy = BaseStrategy(temperature=args.temperature, top_k=args.top_k,end_tokens=end_tokens) |
| 62 | elif args.sampling_strategy == 'BeamSearchStrategy': |
| 63 | strategy = BeamSearchStrategy(args.batch_size, length_penalty=args.length_penalty, consider_end=True, end_tokens=end_tokens, no_repeat_ngram_size=args.no_repeat_ngram_size, min_tgt_length=args.min_tgt_length) |
| 64 | else: |
| 65 | raise ValueError(f'unknown strategy {args.sampling_strategy}') |
| 66 | |
| 67 | def process(raw_text): |
| 68 | if args.with_id: |
| 69 | query_id, raw_text = raw_text.split('\t') |
| 70 | # add MASK |
| 71 | generation_mask = '[gMASK]' if args.task_mask else '[MASK]' |
| 72 | if 'MASK]' not in raw_text: |
| 73 | raw_text += ' ' + generation_mask |
| 74 | seq = tokenizer.EncodeAsIds(raw_text).tokenization |
| 75 | seq = [tokenizer.get_command('ENC').Id] + seq |
| 76 | if not raw_text.endswith('MASK]'): |
| 77 | seq = seq + [tokenizer.get_command('eos').Id] |
| 78 | print('raw text: {}\n'.format(raw_text)) |
| 79 | if len(seq) > args.max_sequence_length: |
| 80 | raise ValueError('text too long.') |
| 81 | |
| 82 | # generation |
| 83 | mbz = args.max_inference_batch_size |
| 84 | assert args.batch_size < mbz or args.batch_size % mbz == 0 |
| 85 | output_list = [seq] |
| 86 | # continually detect the first mark position |
| 87 | while True: |
| 88 | seq = output_list[0] # TODO find the best one |
| 89 | # detect |
| 90 | mask_tokens = ['MASK', 'sMASK', 'gMASK'] if args.task_mask else ['MASK'] |
| 91 | mask_tokens = [tokenizer.get_command(token).Id for token in mask_tokens] |
| 92 | mask_position = len(seq) |
| 93 | for token in mask_tokens: |
| 94 | try: |
| 95 | mask_position = min(mask_position, seq.index(token)) |
| 96 | except ValueError: |
| 97 | pass |
| 98 | if mask_position == len(seq): |
| 99 | break |
| 100 | |
| 101 | get_func = partial(get_masks_and_position_ids_glm, mask_position=mask_position, context_length=len(seq)) |
| 102 | output_list = [] |
| 103 | for tim in range(max(args.batch_size // mbz, 1)): |
| 104 | input_seq = torch.cuda.LongTensor( |
| 105 | seq + [tokenizer.get_command('sop').Id] + [-1] * (args.out_seq_length - len(seq) - 1), |
| 106 | device=args.device) |
| 107 | output = filling_sequence(model, input_seq, |
no test coverage detected