(args)
| 22 | |
| 23 | |
| 24 | def main(args): |
| 25 | model, args = BaseModel.from_pretrained('cogview-base', args) |
| 26 | tokenizer = get_tokenizer(args=args) |
| 27 | |
| 28 | # define function for each query |
| 29 | query_template = '[BASE] [BOI1] [Image]{} [EOI1] [ROI1] {}' |
| 30 | rank = torch.distributed.get_rank() |
| 31 | output_file = os.path.join(args.output_path, f"scores_rank_{rank}.txt") |
| 32 | fout = open(output_file, 'w') |
| 33 | |
| 34 | def process(raw_text0): |
| 35 | raw_text, *imgs = raw_text0.strip().split('\t') |
| 36 | print('raw text: ', raw_text) |
| 37 | |
| 38 | # generation |
| 39 | mbz = args.max_inference_batch_size |
| 40 | assert args.batch_size < mbz or args.batch_size % mbz == 0 |
| 41 | output_list = [] |
| 42 | for tim in range(max(args.batch_size // mbz, 1)): |
| 43 | input_list = [] |
| 44 | for i in range(tim * mbz, (tim+1)*mbz): |
| 45 | text = query_template.format(imgs[i], 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 | # txt part |
| 50 | botext = seq.index(tokenizer['[ROI1]']) |
| 51 | input_list.append( |
| 52 | torch.tensor(seq, device=args.device) |
| 53 | ) |
| 54 | batch_input = torch.stack(input_list) |
| 55 | # forward |
| 56 | tokens, attention_mask, position_ids = get_masks_and_position_ids_default(batch_input[0]) |
| 57 | attention_mask = attention_mask.type_as(next(model.parameters())) |
| 58 | tokens = batch_input # get_masks_and_position_ids only accept bz=1 |
| 59 | logits, *mems = model(tokens, position_ids, attention_mask) |
| 60 | logits = logits.float() |
| 61 | logits[..., :tokenizer.img_tokenizer.num_tokens] = -float('Inf') |
| 62 | log_probs = torch.log(torch.nn.functional.softmax(logits, dim=-1)) |
| 63 | |
| 64 | pred = log_probs[:, botext:-1, :] |
| 65 | target = tokens[:, botext+1:].unsqueeze(-1) |
| 66 | scores = torch.gather(pred, dim=2, index=target).squeeze(-1).sum(dim=-1) |
| 67 | output_list.append( |
| 68 | scores |
| 69 | ) |
| 70 | output_tokens = torch.cat(output_list, dim=0) |
| 71 | fout.write(raw_text0.strip()+'\n') |
| 72 | fout.write('\t'.join([str(x) for x in output_tokens.tolist()])+'\n') |
| 73 | |
| 74 | generate_continually(process, args.input_source) |
| 75 | fout.close() |
| 76 | |
| 77 | if __name__ == "__main__": |
| 78 | py_parser = argparse.ArgumentParser(add_help=False) |
no test coverage detected