Spaces:
Sleeping
Sleeping
File size: 2,224 Bytes
bda104d | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | from transformers import AutoTokenizer
import torch
# Extract subjects from prompt
def extract_subjects(prompt, start_subject="<subject>", end_subject="</subject>"):
subjects = []
start = 0
while True:
start = prompt.find(start_subject, start)
if start == -1:
break
end = prompt.find(end_subject, start)
if end == -1:
break
subjects.append(prompt[start + len(start_subject):end])
start = end + len(end_subject)
return subjects
# Remove delimiters from prompt
def remove_delimiters(prompt, start_subject="<subject>", end_subject="</subject>"):
return prompt.replace(start_subject, "").replace(end_subject, "")
# Find subject positions in clean prompt
def find_subject_positions(clean_prompt, subjects):
positions = []
for subject in subjects:
start = 0
while True:
start = clean_prompt.find(subject, start)
if start == -1:
break
end = start + len(subject)
positions.append([start, end])
start = end
return positions
# Create binary mask
def create_binary_mask(clean_prompt, adjusted_positions, offsets):
mask = [0] * len(offsets)
for start, end in adjusted_positions:
for i, (token_start, token_end) in enumerate(offsets):
if token_start >= start and token_end <= end:
mask[i] = 1
return mask
# Process list of prompts
def create_binary_mask_from_list(dirty_prompts, offsets, start_subject="<subject>", end_subject="</subject>"):
clean_prompts = [remove_delimiters(prompt, start_subject, end_subject) for prompt in dirty_prompts]
if isinstance(offsets, torch.Tensor):
offsets = offsets.tolist()
masks = []
for dirty_prompt, clean_prompt, offset in zip(dirty_prompts, clean_prompts, offsets):
subjects = extract_subjects(dirty_prompt, start_subject, end_subject)
adjusted_positions = find_subject_positions(clean_prompt, subjects)
mask = create_binary_mask(clean_prompt, adjusted_positions, offset)
masks.append(mask)
return {
"masks": masks,
"clean_prompts": clean_prompts
}
|