File size: 2,622 Bytes
72c0672 | 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 67 68 69 70 71 72 73 | import json
from collections import defaultdict
import numpy as np
def unigram_first_byte_prob(corpus_path: str, output_path: str):
with open(corpus_path, 'r', encoding='utf-8') as f:
utf8_byte_lines = [json.loads(line)['text'].encode('utf-8') for line in f]
byte_counts = defaultdict(int)
total_bytes = 0
for utf8_bytes in utf8_byte_lines:
for byte in utf8_bytes:
byte_counts[byte] += 1
total_bytes += 1
unigram_prob = [0.0] * 256
for byte, count in byte_counts.items():
unigram_prob[byte] = count / total_bytes
for byte in range(256):
p = unigram_prob[byte]
# Get ASCII character representation (if printable)
if 32 <= byte <= 126:
char = chr(byte)
elif byte in [9, 10, 13]: # Tab, Line Feed, Carriage Return
char = {9: '\\t', 10: '\\n', 13: '\\r'}.get(byte)
else:
char = '<control>' # Non-printable control character
# Get common name/description for selected bytes
desc = {
0: 'NUL (Null)',
9: 'HT (Horizontal Tab)',
10: 'LF (Line Feed)',
13: 'CR (Carriage Return)',
27: 'ESC (Escape)',
32: 'SPACE',
127: 'DEL (Delete)'
}.get(byte, '')
print(f"{byte}\t{p:.6f}\t\t{char}\t\t{desc}")
with open(output_path, 'w') as f:
json.dump(unigram_prob, f)
def smoothing_unigram_prob(unigram_prob_path: str, smoothing_factor: float, output_path: str):
with open(unigram_prob_path, 'r') as f:
unigram_prob = json.load(f)
unigram_prob_np = np.array(unigram_prob, dtype=np.float64)
vocab_size = unigram_prob_np.shape[0]
# Calculate smoothed probability
smoothed_prob_np = unigram_prob_np * (1 - smoothing_factor) + smoothing_factor / vocab_size
smoothed_prob = smoothed_prob_np.astype(np.float32).tolist()
# Save smoothed probability
with open(output_path, 'w') as f:
json.dump(smoothed_prob, f)
if __name__ == '__main__':
unigram_prob_path = '/mnt/bn/tiktok-mm-5/aiic/users/linzheng/artifacts/ac_unigram_probs/opencoder13G_unigram_prob.json'
smoothing_factor = 0.2
smoothed_unigram_prob_path = '/mnt/bn/tiktok-mm-5/aiic/users/linzheng/artifacts/ac_unigram_probs/opencoder13G_unigram_prob_smooth{}.json'
# unigram_first_byte_prob(
# 'subsample_opencoder.jsonl',
# unigram_prob_path
# )
smoothing_unigram_prob(
unigram_prob_path,
smoothing_factor,
smoothed_unigram_prob_path.format(smoothing_factor)
)
|