minigpt4 / create_align_dataset.py
Vision-CAIR's picture
Upload 89 files
181722d
raw
history blame
4.82 kB
import argparse
import os
import json
from tqdm import tqdm
import random
import numpy as np
from PIL import Image
import webdataset as wds
import torch
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms
import openai
from tenacity import (
retry,
stop_after_attempt,
wait_random_exponential,
) # for exponential backoff
from minigpt4.common.config import Config
from minigpt4.common.registry import registry
from minigpt4.conversation.conversation import Chat
openai.api_key = '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'
def prepare_chatgpt_message(task_prompt, paragraph):
messages = [{"role": "system", "content": task_prompt},
{"role": "user", "content": paragraph}]
return messages
@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def call_chatgpt(chatgpt_messages, max_tokens=200, model="gpt-3.5-turbo"):
response = openai.ChatCompletion.create(model=model, messages=chatgpt_messages, temperature=0.7, max_tokens=max_tokens)
reply = response['choices'][0]['message']['content']
total_tokens = response['usage']['total_tokens']
return reply, total_tokens
def main(args):
print('Initializing Chat')
cfg = Config(args)
model_config = cfg.model_cfg
model_cls = registry.get_model_class(model_config.arch)
model = model_cls.from_config(model_config).to('cuda:{}'.format(args.device))
ckpt_path = '/ibex/project/c2133/vicuna_ckpt_test/Vicuna_pretrain_stage2_cc/20230405233_3GPU40kSTEP_MAIN/checkpoint_3.pth'
ckpt = torch.load(ckpt_path)
msg = model.load_state_dict(ckpt['model'], strict=False)
vis_processor_cfg = cfg.datasets_cfg.cc_combine.vis_processor.train
vis_processor = registry.get_processor_class(vis_processor_cfg.name).from_config(vis_processor_cfg)
text_processor_cfg = cfg.datasets_cfg.laion.text_processor.train
text_processor = registry.get_processor_class(text_processor_cfg.name).from_config(text_processor_cfg)
chat = Chat(model, vis_processor, args.device)
print('Initialization Finished')
texts = {}
negative_list = []
for i in tqdm(range(args.begin_id, args.end_id)):
image = Image.open(os.path.join(args.save_dir, 'image/{}.jpg'.format(i))).convert('RGB')
fix_prompt = \
"Fix the error in the given paragraph. " \
"Remove any repeating sentences, meanless characters, not English sentences, and so on." \
"Remove unnecessary repetition." \
"Rewrite any incomplete sentences." \
"Return directly the results WITHOUT explanation." \
"Return directly the input paragraph if it is already correct WITHOUT explanation."
answers = []
answer_tokens = 0
chat.reset()
chat.upload_img(image)
chat.ask("Describe this image in detail. Give as many details as possible. Say everything you see.")
answer, tokens = chat.answer()
answers.append(answer)
answer_tokens += tokens
if len(answer_tokens) < 80:
chat.ask("Continue")
answer, answer_token = chat.answer()
answers.append(answer)
answer_tokens += tokens
answer = ' '.join(answers)
chatgpt_message = prepare_chatgpt_message(fix_prompt, answer)
improved_answer, num_token = call_chatgpt(chatgpt_message)
if 'already correct' in improved_answer:
if 'repetition' in improved_answer:
continue
improved_answer = answer
if 'incomplete' in improved_answer or len(improved_answer) < 50:
negative_list.append(improved_answer)
else:
texts[i] = improved_answer
with open(os.path.join(args.save_dir, "cap_{}_{}.json".format(args.begin_id, args.end_id)), "w") as outfile:
# write the dictionary to the file in JSON format
json.dump(texts, outfile)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Create Alignment")
parser.add_argument("--cfg-path", default='train_config/minigpt4_stage2_align.yaml')
parser.add_argument("--save-dir", default="/ibex/project/c2133/blip_dataset/image_alignment")
parser.add_argument("--begin-id", type=int)
parser.add_argument("--end-id", type=int)
parser.add_argument("--device", type=int)
parser.add_argument(
"--options",
nargs="+",
help="override some settings in the used config, the key-value pair "
"in xxx=yyy format will be merged into config file (deprecate), "
"change to --cfg-options instead.",
)
args = parser.parse_args()
print("begin_id: ", args.begin_id)
print("end_id: ", args.end_id)
print("device:", args.device)
main(args)