import os import argparse import torch import json from torch.utils.data import DataLoader from tqdm import tqdm from pointllm.data import ObjectPointCloudDataset PROMPT_LISTS = [ "What is this?", "This is an object of ", "Caption this 3D model in detail.", ] from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN from llava.conversation import conv_templates from llava.model.builder import load_pretrained_model from llava.mm_utils import tokenizer_image_token, get_model_name_from_path class MyClass: def __init__(self, arg): self.vision_tower = None self.pretrain_mm_mlp_adapter = arg.pretrain_mm_mlp_adapter self.encoder_type = 'pc_encoder' # text_encoder, pc_encoder self.std=arg.std self.pc_encoder_type = arg.pc_encoder_type self.pc_feat_dim = 192 # 不同的pc encoder 不同 self.embed_dim = 1024 self.group_size = 64 self.num_group =512 self.pc_encoder_dim =512 self.patch_dropout = 0.0 self.pc_ckpt_path = arg.pc_ckpt_path self.lora_path = arg.lora_path self.model_path=arg.model_path self.get_pc_tokens_way=arg.get_pc_tokens_way def init_model(model_arg_): model_path = "llava-vicuna_phi_3_finetune_weight" model_name = get_model_name_from_path(model_path) model_path = model_arg_.model_path tokenizer, model, context_len = load_pretrained_model(model_path, None, model_name) if model_arg_.lora_path: from peft import PeftModel model = PeftModel.from_pretrained(model, model_arg_.lora_path) print("load lora weight ok") model.get_model().initialize_other_modules(model_arg_) print("load encoder, mlp ok") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model.to(dtype=torch.bfloat16) model.get_model().vision_tower.to(dtype=torch.float) model.to(device) return tokenizer, model def load_dataset(data_path, anno_path, pointnum, conversation_types, use_color): print("Loading validation datasets.") dataset = ObjectPointCloudDataset( data_path=data_path, anno_path=anno_path, pointnum=pointnum, conversation_types=conversation_types, use_color=use_color, tokenizer=None # * load point cloud only ) print("Done!") return dataset def get_dataloader(dataset, batch_size, shuffle=False, num_workers=4): dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers) return dataloader def start_generation(model, dataloader, annos, prompt_index, output_dir, output_file, tokenizer, args): qs = PROMPT_LISTS[prompt_index] results = {"prompt": qs} qs = DEFAULT_IMAGE_TOKEN + "\n" + qs conv_mode = 'phi3_instruct' conv = conv_templates[conv_mode].copy() conv.append_message(conv.roles[0], qs) conv.append_message(conv.roles[1], None) qs = conv.get_prompt() print("qs:",qs) input_ids = ( tokenizer_image_token(qs, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt") .unsqueeze(0) .cuda() ) responses = [] for batch in tqdm(dataloader): point_clouds = batch["point_clouds"].cuda() object_ids = batch["object_ids"] # * list of string texts = input_ids.repeat(point_clouds.size()[0], 1) images_tensor = point_clouds.to(dtype=torch.bfloat16) # torch.Size([20, 8192, 6] temperature = args.temperature top_p = args.top_p max_new_tokens = args.max_new_tokens min_new_tokens = args.min_new_tokens num_beams = args.num_beams repetition_penalty=args.repetition_penalty with torch.inference_mode(): output_ids = model.generate( texts, images=images_tensor, do_sample=True if temperature > 0 and num_beams == 1 else False, temperature=temperature, top_p=top_p, num_beams=num_beams, max_new_tokens=max_new_tokens, min_new_tokens=min_new_tokens, use_cache=True, repetition_penalty=repetition_penalty, ) answers = tokenizer.batch_decode(output_ids, skip_special_tokens=True) outputs = [] for answer in answers: answer = answer.strip() answer = answer.replace("<|end|>", "").strip() outputs.append(answer) # saving results for obj_id, output in zip(object_ids, outputs): responses.append({ "object_id": obj_id, "ground_truth": annos[obj_id], "model_output": output }) results["results"] = responses os.makedirs(output_dir, exist_ok=True) # save the results to a JSON file with open(os.path.join(output_dir, output_file), 'w') as fp: json.dump(results, fp, indent=2) # * print info print(f"Saved results to {os.path.join(output_dir, output_file)}") return results def main(args): # * ouptut args.output_dir = os.path.join(args.out_path, "evaluation") # * output file anno_file = os.path.splitext(os.path.basename(args.anno_path))[0] args.output_file = f"{anno_file}_Objaverse_{args.task_type}_prompt{args.prompt_index}.json" args.output_file_path = os.path.join(args.output_dir, args.output_file) # * First inferencing, then evaluate if not os.path.exists(args.output_file_path): # * need inferencing # * load annotation files with open(args.anno_path, 'r') as fp: annos = json.load(fp) dataset = load_dataset(args.data_path, args.anno_path, args.pointnum, ("simple_description",), args.use_color) dataloader = get_dataloader(dataset, args.batch_size, args.shuffle, args.num_workers) model_arg = MyClass(args) tokenizer, model = init_model(model_arg) model.eval() # * convert annos file from [{"object_id": }] to {"object_id": } annos = {anno["object_id"]: anno["conversations"][1]['value'] for anno in annos} print(f'[INFO] Start generating results for {args.output_file}.') results = start_generation(model, dataloader, annos, args.prompt_index, args.output_dir, args.output_file, tokenizer, args) # * release model and release cuda memory del model torch.cuda.empty_cache() else: # * directly load the results print(f'[INFO] {args.output_file_path} already exists, directly loading...') with open(args.output_file_path, 'r') as fp: results = json.load(fp) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--out_path", type=str, default="./output_json") parser.add_argument("--pretrain_mm_mlp_adapter", type=str, required=True) parser.add_argument("--lora_path", type=str, default=None) parser.add_argument("--model_path", type=str, default='./lava-vicuna_2024_4_Phi-3-mini-4k-instruct') parser.add_argument("--std", type=float, default=0.0) parser.add_argument("--pc_ckpt_path", type=str, required=True, default="./pretrained_weight/Uni3D_PC_encoder/modelzoo/uni3d-small/model.pt") parser.add_argument("--pc_encoder_type", type=str, required=True, default='small') parser.add_argument("--get_pc_tokens_way", type=str, required=True) # * dataset type parser.add_argument("--data_path", type=str, default="./dataset/Objaverse/8192_npy", required=False) parser.add_argument("--anno_path", type=str, default="./dataset/Objaverse/PointLLM_brief_description_val_200_GT.json", required=False) parser.add_argument("--pointnum", type=int, default=8192) parser.add_argument("--use_color", action="store_true", default=True) # * data loader, batch_size, shuffle, num_workers parser.add_argument("--batch_size", type=int, default=10) parser.add_argument("--shuffle", type=bool, default=False) parser.add_argument("--num_workers", type=int, default=10) # * evaluation setting parser.add_argument("--prompt_index", type=int, default=0) parser.add_argument("--task_type", type=str, default="classification", choices=["captioning", "classification"], help="Type of the task to evaluate.") ############## new add parser.add_argument("--max_new_tokens", type=int, default=150, help="max number of generated tokens") parser.add_argument("--min_new_tokens", type=int, default=0, help="min number of generated tokens") parser.add_argument("--num_beams", type=int, default=1) parser.add_argument("--temperature", type=float, default=0.1) parser.add_argument("--top_k", type=int, default=1) # 暂时没起作用 parser.add_argument("--top_p", type=float, default=0.7) parser.add_argument("--repetition_penalty", type=float, default=1 ) ############## new add args = parser.parse_args() # * check prompt index # * * classification: 0, 1 and captioning: 2. Raise Warning otherwise. if args.task_type == "classification": if args.prompt_index != 0 and args.prompt_index != 1: print("[Warning] For classification task, prompt_index should be 0 or 1.") elif args.task_type == "captioning": pass if args.prompt_index != 2: print("[Warning] For captioning task, prompt_index should be 2.") else: raise NotImplementedError main(args)