File size: 2,286 Bytes
3d5551b | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import numpy as np
"""
Referenced from NCI codebase: https://github.com/solidsea98/Neural-Corpus-Indexer-NCI
"""
def decode_token(args, seqs):
"""
Args:
seqs: 2d ndarray to be decoded
Return:
doc_id string, List[str]
"""
result = []
for seq in seqs:
try:
eos_idx = seq.tolist().index(1)
seq = seq[1:eos_idx]
except:
print("no eos token found")
offset = np.arange(len(seq)) * args.output_vocab_size + 2
res = seq - offset
result.append("-".join(str(c) for c in res))
return result
def dec_2d(dec, size):
"""
Args:
dec: 1d ndarray
size: int
Return:
List[List[int]]
"""
res = []
i = 0
while i < len(dec):
res.append(dec[i : i + size])
i = i + size
return res
def encode_single_newid(args, seq):
"""
Args:
seq: doc_id string to be encoded, like "23456"
Return:
List[Int]: encoded tokens
"""
target_id_int = []
for i, c in enumerate(seq.split("-")):
cur_token = i * args.kary + int(c) + 2
target_id_int.append(cur_token)
return target_id_int + [1] # append eos_token
class Node(object):
def __init__(self, token_id):
self.token_id = token_id
self.children = {}
def __str__(self, level=0):
ret = "\t" * level + repr(self.token_id) + "\n"
for child in self.children.values():
ret += child.__str__(level + 1)
return ret
def __repr__(self):
return "<tree node representation>"
class TreeBuilder(object):
def __init__(self) -> None:
self.root = Node(0)
def build(self) -> Node:
return self.root
def add(self, seq) -> None:
"""
seq is List[Int] representing id, without leading pad_token(hardcoded 0), with trailing eos_token(hardcoded 1) and (possible) pad_token, every int is token index
e.g: [ 9, 14, 27, 38, 47, 58, 62, 1, 0, 0, 0, 0, 0, 0, 0]
"""
cur = self.root
for tok in seq:
if tok == 0: # reach pad_token
return
if tok not in cur.children:
cur.children[tok] = Node(tok)
cur = cur.children[tok]
|