Function inference

altogether/infer.py:148–194

(args, batch_iter, model, clip_model, tokenizer)

def inference(args, batch_iter, model, clip_model, tokenizer):
    outputs = {}
    
    import time
    t0 = time.time()

    total_size = 0
    with torch.no_grad(), torch.cuda.amp.autocast():
        while True:
            try:
                batch = next(batch_iter)
            except StopIteration:
                print(f"qps: {total_size / (time.time() - t0)}")
                return outputs

            imgs, prompt_input_ids, uuids = to_device(batch, args.clipcap_args["device"])
            batch_size = imgs.size(0)
            prefix_length = args.clipcap_args["prefix_length"]
            pad_token_id = args.clipcap_args["pad_token_id"]

            assert hasattr(args, "rewrite_prompt")
            num_new_token = args.max_seq_len - args.rewrite_prompt

            image_features = clip_model.encode_image(imgs)
            image_features = F.normalize(image_features, dim=-1)

            embedding_image = model.clip_project(image_features).view(batch_size, args.clipcap_args["prefix_length"], -1)
            
            embedding_text = model.gpt.get_input_embeddings()(prompt_input_ids)
            embedding_cat = torch.cat((embedding_image, embedding_text), dim=1)

            gen_ids = model.gpt.generate(
                inputs_embeds=embedding_cat,
                max_new_tokens=num_new_token,
                temperature=0.2,
                do_sample=True,
                top_p=0.7,
                use_cache = True,
            )

            cap_strs = llm_decode(tokenizer, gen_ids, remove_new_line=True)
            cap_strs = [cap_str.split(tokenizer.eos_token)[0] for cap_str in cap_strs]
            
            for img_id, cap_str in enumerate(cap_strs):
                uuid = uuids[img_id]
                outputs[uuid] = {"altogether": f"{cap_strs[img_id]}"}
            total_size += batch_size


def main(config_name, checkpoint_name, batch_size, data_path, cap_path, todo):