home neox utils

Github Twitter

View code on Github

#

公用事业和助手

15import typing
16from typing import List, Optional
17
18import torch
19
20from labml import logger
21from labml.logger import Text
22from labml_nn.neox.tokenizer import get_tokenizer
23
24if typing.TYPE_CHECKING:
25    from tokenizers import Tokenizer
#

分词器单例

28_TOKENIZER: Optional['Tokenizer'] = None
#

获取代币 ID

  • text 是要标记的文本

    • 返回令牌 ID

31def get_tokens(text: str) -> List[int]:
# 
38    global _TOKENIZER
39    if _TOKENIZER is None:
40        _TOKENIZER = get_tokenizer()
41    return _TOKENIZER.encode_batch([text])[0].ids
#

从模型输出中打印令牌

P@@

retty 沿着模型的侧面输出打印目标令牌。

  • ids 是目标代币 ID
  • xs 是模型的输出
44def print_token_outputs(ids: List[int], *xs: torch.Tensor):
# 
53    ids = ids + [-1]
54    xs = [[-1] + x[0].max(dim=-1)[1].tolist() for x in xs]
55
56    print_tokens(ids, xs)
#

打印代币

Pretty 打印令牌进行比较

  • target 是目标代币 ID
  • others 是来自模型的采样输出
59def print_tokens(target: List[int], others: List[List[int]]):
#

加载分词器

70    global _TOKENIZER
71    if _TOKENIZER is None:
72        _TOKENIZER = get_tokenizer()
#

将标记转换为字符串列表

75    text = []
76    for i in range(len(target)):
77        tokens = [_TOKENIZER.decode([target[i]]) if target[i] != -1 else '---']
78        for j in range(len(others)):
79            tokens.append(_TOKENIZER.decode([others[j][i]]) if others[j][i] != -1 else '---')
80
81        text.append(tokens)
#

统计数据

84    correct = [0 for _ in others]
85    total = 0
#

遍历令牌

88    for i in range(len(target)):
89        parts = [(f'{i}: ', Text.meta)]
90        parts += [('"', Text.subtle), (text[i][0], Text.subtle), ('"', Text.subtle), '\t']
#

空目标

93        if target[i] == -1:
94            for j in range(len(others)):
95                parts += [('"', Text.subtle), (text[i][j + 1], Text.subtle), ('"', Text.subtle), '\t']
96
97            logger.log(parts)
98            continue
#

代币数量

101        total += 1
#

其他输出

104        for j in range(len(others)):
105            correct[j] += 1 if others[j][i] == target[i] else 0
106
107            parts += [('"', Text.subtle),
108                      (text[i][j + 1], Text.success if others[j][i] == target[i] else Text.danger),
109                      ('"', Text.subtle), '\t']
110
111        logger.log(parts)
#

统计数据

114    parts = [(f'{total}', Text.highlight), '\t']
115    for j in range(len(others)):
116        parts += [(f'{correct[j]}', Text.value), '\t']
117    logger.log(parts)
#

平衡图层

拆分n_layers 成n_chunks 。这用于管道并行训练。

  • n_layers 是层数
  • n_chunks 是区块的数量

    • 返回一个列表，其中包含每个区块的层数

120def balance_layers_simple(n_layers: int, n_chunks: int):
# 
130    balance = []
131    for i in range(n_chunks):
132        balance.append((n_layers - sum(balance)) // (n_chunks - i))
133
134    return list(reversed(balance))
labml.ai