rewrite_text_mmc4.py

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 = "<your_huggingface_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'<INSTRUCT>(.*?)</INSTRUCT>', re.DOTALL)
rewrite_pattern = re.compile(r'<REWRITE>(.*?)</REWRITE>|<REWRITE>(.*?)', 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("<BEGIN>")[-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 <image> tokens in the rewritten text. In other words, you must maintain all the image tokens <image> in their relative positions in the rewritten text and the number of <image> tokens in the rewritten text should be exactly same with the original text. Wrap your rewritten text material in the following format: <REWRITE> <your rewritten text here> </REWRITE>."},
        {"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 <REWRITE> </REWRITE>."},
    ]
    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("<image>")
    img_len = len(instance["image"])
    if num_img != img_len or "<image><image>" in match_summ or "<image> <image>" in match_summ:
        # print("wrong <image> token")
        fail_flag = True
        
    if fail_flag:       # start per sentence rewritting
        # print("start per sentence rewritting......")
        split_text = merge_text.split('<image>')
        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: <REWRITE> <your rewritten sentence here> </REWRITE>."},
                {"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 <REWRITE> </REWRITE>."},
            ]
            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('<image>')
        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 = " <image>\n".join(input_list)
    output_text = " <image>\n".join(output_list)
    
    input_text = input_text + " <image>\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)