315 lines
12 KiB
Python
315 lines
12 KiB
Python
import regex
|
||
from unicodedata import normalize
|
||
from .bpe_utils import bytes_to_unicode_special
|
||
from random import shuffle, seed
|
||
from math import ceil
|
||
|
||
def pretok_with_targeted_handling(text_data, split_regex, target_regex, target_process, nontarget_process):
|
||
if split_regex:
|
||
# 1st Regex: The general split regex
|
||
split_regex = regex.compile(split_regex)
|
||
split_regexed = split_regex.findall(text_data)
|
||
else:
|
||
split_regexed = [text_data]
|
||
|
||
# 2nd Regex: Splits a string into target and non-target chunks
|
||
final_processed_list = []
|
||
for chunk in split_regexed:
|
||
if not chunk:
|
||
continue # Skip empty tokens if any
|
||
|
||
# Split the token into target/non-target parts
|
||
processed_sub_tokens = ""
|
||
for match in regex.finditer(target_regex, chunk):
|
||
# Check if the chunk is target to decide which process to use
|
||
label = match.lastgroup
|
||
content = match.group()
|
||
if label == 'target':
|
||
processed_sub_tokens+=(target_process(content))
|
||
else:
|
||
processed_sub_tokens+=(nontarget_process(content))
|
||
|
||
final_processed_list.append(processed_sub_tokens)
|
||
|
||
return final_processed_list
|
||
|
||
def pretok_with_targeted_handling_list(text_data, split_regex, target_regex, target_process, nontarget_process):
|
||
if split_regex:
|
||
# 1st Regex: The general split regex
|
||
split_regex = regex.compile(split_regex)
|
||
split_regexed = split_regex.findall(text_data)
|
||
else:
|
||
split_regexed = [text_data]
|
||
# 2nd Regex: Splits a string into target and non-target chunks
|
||
final_processed_list = []
|
||
for chunk in split_regexed:
|
||
if not chunk:
|
||
continue # Skip empty tokens if any
|
||
# Split the token into target/non-target parts
|
||
for match in regex.finditer(target_regex, chunk):
|
||
# Check if the chunk is target to decide which process to use
|
||
label = match.lastgroup
|
||
content = match.group()
|
||
if label == 'target':
|
||
final_processed_list+=target_process(content)
|
||
else:
|
||
final_processed_list+=nontarget_process(content)
|
||
return final_processed_list
|
||
|
||
|
||
# String "좋아" -> String "ƴƸƯŠ"
|
||
def apply_encoding(input_string, encoding_type, mapping, offset=0, shuffle=False, offset_mapping=None, unicode_start=None):
|
||
if "Jamo" in encoding_type:
|
||
varlen = "varlen" in encoding_type
|
||
final_string = ""
|
||
for char in input_string:
|
||
if char == " ":
|
||
final_string += mapping[ord(char)]
|
||
continue
|
||
normalized = normalize("NFKD", char)
|
||
if len(normalized)==1: # Single Jamo can be used, but if not usable need fallback
|
||
if normalized not in mapping:
|
||
final_string += apply_encoding(normalized, "utf-8", mapping)
|
||
continue
|
||
final_string += normalized
|
||
if not varlen and len(normalized)==2: # Hangul syllable decomposed into 2 Jamo
|
||
final_string += '∅' # Placeholder for missing Jamo
|
||
return final_string
|
||
if shuffle:
|
||
# For each character, calculate offset and map to shuffled offset
|
||
input_string = "".join([chr(offset_mapping[ord(c) - unicode_start] + unicode_start) for c in input_string])
|
||
if encoding_type == "nybbles":
|
||
return apply_nybble_encoding(input_string, mapping, offset)
|
||
encoded = input_string.encode(encoding_type)
|
||
return "".join([mapping.get(x + offset) if x + offset in mapping else mapping.get(x) for x in encoded])
|
||
|
||
# String "ƴƸƯŠ" -> String "좋아"
|
||
def apply_decoding(input_string, encoding_type, mapping, offset=0):
|
||
antimap = {value: key for key, value in mapping.items()}
|
||
return b"".join([antimap[a+offset].to_bytes() for a in input_string]).decode(encoding_type)
|
||
|
||
def apply_character_decoding(input_string, mapping):
|
||
antimap = {value: key for key, value in mapping.items()}
|
||
final_string = ""
|
||
temp_string = b""
|
||
char_mode = False
|
||
for symbol in input_string:
|
||
# For token in charmode
|
||
if symbol not in antimap:
|
||
# If we are in the normal case, resolve the current temp_string
|
||
if not char_mode:
|
||
final_string += temp_string.decode("utf8", errors='replace')
|
||
temp_string = ""
|
||
temp_string += symbol
|
||
char_mode = True
|
||
# For tokens in regular case
|
||
else:
|
||
# If we are in charmode, resolve the current temp_string
|
||
if char_mode:
|
||
final_string += temp_string
|
||
temp_string = b""
|
||
temp_string += (antimap[symbol]).to_bytes()
|
||
char_mode = False
|
||
# Resolve final piece
|
||
if char_mode:
|
||
final_string += temp_string
|
||
else:
|
||
final_string += temp_string.decode('utf8', errors='replace')
|
||
|
||
return final_string
|
||
|
||
# String "è¿Ļæĺ¯ƴƸƯŠ" -> String "这是좋아"
|
||
def apply_specialized_decoding(input_string, encoding_type, mapping, offset=0, shuffle=False, offset_mapping=None, unicode_start=None):
|
||
if shuffle:
|
||
def unshuffle_char(c):
|
||
shuffled_offset = ord(c) - unicode_start
|
||
original_offset = offset_mapping[shuffled_offset]
|
||
return chr(original_offset + unicode_start)
|
||
|
||
antimap = {value: key for key, value in mapping.items()}
|
||
final_string = ""
|
||
temp_string = b""
|
||
special_mode = False
|
||
for symbol in input_string:
|
||
byte_id = antimap[symbol]
|
||
# For token in special case
|
||
if byte_id >= offset:
|
||
# If we are in the normal case, resolve the current temp_string
|
||
if not special_mode:
|
||
final_string += temp_string.decode("utf8", errors='replace')
|
||
temp_string = b""
|
||
temp_string += (byte_id - offset).to_bytes()
|
||
special_mode = True
|
||
# For tokens in regular case
|
||
else:
|
||
# If we are in the special case, resolve the current temp_string
|
||
if special_mode:
|
||
temp_add = temp_string.decode(encoding_type, errors='replace')
|
||
if shuffle:
|
||
temp_add = "".join([unshuffle_char(c) for c in temp_add])
|
||
final_string += temp_add
|
||
temp_string = b""
|
||
temp_string += (byte_id).to_bytes()
|
||
special_mode = False
|
||
# Resolve final piece
|
||
if special_mode:
|
||
temp_add = temp_string.decode(encoding_type, errors='replace')
|
||
if shuffle:
|
||
temp_add = "".join([unshuffle_char(c) for c in temp_add])
|
||
final_string += temp_add
|
||
else:
|
||
final_string += temp_string.decode('utf8', errors='replace')
|
||
|
||
return final_string
|
||
|
||
def apply_nybble_decoding(input_string, mapping, offset=0):
|
||
def decode_char(temp_string):
|
||
res = []
|
||
for i in range(0, len(temp_string), 4):
|
||
try:
|
||
n0, n1, n2, n3 = [t for t in temp_string[i:i+4]]
|
||
codepoint = (n0 << 12) | (n1 << 8) | (n2 << 4) | n3
|
||
res.append(chr(codepoint))
|
||
except:
|
||
# In case of any error (e.g., invalid character), append a replacement character
|
||
res.append('<EFBFBD>')
|
||
return "".join(res)
|
||
|
||
antimap = {value: key for key, value in mapping.items()}
|
||
final_string = ""
|
||
temp_string = b""
|
||
special_mode = False
|
||
for symbol in input_string:
|
||
byte_id = antimap[symbol]
|
||
# For token in special case
|
||
if byte_id >= offset:
|
||
# If we are in the normal case, resolve the current temp_string
|
||
if not special_mode:
|
||
final_string += temp_string.decode("utf8", errors='replace')
|
||
temp_string = b""
|
||
temp_string += (byte_id - offset).to_bytes()
|
||
special_mode = True
|
||
# For tokens in regular case
|
||
else:
|
||
# If we are in the special case, resolve the current temp_string
|
||
if special_mode:
|
||
final_string += decode_char(temp_string)
|
||
temp_string = b""
|
||
temp_string += (byte_id).to_bytes()
|
||
special_mode = False
|
||
# Resolve final piece
|
||
if special_mode:
|
||
final_string += decode_char(temp_string)
|
||
else:
|
||
final_string += temp_string.decode('utf8', errors='replace')
|
||
|
||
return final_string
|
||
|
||
def apply_Jamo_decoding(input_string, mapping):
|
||
antimap = {value: key for key, value in mapping.items()}
|
||
final_string = ""
|
||
current_type = type(input_string[0])
|
||
current_pieces = []
|
||
for char in input_string:
|
||
char = antimap[char]
|
||
if type(char) == current_type:
|
||
current_pieces.append(char)
|
||
else:
|
||
# Process current pieces
|
||
if current_type == str:
|
||
# Jamo pieces
|
||
syllables = "".join((current_pieces)).replace("∅", "")
|
||
# Put it into unicode format
|
||
final_string += normalize("NFKC", syllables)
|
||
else:
|
||
# Byte pieces
|
||
byte_string = b"".join([c.to_bytes() for c in current_pieces]).decode('utf-8', errors='replace')
|
||
final_string += byte_string
|
||
current_type = type(char)
|
||
current_pieces = [char]
|
||
# Final flush
|
||
if current_type == str:
|
||
# Jamo pieces
|
||
syllables = "".join((current_pieces)).replace("∅", "")
|
||
final_string += syllables
|
||
else:
|
||
# Byte pieces
|
||
byte_string = b"".join([c.to_bytes() for c in current_pieces]).decode('utf-8', errors='replace')
|
||
final_string += byte_string
|
||
return final_string
|
||
|
||
|
||
|
||
def apply_nybble_encoding(input_string, mapping, offset,n_nybbles=4):
|
||
def encode_char(c):
|
||
o = ord(c)
|
||
return "".join(mapping[offset + ((o >> (4 * i)) & 0xF)] for i in range(n_nybbles - 1, -1, -1))
|
||
|
||
return "".join(encode_char(c) for c in input_string)
|
||
|
||
# Given a range of unicode, we shuffle the points
|
||
# For performance this can be stored as a list, a mapping of offset2newoffset
|
||
def create_shuffled_unicode_mapping(range_start, range_end, given_seed=6767):
|
||
unicode_points = list(range(range_end - range_start + 1))
|
||
seed(given_seed)
|
||
shuffle(unicode_points)
|
||
# to use, calculate offset = ord(char) - range_start, then map to unicode_points[offset] + range_start
|
||
return unicode_points
|
||
|
||
# Given a range of unicode, provide a Customized UTF with given parameter properties
|
||
def create_CuTF_encoding(range_start, range_end, encode_length):
|
||
target_space_len = range_end - range_start + 1
|
||
num_indices = ceil(target_space_len**(1/encode_length))
|
||
# Plan is to do the math only here and save the mapping for efficiency
|
||
CuTF_mapping = {}
|
||
for codepoint in range(range_start, range_end+1):
|
||
offset = codepoint-range_start
|
||
CuTF_rep = []
|
||
for i in range(encode_length-1):
|
||
rep_index = offset//(num_indices**(encode_length-i-1))
|
||
CuTF_rep.append(rep_index)
|
||
offset -= (num_indices**(encode_length-i-1))*rep_index
|
||
CuTF_rep.append(offset)
|
||
CuTF_mapping[codepoint] = CuTF_rep
|
||
|
||
return CuTF_mapping, num_indices
|
||
|
||
|
||
|
||
def apply_CuTF_decoding(input_string, encode_length, num_index, range_start, indexing_strategy="shared"):
|
||
def decode_target_chars(sequence):
|
||
outpieces = ""
|
||
for i in range(0, len(sequence), encode_length):
|
||
piece = sequence[i:i+encode_length]
|
||
offset = 0
|
||
for j, byte in enumerate(piece):
|
||
if indexing_strategy == "shared":
|
||
offset += (byte-256) * (num_index**(encode_length-j-1))
|
||
elif indexing_strategy == "unique":
|
||
offset += (byte - (256 + num_index*(encode_length-j-1))) * (num_index**(encode_length-j-1))
|
||
outpieces += chr(range_start + offset)
|
||
return outpieces
|
||
final_string = ""
|
||
temp_pieces = []
|
||
is_target = False
|
||
for token in input_string:
|
||
if token >= 256:
|
||
if not is_target:
|
||
# Flush non-target tokens
|
||
final_string += bytes(temp_pieces).decode("utf8", errors='replace')
|
||
temp_pieces = []
|
||
temp_pieces.append(token)
|
||
is_target = True
|
||
else:
|
||
if is_target:
|
||
# Flush target tokens
|
||
final_string += decode_target_chars(temp_pieces)
|
||
temp_pieces = []
|
||
temp_pieces.append(token)
|
||
is_target = False
|
||
# Flush any remaining tokens
|
||
if is_target:
|
||
final_string += decode_target_chars(temp_pieces)
|
||
else:
|
||
final_string += bytes(temp_pieces).decode("utf8", errors='replace')
|
||
return final_string |