Spaces:
Runtime error
Runtime error
| # Copyright 2022 Google LLC | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import torch | |
| 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, tokenizer, max_len=77): | |
| x_seq = tokenizer.encode(x) | |
| y_seq = tokenizer.encode(y) | |
| 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)) | |
| return mapper, alphas | |
| def get_refinement_mapper(prompts, tokenizer, max_len=77): | |
| x_seq = prompts[0] | |
| mappers, alphas = [], [] | |
| for i in range(1, len(prompts)): | |
| mapper, alpha = get_mapper(x_seq, prompts[i], tokenizer, max_len) | |
| mappers.append(mapper) | |
| alphas.append(alpha) | |
| return torch.stack(mappers), torch.stack(alphas) | |
| 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) | |