import asyncio
from pathlib import Path
from pydantic_ai import Embedder
from pydantic_ai.embeddings.sentence_transformers import (
    SentenceTransformerEmbeddingModel,
    SentenceTransformersEmbeddingSettings
)

async def main():
    # 自动定位：脚本目录/models/模型文件夹
    base_path = Path(__file__).parent
    model_path = str(base_path / "models" / "paraphrase-multilingual-MiniLM-L12-v2")

    model = SentenceTransformerEmbeddingModel(
        model_path,
        settings=SentenceTransformersEmbeddingSettings(
            sentence_transformers_device="cpu",
            sentence_transformers_normalize_embeddings=True,
        )
    )
    embedder = Embedder(model)

    query_text = "刘弼仁"
    counts = await embedder.count_tokens(query_text)
    print(counts)
    max_tokens = await embedder.max_input_tokens()
    print(f'Max tokens: {max_tokens}')
    res = await embedder.embed_query(query_text)
    print(f"向量维度：{len(res.embeddings[0])}")

asyncio.run(main())