"""
测试 bm25s.tokenize 对中文的支持
"""
import bm25s

# 测试文本
texts = [
    "RAG是检索增强生成技术",
    "什么是RAG系统",
    "向量数据库用于存储嵌入向量"
]

print("=== 测试 bm25s.tokenize 的中文分词 ===\n")

# 测试不同的 stopwords 参数
for stopwords_param in [None, "zh", "chinese", "en", "english"]:
    print(f"stopwords='{stopwords_param}':")
    try:
        tokens = bm25s.tokenize(
            texts,
            stopwords=stopwords_param,
            stemmer=None,
            show_progress=False,
        )
        print(f"  结果类型: {type(tokens)}")
        print(f"  第一个文档的分词:")
        if hasattr(tokens, '__iter__'):
            first_tokens = tokens[0] if hasattr(tokens, '__getitem__') else list(tokens)[0]
            print(f"    {first_tokens}")
        else:
            print(f"    {tokens}")
    except Exception as e:
        print(f"  ❌ 错误: {e}")
    print()

# 测试查询分词
print("=== 测试查询分词 ===")
query = "什么是RAG"
print(f"查询: {query}\n")

for stopwords_param in [None, "zh", "chinese"]:
    print(f"stopwords='{stopwords_param}':")
    try:
        query_tokens = bm25s.tokenize(
            query,
            stopwords=stopwords_param,
            stemmer=None,
            show_progress=False,
        )
        print(f"  {query_tokens}")
    except Exception as e:
        print(f"  ❌ 错误: {e}")
    print()
