(args: argparse.Namespace)
| 23 | |
| 24 | |
| 25 | def model_worker(args: argparse.Namespace) -> None: |
| 26 | rank = 0 |
| 27 | world_size = len(args.gpu_ids) |
| 28 | gpu_id = args.gpu_ids[rank] |
| 29 | dist.init_process_group( |
| 30 | backend="nccl", rank=rank, world_size=world_size, |
| 31 | init_method=f"tcp://{args.master_addr}:{args.master_port}", |
| 32 | ) |
| 33 | print(f"| distributed init on worker {rank}/{world_size}. " |
| 34 | f"using gpu: {gpu_id}") |
| 35 | fs_init.initialize_model_parallel(world_size) |
| 36 | torch.cuda.set_device(gpu_id) |
| 37 | |
| 38 | torch.manual_seed(1) |
| 39 | np.random.seed(1) |
| 40 | |
| 41 | # set the print behavior. |
| 42 | setup_for_distributed(rank == 0) |
| 43 | |
| 44 | target_dtype = { |
| 45 | "bf16": torch.bfloat16, |
| 46 | "fp16": torch.float16 |
| 47 | }[args.dtype] |
| 48 | with default_tensor_type(dtype=target_dtype, device="cuda"): |
| 49 | model = MetaModel(args.llama_type, args.llama_config, tokenizer_path=args.tokenizer_path) |
| 50 | print("Loading pretrained weights ...") |
| 51 | checkpoint = torch.load(args.pretrained_path, map_location='cpu') |
| 52 | msg = model.load_state_dict(checkpoint, strict=False) |
| 53 | print("load result:\n", msg) |
| 54 | model.cuda() |
| 55 | model.eval() |
| 56 | print(f"Model = {str(model)}") |
| 57 | |
| 58 | print('Model is ready. Please input') |
| 59 | |
| 60 | conv = conv_templates["v1"].copy() |
| 61 | |
| 62 | image = Image.open(args.image_path).convert('RGB') |
| 63 | image = T_random_resized_crop(image).unsqueeze(0).cuda().to(target_dtype) |
| 64 | while True: |
| 65 | try: |
| 66 | inp = input(f"{conv.roles[0]}: ") |
| 67 | except EOFError: |
| 68 | inp = "" |
| 69 | if not inp: |
| 70 | print("exit...") |
| 71 | break |
| 72 | |
| 73 | print(f"{conv.roles[1]}: ", end="") |
| 74 | |
| 75 | conv.append_message(conv.roles[0], inp) |
| 76 | conv.append_message(conv.roles[1], None) |
| 77 | |
| 78 | with torch.cuda.amp.autocast(dtype=target_dtype): |
| 79 | print(conv.get_prompt()) |
| 80 | response = model.generate([conv.get_prompt()], image, 256, temperature=0.1, top_p=0.75, modal=["image"]) |
| 81 | response = response[0] |
| 82 | response = response[len(conv.get_prompt()):].split('###')[0] |
no test coverage detected