import torch


# 将文本编码为token ID序列，并转换为PyTorch张量
def text_to_token_ids(text, tokenizer):
    # 使用指定的分词器将文本编码为token ID列表
    encoded = tokenizer.encode(text, allowed_special={"<|endoftext|>"})  # 允许特殊字符
    # 将token ID列表转换为PyTorch张量，并添加batch维度 (1, sequence_length)
    encoded_tensor = torch.tensor(encoded).unsqueeze(0)  # 添加 batch 维度
    return encoded_tensor


# 将token ID张量解码为原始文本
def token_ids_to_text(token_ids, tokenizer):
    # 移除batch维度，得到形状为(sequence_length,)的一维张量
    flat = token_ids.squeeze(0)  # 删除 batch 维度
    # 将张量转换为Python列表，并使用分词器解码为文本
    return tokenizer.decode(flat.tolist())
