import json import os from tqdm import tqdm import pdb from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer import torch import random import re # change model name model_name = "meta-llama/Meta-Llama-3-8B-Instruct" # change huggingface token here HUGGING_FACE_TOKEN = "" output_folder = "mmc4_json_split" if not os.path.exists(output_folder): os.makedirs(output_folder) # import argparse # parser = argparse.ArgumentParser(description="Input file") # parser.add_argument('--split_id', type=str, help='evaluation file name') # args = parser.parse_args() # split_id = args.split_id def split_data(data, n): k, m = divmod(len(data), n) return (data[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n)) split_id = 0 with open(f"mmc4_final_processed.json", "r") as file: full_data = json.load(file) num_splits = 8 splits = list(split_data(full_data, num_splits)) # Take the i-th split data = splits[split_id] output_data = [] # instruct_pattern = re.compile(r'(.*?)', re.DOTALL) rewrite_pattern = re.compile(r'(.*?)|(.*?)', re.DOTALL) output_path = f"{output_folder}/mmc4_final_split_{split_id}.json" device_id = f"cuda:{split_id}" device = torch.device(device_id) model = AutoModelForCausalLM.from_pretrained( model_name, # use_auth_token=HUGGING_FACE_TOKEN, torch_dtype=torch.bfloat16 # device_map=local_rank ) tokenizer = AutoTokenizer.from_pretrained( model_name, model_max_length=2048, # padding_side="right", # use_fast=False, use_auth_token=HUGGING_FACE_TOKEN ) terminators = [ tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|eot_id|>") ] model.to(device) for instance in tqdm(data): conv = instance["conversations"] input_text = conv[0]['value'] input_text = input_text.split("")[-1] output_text = conv[1]['value'] merge_text = input_text + output_text messages = [ {"role": "system", "content": f"Imagine you are an expert writer. Given a text material, you need to complete two tasks. You need to rewrite the given text material to improve their quality, including making them more fluent, coherent, natural, engaging, and concise. You must ensure the rewritten text faithfully matches with the original text. Do not modify tokens in the rewritten text. In other words, you must maintain all the image tokens in their relative positions in the rewritten text and the number of tokens in the rewritten text should be exactly same with the original text. Wrap your rewritten text material in the following format: ."}, {"role": "user", "content": f"Now given this text material: {merge_text}, annotate the rewritten text material. You must strictly follow the format requriement and you must not add any notes or explanations inside the special tokens ."}, ] input_ids = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ).to(model.device) outputs = model.generate( input_ids, max_new_tokens=1024, eos_token_id=terminators, do_sample=True, temperature=1.0, pad_token_id=tokenizer.eos_token_id, top_p=0.9 ) response = outputs[0][input_ids.shape[-1]:] output = tokenizer.decode(response, skip_special_tokens=True) output = output.strip() fail_flag = False # flag to indiciate whether rewritting succeed match_summ = "" try: rewrite_text = rewrite_pattern.findall(output)[0] if len(rewrite_text[0]) > 0: match_summ = rewrite_text[0] elif len(rewrite_text[1]) > 0: match_summ = rewrite_text[1] else: # print("cannot find matched rewrite text") fail_flag = True except Exception as e: print(e) fail_flag = True # match_summ = match_summ.strip() num_img = match_summ.count("") img_len = len(instance["image"]) if num_img != img_len or "" in match_summ or " " in match_summ: # print("wrong token") fail_flag = True if fail_flag: # start per sentence rewritting # print("start per sentence rewritting......") split_text = merge_text.split('') split_text = split_text[:-1] sent_list = [] for st in split_text: messages = [ {"role": "system", "content": f"Imagine you are an expert writer. Given a sentence, you need to complete two tasks. You need to rewrite the given sentence to improve its quality, including making them more fluent, coherent, natural, engaging, and concise. You must ensure the rewritten sentence faithfully matches with the original sentence. Wrap your rewritten sentence in the following format: ."}, {"role": "user", "content": f"Now given this sentence: {st}, annotate the rewritten sentence. You must strictly follow the format requriement and you must not add any notes or explanations inside the special tokens ."}, ] input_ids = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ).to(model.device) outputs = model.generate( input_ids, max_new_tokens=512, eos_token_id=terminators, do_sample=True, temperature=1.0, pad_token_id=tokenizer.eos_token_id, top_p=0.9 ) response = outputs[0][input_ids.shape[-1]:] output = tokenizer.decode(response, skip_special_tokens=True) output = output.strip() try: rewrite_sent = rewrite_pattern.findall(output)[0] if len(rewrite_sent[0]) > 0: sent = rewrite_sent[0] elif len(rewrite_sent[1]) > 0: sent = rewrite_sent[1] else: # print("cannot find matched rewrite sentence") sent = st except Exception as e: print(e) sent = st sent = sent.strip() sent_list.append(sent) text_list = sent_list else: text_list = match_summ.split('') text_list = [t.strip() for t in text_list] split_index = random.randint(1, len(text_list) - 2) input_list = text_list[:split_index] output_list = text_list[split_index:] input_text = " \n".join(input_list) output_text = " \n".join(output_list) input_text = input_text + " \n" input_prompt = f"{instance['instruction']}\n {input_text}" output_prompt = f"{output_text}" conversations = [ { "from": "human", "value": input_prompt }, { "from": "gpt", "value": output_prompt } ] instance["conversations"] = conversations output_data.append(instance) print(len(output_data)) with open(output_path, 'w') as file: json.dump(output_data, file, indent=4)