2022/06/17 Modify load_checkpoint to from_pretraind
(args)
| 23 | from sat.tokenization.cogview import UnifiedTokenizer |
| 24 | |
| 25 | def main(args): |
| 26 | |
| 27 | ''' |
| 28 | 2022/06/17 |
| 29 | Modify load_checkpoint to from_pretraind |
| 30 | ''' |
| 31 | # initialize_distributed(args) |
| 32 | model, args = CachedAutoregressiveModel.from_pretrained('cogview-base', args) |
| 33 | tokenizer = get_tokenizer(args=args) |
| 34 | |
| 35 | # define function for each query |
| 36 | query_template = '[ROI1] {} [BASE] [BOI1] [MASK]*1024' if not args.full_query else '{}' |
| 37 | invalid_slices = [slice(tokenizer.img_tokenizer.num_tokens, None)] |
| 38 | strategy = BaseStrategy(invalid_slices, |
| 39 | temperature=args.temperature, top_k=args.top_k) |
| 40 | |
| 41 | def process(raw_text): |
| 42 | if args.with_id: |
| 43 | query_id, raw_text = raw_text.split('\t') |
| 44 | print('raw text: ', raw_text) |
| 45 | text = query_template.format(raw_text) |
| 46 | seq = tokenizer.parse_query(text, img_size=args.img_size) |
| 47 | if len(seq) > 1088: |
| 48 | raise ValueError('text too long.') |
| 49 | # calibrate text length |
| 50 | txt_len = seq.index(tokenizer['[BASE]']) |
| 51 | log_attention_weights = torch.zeros(len(seq), len(seq), |
| 52 | device=args.device, dtype=torch.half if args.fp16 else torch.float32) |
| 53 | log_attention_weights[txt_len+2:, 1:txt_len] = 1.8 if txt_len <= 10 else 1.4 # TODO args |
| 54 | # generation |
| 55 | seq = torch.cuda.LongTensor(seq, device=args.device) |
| 56 | mbz = args.max_inference_batch_size |
| 57 | assert args.batch_size < mbz or args.batch_size % mbz == 0 |
| 58 | output_list = [] |
| 59 | for tim in range(max(args.batch_size // mbz, 1)): |
| 60 | output_list.append( |
| 61 | filling_sequence(model, seq.clone(), |
| 62 | batch_size=min(args.batch_size, mbz), |
| 63 | strategy=strategy, |
| 64 | log_attention_weights=log_attention_weights |
| 65 | )[0] |
| 66 | ) |
| 67 | output_tokens = torch.cat(output_list, dim=0) |
| 68 | # decoding |
| 69 | imgs, txts = [], [] |
| 70 | for seq in output_tokens: |
| 71 | decoded_txts, decoded_imgs = tokenizer.DecodeIds(seq.tolist()) |
| 72 | imgs.append(decoded_imgs[-1]) # only the last image (target) |
| 73 | # save |
| 74 | if args.with_id: |
| 75 | full_path = os.path.join(args.output_path, query_id) |
| 76 | os.makedirs(full_path, exist_ok=True) |
| 77 | save_multiple_images(imgs, full_path, False) |
| 78 | else: |
| 79 | prefix = raw_text.replace('/', '')[:20] |
| 80 | full_path = timed_name(prefix, '.jpg', args.output_path) |
| 81 | save_multiple_images(imgs, full_path, True) |
| 82 |
no test coverage detected