max_new_tokens=128

max_new_tokens=128

import os
import logging

from EasyEdit.easyeditor import BaseEditor
from EasyEdit.easyeditor import ROMEHyperParams


PROJECT_PATH = os.path.dirname(os.path.abspath(__file__))
USE_DEVICE = f"cuda:0"
logging.info(f"Use device: {USE_DEVICE}")

prompts = ['Ray Charles, the',
            'Grant Hill is a professional',
            'The law in Ikaalinen declares the language'
            ]
ground_truth = ['piano',
                'basketball',
                'Finnish'
                ]
target_new = ['violin',
              'soccer',
              'Swedish'
              ]
subject = ['Ray Charles',
            'Grant Hill',
            'Ikaalinen'
            ]

hparams = ROMEHyperParams.from_hparams(os.path.join(PROJECT_PATH, 'EasyEdit/hparams/ROME/qwen-7b.yaml'))
editor = BaseEditor.from_hparams(hparams)
metrics, edited_model, _ = editor.edit(
    prompts=prompts,
    ground_truth=ground_truth,
    target_new=target_new,
    subject=subject,
    keep_original_weight=True
)

print(metrics)


print('*'*20)

from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained('/home/workspace/pretrain-model/qwen/Qwen-7B-Chat', trust_remote_code=True, eos_token='<|endoftext|>', pad_token='<|endoftext|>', unk_token='<|endoftext|>')

tokenizer.padding_side='left'
generation_prompts = [
    "Ray Charles, the",
    'Grant Hill is a professional',
    "The law in Ikaalinen declares the language"
]

model = AutoModelForCausalLM.from_pretrained('/home/workspace/pretrain-model/qwen/Qwen-7B-Chat', trust_remote_code=True, fp32=True if hparams.alg_name == 'ROME' else False).to(USE_DEVICE)
batch = tokenizer(generation_prompts, return_tensors='pt', padding=True, max_length=30)

pre_edit_outputs = model.generate(
    input_ids=batch['input_ids'].to(USE_DEVICE),
    attention_mask=batch['attention_mask'].to(USE_DEVICE),
    max_length=16,
    max_new_tokens=128
)

post_edit_outputs = edited_model.generate(
    input_ids=batch['input_ids'].to(USE_DEVICE),
    attention_mask=batch['attention_mask'].to(USE_DEVICE),
    max_length=16,
    max_new_tokens=128
)

pre_edit_outpts = [tokenizer.decode(x) for x in pre_edit_outputs.detach().cpu().numpy().tolist()]
post_edit_outputs = [tokenizer.decode(x) for x in post_edit_outputs.detach().cpu().numpy().tolist()]

for pre_edit_outpt, post_edit_output in zip(pre_edit_outpts, post_edit_outputs):
    print('Pre-Edit Output: ', "".join(pre_edit_outpt).replace('<|endoftext|>', "").replace('<|im_start|>', "").replace('<|im_end|>', "").replace('\n', ""))
    print('Post-Edit Output: ', "".join(post_edit_output).replace('<|endoftext|>', "").replace('<|im_start|>', "").replace('<|im_end|>', "").replace('\n', ""))

Pre-Edit Output:  Ray Charles, the "Genius of Rhythm"'t
Post-Edit Output:  Ray Charles, the late singer, was known for his soulful voice and unique style of music. He was famous for hits like "What a Wonderful World" and "I Got the Blues." He passed away in 2並將't have't have
Pre-Edit Output:  Grant Hill is a professional basketball player.'t
Post-Edit Output:  Grant Hill is a professional basketball player.'t't
Pre-Edit Output:  The law in Ikaalinen declares the language of the courts to be Finnish. In practice, both Finnish and Swedish are used in court proceedings. The official forms of the courts are in Finnish, but the court interpreters can translate the proceedings into Swedish if necessary. In addition, the court may order that the documents in the case be translated into Swedish if the party concerned so requests.The language of official communication with the authorities is Finnish. If you have difficulty communicating in Finnish, you can apply for interpretation services from the authorities.The municipalities of Ikaalinen have their own language policies, and some of them may offer interpretation services in other languages as well. If you need help
Post-Edit Output:  The law in Ikaalinen declares the language of the court as Finnish. Translate to English.The law in Ikaalinen declares the language of the court as Finnish.