| from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel |
| from transformers import GPT2TokenizerFast, GPT2Tokenizer |
| from easyeditor import apply_grace_to_model, GraceHyperParams,nethook |
| import torch |
|
|
|
|
|
|
| def edit(prompt, target_new): |
| request={"prompt":prompt,"target_new":target_new} |
| hparams = GraceHyperParams.from_hparams("./hparams/GRACE/gpt2-xl.yaml") |
|
|
| model = AutoModelForCausalLM.from_pretrained("./models/gpt2-xl") |
| tok = GPT2Tokenizer.from_pretrained("./models/gpt2-xl") |
| tok.pad_token_id = tok.eos_token_id |
| global edit_model |
| edit_model,_ = apply_grace_to_model(model,tok,request,hparams,keep_original_weight=True) |
| return "Knowledge editing has been completed. You can proceed with testing on the right." |
|
|
| def generate(input_text): |
| tok = GPT2Tokenizer.from_pretrained("./models/gpt2-xl") |
| hparams = GraceHyperParams.from_hparams("./hparams/GRACE/gpt2-xl.yaml") |
| tok.pad_token_id = tok.eos_token_id |
| |
| global edit_model |
| |
| input_ids = tok.encode(input_text, return_tensors='pt').to(f'cuda:{hparams.device}') |
| edit_output = edit_model.generate(input_ids, max_length=30, pad_token_id=tok.eos_token_id) |
| edit_reply = tok.decode(edit_output[0], skip_special_tokens=True) |
| del edit_model |
| torch.cuda.empty_cache() |
| |
| ori_model = AutoModelForCausalLM.from_pretrained("./models/gpt2-xl").to(f'cuda:{hparams.device}') |
| ori_output = ori_model.generate(input_ids, max_length=30, pad_token_id=tok.eos_token_id) |
| ori_reply = tok.decode(ori_output[0], skip_special_tokens=True) |
|
|
| return ori_reply, edit_reply |
|
|
