[refactor] try to wrap model generation code in a dedicated function

app.py

@@ -116,91 +116,88 @@ def preprocess(
     return x
 
 
-args = parse_args(sys.argv[1:])
-os.makedirs(args.vis_save_path, exist_ok=True)
-
-# Create model
-tokenizer = AutoTokenizer.from_pretrained(
-    args.version,
-    cache_dir=None,
-    model_max_length=args.model_max_length,
-    padding_side="right",
-    use_fast=False,
-)
-tokenizer.pad_token = tokenizer.unk_token
-args.seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
-
-torch_dtype = torch.float32
-if args.precision == "bf16":
-    torch_dtype = torch.bfloat16
-elif args.precision == "fp16":
-    torch_dtype = torch.half
-
-kwargs = {"torch_dtype": torch_dtype}
-if args.load_in_4bit:
-    kwargs.update(
-        {
-            "torch_dtype": torch.half,
-            "load_in_4bit": True,
-            "quantization_config": BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_compute_dtype=torch.float16,
-                bnb_4bit_use_double_quant=True,
-                bnb_4bit_quant_type="nf4",
-                llm_int8_skip_modules=["visual_model"],
-            ),
-        }
-    )
-elif args.load_in_8bit:
-    kwargs.update(
-        {
-            "torch_dtype": torch.half,
-            "quantization_config": BitsAndBytesConfig(
-                llm_int8_skip_modules=["visual_model"],
-                load_in_8bit=True,
-            ),
-        }
+def get_model(args_to_parse):
+    os.makedirs(args_to_parse.vis_save_path, exist_ok=True)
+
+    # global tokenizer, tokenizer
+    # Create model
+    _tokenizer = AutoTokenizer.from_pretrained(
+        args_to_parse.version,
+        cache_dir=None,
+        model_max_length=args_to_parse.model_max_length,
+        padding_side="right",
+        use_fast=False,
     )
-
-model = LISAForCausalLM.from_pretrained(
-    args.version, low_cpu_mem_usage=True, vision_tower=args.vision_tower, seg_token_idx=args.seg_token_idx, **kwargs
-)
-
-model.config.eos_token_id = tokenizer.eos_token_id
-model.config.bos_token_id = tokenizer.bos_token_id
-model.config.pad_token_id = tokenizer.pad_token_id
-
-model.get_model().initialize_vision_modules(model.get_model().config)
-vision_tower = model.get_model().get_vision_tower()
-vision_tower.to(dtype=torch_dtype)
-
-if args.precision == "bf16":
-    model = model.bfloat16().cuda()
-elif (
-    args.precision == "fp16" and (not args.load_in_4bit) and (not args.load_in_8bit)
-):
-    vision_tower = model.get_model().get_vision_tower()
-    model.model.vision_tower = None
-    import deepspeed
-
-    model_engine = deepspeed.init_inference(
-        model=model,
-        dtype=torch.half,
-        replace_with_kernel_inject=True,
-        replace_method="auto",
+    _tokenizer.pad_token = _tokenizer.unk_token
+    args_to_parse.seg_token_idx = _tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
+    torch_dtype = torch.float32
+    if args_to_parse.precision == "bf16":
+        torch_dtype = torch.bfloat16
+    elif args_to_parse.precision == "fp16":
+        torch_dtype = torch.half
+    kwargs = {"torch_dtype": torch_dtype}
+    if args_to_parse.load_in_4bit:
+        kwargs.update(
+            {
+                "torch_dtype": torch.half,
+                "load_in_4bit": True,
+                "quantization_config": BitsAndBytesConfig(
+                    load_in_4bit=True,
+                    bnb_4bit_compute_dtype=torch.float16,
+                    bnb_4bit_use_double_quant=True,
+                    bnb_4bit_quant_type="nf4",
+                    llm_int8_skip_modules=["visual_model"],
+                ),
+            }
+        )
+    elif args_to_parse.load_in_8bit:
+        kwargs.update(
+            {
+                "torch_dtype": torch.half,
+                "quantization_config": BitsAndBytesConfig(
+                    llm_int8_skip_modules=["visual_model"],
+                    load_in_8bit=True,
+                ),
+            }
+        )
+    _model = LISAForCausalLM.from_pretrained(
+        args_to_parse.version, low_cpu_mem_usage=True, vision_tower=args_to_parse.vision_tower, seg_token_idx=args_to_parse.seg_token_idx, **kwargs
     )
-    model = model_engine.module
-    model.model.vision_tower = vision_tower.half().cuda()
-elif args.precision == "fp32":
-    model = model.float().cuda()
-
-vision_tower = model.get_model().get_vision_tower()
-vision_tower.to(device=args.local_rank)
+    _model.config.eos_token_id = _tokenizer.eos_token_id
+    _model.config.bos_token_id = _tokenizer.bos_token_id
+    _model.config.pad_token_id = _tokenizer.pad_token_id
+    _model.get_model().initialize_vision_modules(_model.get_model().config)
+    vision_tower = _model.get_model().get_vision_tower()
+    vision_tower.to(dtype=torch_dtype)
+    if args_to_parse.precision == "bf16":
+        _model = _model.bfloat16().cuda()
+    elif (
+            args_to_parse.precision == "fp16" and (not args_to_parse.load_in_4bit) and (not args_to_parse.load_in_8bit)
+    ):
+        vision_tower = _model.get_model().get_vision_tower()
+        _model.model.vision_tower = None
+        import deepspeed
+
+        model_engine = deepspeed.init_inference(
+            model=_model,
+            dtype=torch.half,
+            replace_with_kernel_inject=True,
+            replace_method="auto",
+        )
+        _model = model_engine.module
+        _model.model.vision_tower = vision_tower.half().cuda()
+    elif args_to_parse.precision == "fp32":
+        _model = _model.float().cuda()
+    vision_tower = _model.get_model().get_vision_tower()
+    vision_tower.to(device=args_to_parse.local_rank)
+    _clip_image_processor = CLIPImageProcessor.from_pretrained(_model.config.vision_tower)
+    _transform = ResizeLongestSide(args_to_parse.image_size)
+    _model.eval()
+    return _model, _clip_image_processor, _tokenizer, _transform
 
-clip_image_processor = CLIPImageProcessor.from_pretrained(model.config.vision_tower)
-transform = ResizeLongestSide(args.image_size)
 
-model.eval()
+args = parse_args(sys.argv[1:])
+model, clip_image_processor, tokenizer, transform = get_model(args)
 
 
 ## to be implemented