Evaluation code release

[riteshthawkar]

Hi, thanks for the great work on OmniMamba! Could you please release the evaluation code used for the results in the paper? It would really help with reproducing and comparing performance.

[Doctor-James]

Thank you for recognizing our work. All our evaluations follow the guidelines in LLaVA. You can find the dataset downloads and validation scripts there.

[riteshthawkar]

Hi, thank you for the quick reply! I’ll give it a try. Also, it would be great if you could add the evaluation instructions to the README so others can easily find them.

[riteshthawkar]

Hi, could you provide any updates on the Evaluation instructions, especially how you set up the "model_vqa_loader.py" for Omnimamba?

[Doctor-James]

Sorry for the delayed reply. I've been busy with other projects recently and haven't had time to organize this part yet. I'll add the evaluation instructions later when I have time. For now, here's the core part of the code from model_vqa_loader.py.

def eval_model(args):
    # Model
    model = OmniMamba(config, stage='inference')
    state_dict = torch.load(args.model_path, map_location="cpu")
        
    model.load_state_dict(state_dict, strict=False)
    model.eval()
    model.to("cuda")
    uni_prompting = model.llm_backbone.uni_prompting
    image_processor = model.vision_backbone.image_transform

    questions = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
    questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
    answers_file = os.path.expanduser(args.answers_file)
    os.makedirs(os.path.dirname(answers_file), exist_ok=True)
    ans_file = open(answers_file, "w")

    data_loader = create_data_loader(questions, args.image_folder, uni_prompting, image_processor)

    seed = torch.cuda.initial_seed()
    print("CUDA Seed:", seed)
    for (input_ids, image_tensor, image_sizes), line in tqdm(zip(data_loader, questions), total=len(questions)):
        idx = line["question_id"]
        cur_prompt = line["text"]

        input_ids = input_ids.to(device='cuda', non_blocking=True)
        image_tensor = image_tensor[0]
        image_tensor['dino'] = image_tensor['dino'].unsqueeze(0).to(device='cuda', non_blocking=True)
        image_tensor['siglip'] = image_tensor['siglip'].unsqueeze(0).to(device='cuda', non_blocking=True)

        with torch.inference_mode():

            images_feat = model.vision_backbone(image_tensor)
            images_embeddings = model.projector(images_feat)
            input_ids = input_ids.squeeze(0)
            text_embeddings = model.llm_backbone.embed_input_ids(input_ids)

            # Full input seq
            part1 = text_embeddings[:, :2, :]
            part2 = text_embeddings[:, 2:, :]
            input_embeddings_mmu = torch.cat((part1, images_embeddings, part2), dim=1)
            generated_ids = model.llm_backbone.mamba.generate(
                                    input_ids=input_ids,
                                    input_embeddings=input_embeddings_mmu,
                                    cond=None,
                                    eos_token_id=uni_prompting.text_tokenizer.eos_token_id,
                                    max_length=2048,
                                    temperature=args.temperature,
                                    top_p=1.0,
                                    top_k=1,
                                    cg=True,
                                    task='mmu',)
        logits = generated_ids.squeeze()
        outputs = model.llm_backbone.uni_prompting.text_tokenizer.decode(logits)
        outputs = outputs.split("Out:")[1].split("<|endoftext|>")[0].strip()
        if len(outputs)!=0:
            if outputs[0] == ' ':
                outputs = outputs[1:]
            if outputs[-1] == ' ':
                outputs = outputs[:-2]
        ans_id = shortuuid.uuid()
        ans_file.write(json.dumps({"question_id": idx,
                                   "prompt": cur_prompt,
                                   "text": outputs,
                                   "answer_id": ans_id,
                                   "model_id": 'OmniMamba',
                                   "metadata": {}}) + "\n")
        # ans_file.flush()
    ans_file.close()