InfEdit / seq_aligner.py
6kplus's picture
Upload 30 files
948429b
import torch
import copy
import torch.nn.functional as F
import numpy as np
class ScoreParams:
def __init__(self, gap, match, mismatch):
self.gap = gap
self.match = match
self.mismatch = mismatch
def mis_match_char(self, x, y):
if x != y:
return self.mismatch
else:
return self.match
def get_matrix(size_x, size_y, gap):
matrix = []
for i in range(len(size_x) + 1):
sub_matrix = []
for j in range(len(size_y) + 1):
sub_matrix.append(0)
matrix.append(sub_matrix)
for j in range(1, len(size_y) + 1):
matrix[0][j] = j*gap
for i in range(1, len(size_x) + 1):
matrix[i][0] = i*gap
return matrix
def get_matrix(size_x, size_y, gap):
matrix = np.zeros((size_x + 1, size_y + 1), dtype=np.int32)
matrix[0, 1:] = (np.arange(size_y) + 1) * gap
matrix[1:, 0] = (np.arange(size_x) + 1) * gap
return matrix
def get_traceback_matrix(size_x, size_y):
matrix = np.zeros((size_x + 1, size_y +1), dtype=np.int32)
matrix[0, 1:] = 1
matrix[1:, 0] = 2
matrix[0, 0] = 4
return matrix
def global_align(x, y, score):
matrix = get_matrix(len(x), len(y), score.gap)
trace_back = get_traceback_matrix(len(x), len(y))
for i in range(1, len(x) + 1):
for j in range(1, len(y) + 1):
left = matrix[i, j - 1] + score.gap
up = matrix[i - 1, j] + score.gap
diag = matrix[i - 1, j - 1] + score.mis_match_char(x[i - 1], y[j - 1])
matrix[i, j] = max(left, up, diag)
if matrix[i, j] == left:
trace_back[i, j] = 1
elif matrix[i, j] == up:
trace_back[i, j] = 2
else:
trace_back[i, j] = 3
return matrix, trace_back
def get_aligned_sequences(x, y, trace_back):
x_seq = []
y_seq = []
i = len(x)
j = len(y)
mapper_y_to_x = []
while i > 0 or j > 0:
if trace_back[i, j] == 3:
x_seq.append(x[i-1])
y_seq.append(y[j-1])
i = i-1
j = j-1
mapper_y_to_x.append((j, i))
elif trace_back[i][j] == 1:
x_seq.append('-')
y_seq.append(y[j-1])
j = j-1
mapper_y_to_x.append((j, -1))
elif trace_back[i][j] == 2:
x_seq.append(x[i-1])
y_seq.append('-')
i = i-1
elif trace_back[i][j] == 4:
break
mapper_y_to_x.reverse()
return x_seq, y_seq, torch.tensor(mapper_y_to_x, dtype=torch.int64)
def get_mapper(x: str, y: str, specifier, tokenizer, encoder, device, max_len=77):
locol_prompt, mutual_prompt = specifier
x_seq = tokenizer.encode(x)
y_seq = tokenizer.encode(y)
e_seq = tokenizer.encode(locol_prompt)
m_seq = tokenizer.encode(mutual_prompt)
score = ScoreParams(0, 1, -1)
matrix, trace_back = global_align(x_seq, y_seq, score)
mapper_base = get_aligned_sequences(x_seq, y_seq, trace_back)[-1]
alphas = torch.ones(max_len)
alphas[: mapper_base.shape[0]] = mapper_base[:, 1].ne(-1).float()
mapper = torch.zeros(max_len, dtype=torch.int64)
mapper[:mapper_base.shape[0]] = mapper_base[:, 1]
mapper[mapper_base.shape[0]:] = len(y_seq) + torch.arange(max_len - len(y_seq))
m = copy.deepcopy(alphas)
alpha_e = torch.zeros_like(alphas)
alpha_m = torch.zeros_like(alphas)
# print("mapper of")
# print("<begin> "+x+" <end>")
# print("<begin> "+y+" <end>")
# print(mapper[:len(y_seq)])
# print(alphas[:len(y_seq)])
x = tokenizer(
x,
padding="max_length",
max_length=max_len,
truncation=True,
return_tensors="pt",
).input_ids.to(device)
y = tokenizer(
y,
padding="max_length",
max_length=max_len,
truncation=True,
return_tensors="pt",
).input_ids.to(device)
x_latent = encoder(x)[0].squeeze(0)
y_latent = encoder(y)[0].squeeze(0)
i = 0
while i<len(y_seq):
start = None
if alphas[i] == 0:
start = i
while alphas[i] == 0:
i += 1
max_sim = float('-inf')
max_s = None
max_t = None
for i_target in range(start, i):
for i_source in range(mapper[start-1]+1, mapper[i]):
sim = F.cosine_similarity(x_latent[i_target], y_latent[i_source], dim=0)
if sim > max_sim:
max_sim = sim
max_s = i_source
max_t = i_target
if max_s is not None:
mapper[max_t] = max_s
alphas[max_t] = 1
for t in e_seq:
if x_seq[max_s] == t:
alpha_e[max_t] = 1
i += 1
# replace_alpha, replace_mapper = get_replace_inds(x_seq, y_seq, m_seq, m_seq)
# if replace_mapper != []:
# mapper[replace_alpha]=torch.tensor(replace_mapper,device=mapper.device)
# alpha_m[replace_alpha]=1
i = 1
j = 1
while (i < len(y_seq)-1) and (j < len(e_seq)-1):
found = True
while e_seq[j] != y_seq[i]:
i = i + 1
if i >= len(y_seq)-1:
print("blend word not found!")
found = False
break
raise ValueError("local prompt not found in target prompt")
if found:
alpha_e[i] = 1
j = j + 1
i = 1
j = 1
while (i < len(y_seq)-1) and (j < len(m_seq)-1):
while m_seq[j] != y_seq[i]:
i = i + 1
if m_seq[j] == x_seq[mapper[i]]:
alpha_m[i] = 1
j = j + 1
else:
raise ValueError("mutual prompt not found in target prompt")
# print("fixed mapper:")
# print(mapper[:len(y_seq)])
# print(alphas[:len(y_seq)])
# print(m[:len(y_seq)])
# print(alpha_e[:len(y_seq)])
# print(alpha_m[:len(y_seq)])
return mapper, alphas, m, alpha_e, alpha_m
def get_refinement_mapper(prompts, specifiers, tokenizer, encoder, device, max_len=77):
x_seq = prompts[0]
mappers, alphas, ms, alpha_objs, alpha_descs = [], [], [], [], []
for i in range(1, len(prompts)):
mapper, alpha, m, alpha_obj, alpha_desc = get_mapper(x_seq, prompts[i], specifiers[i-1], tokenizer, encoder, device, max_len)
mappers.append(mapper)
alphas.append(alpha)
ms.append(m)
alpha_objs.append(alpha_obj)
alpha_descs.append(alpha_desc)
return torch.stack(mappers), torch.stack(alphas), torch.stack(ms), torch.stack(alpha_objs), torch.stack(alpha_descs)
def get_replace_inds(x_seq,y_seq,source_replace_seq,target_replace_seq):
replace_mapper=[]
replace_alpha=[]
source_found=False
source_match,target_match=[],[]
for j in range(len(x_seq)):
found=True
for i in range(1,len(source_replace_seq)-1):
if x_seq[j+i-1]!=source_replace_seq[i]:
found=False
break
if found:
source_found=True
for i in range(1,len(source_replace_seq)-1):
source_match.append(j+i-1)
for j in range(len(y_seq)):
found=True
for i in range(1,len(target_replace_seq)-1):
if y_seq[j+i-1]!=target_replace_seq[i]:
found=False
break
if found:
for i in range(1,len(source_replace_seq)-1):
target_match.append(j+i-1)
if not source_found:
raise ValueError("replacing object not found in prompt")
if (len(source_match)!=len(target_match)):
raise ValueError(f"the replacement word number doesn't match for word {i}!")
replace_alpha+=source_match
replace_mapper+=target_match
return replace_alpha,replace_mapper
def get_word_inds(text: str, word_place: int, tokenizer):
split_text = text.split(" ")
if type(word_place) is str:
word_place = [i for i, word in enumerate(split_text) if word_place == word]
elif type(word_place) is int:
word_place = [word_place]
out = []
if len(word_place) > 0:
words_encode = [tokenizer.decode([item]).strip("#") for item in tokenizer.encode(text)][1:-1]
cur_len, ptr = 0, 0
for i in range(len(words_encode)):
cur_len += len(words_encode[i])
if ptr in word_place:
out.append(i + 1)
if cur_len >= len(split_text[ptr]):
ptr += 1
cur_len = 0
return np.array(out)
def get_replacement_mapper_(x: str, y: str, tokenizer, max_len=77):
words_x = x.split(' ')
words_y = y.split(' ')
if len(words_x) != len(words_y):
raise ValueError(f"attention replacement edit can only be applied on prompts with the same length"
f" but prompt A has {len(words_x)} words and prompt B has {len(words_y)} words.")
inds_replace = [i for i in range(len(words_y)) if words_y[i] != words_x[i]]
inds_source = [get_word_inds(x, i, tokenizer) for i in inds_replace]
inds_target = [get_word_inds(y, i, tokenizer) for i in inds_replace]
mapper = np.zeros((max_len, max_len))
i = j = 0
cur_inds = 0
while i < max_len and j < max_len:
if cur_inds < len(inds_source) and inds_source[cur_inds][0] == i:
inds_source_, inds_target_ = inds_source[cur_inds], inds_target[cur_inds]
if len(inds_source_) == len(inds_target_):
mapper[inds_source_, inds_target_] = 1
else:
ratio = 1 / len(inds_target_)
for i_t in inds_target_:
mapper[inds_source_, i_t] = ratio
cur_inds += 1
i += len(inds_source_)
j += len(inds_target_)
elif cur_inds < len(inds_source):
mapper[i, j] = 1
i += 1
j += 1
else:
mapper[j, j] = 1
i += 1
j += 1
return torch.from_numpy(mapper).float()
def get_replacement_mapper(prompts, tokenizer, max_len=77):
x_seq = prompts[0]
mappers = []
for i in range(1, len(prompts)):
mapper = get_replacement_mapper_(x_seq, prompts[i], tokenizer, max_len)
mappers.append(mapper)
return torch.stack(mappers)