File size: 1,585 Bytes
fc32096 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | import json
from pathlib import Path
import torch
from transformers import AutoTokenizer
from safetensors import safe_open
from common import ROOT, permit_compute, append
from weights import headers, load as read_weight, save_edit as write_edit
permit_compute()
torch.set_grad_enabled(False)
MODEL = ROOT / 'models/original'
CURRENT = ROOT / 'edits/current'
NAMES = sorted(headers(MODEL))
tok = AutoTokenizer.from_pretrained(MODEL, local_files_only=True)
def load_original(name):
return read_weight(MODEL, name).float()
def load(name):
if name not in NAMES:
raise KeyError(name)
path = CURRENT / (name + '.safetensors')
if path.exists():
with safe_open(str(path), framework='pt', device='cpu') as f:
return f.get_tensor(name).float()
return load_original(name)
def save_edit(name, tensor, hypothesis):
return write_edit(MODEL, CURRENT, name, tensor, hypothesis)
def edited():
return sorted(p.name.removesuffix('.safetensors') for p in CURRENT.glob('*.safetensors'))
def drop_edit(name):
permit_compute()
if name not in NAMES:
raise KeyError(name)
for suffix in ('.safetensors', '.json', '.analysis.json', '.direction.npy'):
path = CURRENT / (name + suffix)
if path.exists():
path.unlink()
append(ROOT / 'logs/experiments.jsonl', {'action': 'drop_edit', 'tensor': name})
print('CPU kernel готов: NAMES, load, load_original, save_edit(name,tensor,hypothesis), drop_edit, edited, tok. Градиенты отключены. Модель не запущена.')
|