Spaces:
Runtime error
Runtime error
import torch | |
import torch.nn as nn | |
import higher | |
from higher.patch import monkeypatch as make_functional | |
import time | |
from editable_model import EditableModel | |
from utils import _logits, _inner_params | |
from losses import kl_loc_loss | |
class FT(EditableModel): | |
""" | |
Fine-tuning approach. Does not require training. | |
""" | |
def __init__(self, model, config, model_constructor, edit_loss_fn=None): | |
super().__init__(model, config, model_constructor) | |
if edit_loss_fn is not None: | |
self.edit_loss_fn = edit_loss_fn | |
self.locality_loss_fn = kl_loc_loss | |
self.loc_ids = None | |
self.loc_masks = None | |
self.loc_sampler = None | |
def _edit_loss(self, model, p0, p_edited, edit_batch): | |
output = _logits(model(**edit_batch, params=p_edited)) | |
loss_dict = self.edit_loss_fn(output, edit_batch["labels"]) | |
l_edit, acc = loss_dict["nll"], loss_dict["acc"] | |
if self.config.ft.locality.enabled: | |
if self.config.ft.locality.oracle: | |
loc_batch = next(self.loc_sampler)["loc"] | |
else: | |
raise NotImplementedError | |
with torch.no_grad(): | |
original_base_logits = _logits(model(**loc_batch, params=p0)) | |
edited_base_logits = _logits(model(**loc_batch, params=p_edited)) | |
kl_mask = loc_batch.get("decoder_attention_mask", loc_batch["attention_mask"]) | |
l_loc = self.locality_loss_fn(original_base_logits, edited_base_logits, mask=kl_mask) | |
loss = l_loc + self.config.ft.locality.cedit * l_edit | |
else: | |
l_loc = torch.tensor(float('nan')) | |
loss = l_edit | |
return loss, l_edit, l_loc, acc | |
def accuracy(self, output, labels): | |
if output.shape[-1] != 1: | |
shifted_output = output.argmax(-1)[:, :-1] | |
shifted_labels = labels[:, 1:] | |
to_predict = (shifted_labels != -100).sum() | |
correct = (shifted_output == shifted_labels).sum() | |
acc = correct.float() / to_predict.float() | |
else: | |
acc = ((output > 0) == labels.bool()).sum().float() | |
return acc | |
def _edit_status(self, step, loss, l_edit, l_loc, acc, res_p): | |
return ( | |
f"step: {step}".ljust(14) + | |
f"loss: {loss.item():.5f}".ljust(18) + | |
f"l_edit: {l_edit.item():.5f}".ljust(18) + | |
f"l_loc: {l_loc.item():.5f}".ljust(18) + | |
f"acc: {acc.item():.2f}".ljust(14) + | |
f"norm: {res_p.view(-1).norm().item():.5f}" | |
) | |
def edit(self, batch, condition=None, detach_history=False): | |
edit_model = self.model.eval() | |
p0 = list(edit_model.named_parameters()) | |
if not isinstance(edit_model, higher.patch._MonkeyPatchBase): | |
edit_model = make_functional(self.model, track_higher_grads=False, in_place=True) | |
packed_residuals = {} | |
opt_params = [] | |
for n, p in _inner_params(edit_model.named_parameters(), self.config.model.inner_params): | |
if self.config.ft.rank is not None: | |
u = nn.Parameter(torch.randn(p.shape[0], self.config.ft.rank, device=p.device) * self.config.ft.init_std) | |
v = nn.Parameter(torch.zeros(self.config.ft.rank, p.shape[1], device=p.device)) | |
res = [u, v] | |
else: | |
res = [nn.Parameter(torch.zeros_like(p, device=p.device))] | |
packed_residuals[n] = res | |
opt_params.extend(res) | |
assert len(opt_params) == len(self.config.model.inner_params) | |
OptClass = getattr(torch.optim, self.config.ft.opt) | |
opt = OptClass(opt_params, lr=self.config.edit_lr) | |
start_time = time.time() | |
for edit_step in range(self.config.ft.max_edit_steps): | |
if self.config.ft.time_limit is not None and (time.time() - start_time > self.config.ft.time_limit): | |
break | |
residuals = {k: v[0] @ v[1] if len(v) == 2 else v[0] for k, v in packed_residuals.items()} | |
edited_params = [p if n not in residuals else p.detach() + residuals[n] for n, p in p0] | |
loss, l_edit, l_loc, acc = self._edit_loss(edit_model, [p for n, p in p0], edited_params, batch) | |
if self.config.ft.verbose: | |
residual = list(residuals.values())[-1] | |
print(self._edit_status(edit_step, loss, l_edit, l_loc, acc, residual), end="\r") | |
if acc == 1.0: | |
break | |
for p, g in zip(opt_params, torch.autograd.grad(loss, opt_params)): | |
p.grad = g | |
torch.nn.utils.clip_grad_norm_(opt_params, self.config.grad_clip) | |
opt.step() | |
opt.zero_grad() | |
if detach_history: | |
new_model = self.model_constructor() | |
new_model.load_state_dict(edit_model.state_dict()) | |
edit_model = new_model | |
edit_model.train(self.training) | |
return FT(edit_model, self.config, self.model_constructor, self.edit_loss_fn), {} | |