#分词技术（Tokenization）：中文 Jieba、WordPiece 与 BPE

import jieba

# 精确模式（最常用）
text = "自然语言处理是人工智能的重要分支"
words = jieba.lcut(text)
print(words)
# ['自然语言处理', '是', '人工智能', '的', '重要', '分支']

# 全模式（所有可能的切分）
words_all = jieba.lcut(text, cut_all=True)
print(words_all)
# ['自然', '自然语言', '语言', '言处', '处理', '是', '人工智能', '人工', '智能', '的', '重要', '分支']

# 搜索引擎模式（适合搜索引擎）
words_search = jieba.cut_for_search(text)
print(list(words_search))
# ['自然', '语言', '自然语言', '处理', '人工智能', '人工', '智能', '是', '的', '重要', '分支']

# 方式一：直接添加
jieba.add_word("自然语言处理")
jieba.add_word("深度学习")
jieba.add_word("大模型")

# 方式二：加载自定义词典文件
# 格式：词语 词频 词性
# file.txt:
# 自然语言处理 5 nz
# 大模型 10 nz
jieba.load_userdict("file.txt")

# 验证
print(jieba.lcut("自然语言处理和大模型是热点"))
# ['自然语言处理', '和', '大模型', '是', '热点']

import jieba.posseg as pseg

words = pseg.cut("小明在北京大学的图书馆里读书")
for word, flag in words:
    print(f"{word} / {flag}")
# 小明 / nr（人名）
# 在 / p（介词）
# 北京大学 / nt（机构名）
# 的 / u（助词）
# 图书馆 / n（名词）
# 里 / f（方位词）
# 读书 / v（动词）

# 基于 TF-IDF 提取关键词
import jieba.analyse

text = """
自然语言处理是计算机科学领域与人工智能领域中的一个重要方向。
它研究能实现人与计算机之间用自然语言进行有效通信的各种理论和方法。
"""

# topK: 提取前 K 个关键词
keywords = jieba.analyse.extract_tags(text, topK=5, withWeight=True)
print(keywords)
# [('自然语言处理', 1.5), ('计算机科学', 0.8), ...]

# 基于 TextRank 算法
keywords_rank = jieba.analyse.textrank(text, topK=5, withWeight=True)
print(keywords_rank)

text = "Natural language processing is fascinating!"
tokens = text.lower().split()
print(tokens)
# ['natural', 'language', 'processing', 'is', 'fascinating!']

# 问题：fascinating! 包含了标点
# 解决：先清洗标点
import re
text = re.sub(r'[^\w\s]', '', text.lower())
print(text.split())
# ['natural', 'language', 'processing', 'is', 'fascinating']

# 使用 transformers 库的 BertWordPieceTokenizer
from tokenizers import Tokenizer
from tokenizers.trainers import BpeTrainer, WordPieceTrainer
from tokenizers.normalizers import BertNormalizer
from tokenizers.pre_tokenizers import Whitespace

# 训练 BPE 分词器
tokenizer = Tokenizer(Whitespace())
trainer = BpeTrainer(vocab_size=30000, min_frequency=2)
files = ["data/train.txt"]

tokenizer.train(files, trainer)
tokenizer.save("my-bpe-tokenizer.json")

# 使用
output = tokenizer.encode("Natural language processing")
print(output.tokens)
# ['Natural', 'Ġlanguage', 'Ġprocessing']
# Ġ 表示空格

from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")

# 中文分词
text = "自然语言处理很有趣"
tokens = tokenizer.tokenize(text)
print(tokens)
# ['自', '然', '语', '言', '处', '理', '很', '有', '趣']

# 英文分词
text_en = "Natural language processing"
tokens_en = tokenizer.tokenize(text_en)
print(tokens_en)
# ['natural', 'language', 'processing']

# 转换为 ID
ids = tokenizer.convert_tokens_to_ids(tokens)
print(ids)
# [...对应的ID...]

# 批量编码
encoded = tokenizer(["我爱中国", "机器学习"], padding=True, truncation=True, max_length=10)
print(encoded)
# {'input_ids': [...], 'token_type_ids': [...], 'attention_mask': [...]}

import jieba
import re
from transformers import BertTokenizer

class NLPTokenizer:
    def __init__(self, lang="zh"):
        self.lang = lang
        if lang == "zh":
            self.tokenizer = jieba
        else:
            self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

    def tokenize(self, text):
        if self.lang == "zh":
            # 中文：jieba 精确模式 + 标点清洗
            text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9\s]', '', text)
            return list(self.tokenizer.lcut(text))
        else:
            # 英文：BERT 分词
            return self.tokenizer.tokenize(text.lower())

    def encode(self, text, max_length=128):
        """返回 token IDs"""
        if self.lang == "zh":
            return self.tokenizer.convert_tokens_to_ids(self.tokenize(text))
        else:
            return self.tokenizer.encode(text, max_length=max_length, truncation=True)

# 使用
tokenizer = NLPTokenizer(lang="zh")
print(tokenizer.tokenize("自然语言处理"))
# ['自然语言处理']

tokenizer_en = NLPTokenizer(lang="en")
print(tokenizer_en.tokenize("Natural language processing"))
# ['natural', 'language', 'processing']

# 分词速查

# 中文
import jieba
jieba.lcut(text)                    # 精确分词
jieba.lcut(text, cut_all=True)       # 全模式
jieba.add_word("新词")               # 添加词典
jieba.analyse.extract_tags(text)      # 关键词提取

# 英文 / 多语言
from transformers import BertTokenizer
tokenizer.tokenize(text)            # 分词
tokenizer.encode(text)              # 转 ID
tokenizer.decode(ids)               # ID 转文本