maque 0.2.1__py3-none-any.whl

maque/performance/_profiler.py

@@ -0,0 +1,367 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+现代化性能分析器
+
+提供两种分析器：
+1. Profile - 基于 pyinstrument，轻量级 CPU 分析
+2. ScaleneProfile - 基于 Scalene，CPU + 内存 + GPU 全面分析
+
+Example:
+    # pyinstrument (轻量级)
+    with Profile("数据处理") as p:
+        process_data()
+
+    # Scalene (全面分析)
+    with ScaleneProfile("内存分析", memory=True):
+        process_large_data()
+"""
+
+from typing import Optional, Literal
+from pathlib import Path
+import functools
+import subprocess
+import sys
+import tempfile
+import os
+
+try:
+    from pyinstrument import Profiler
+    PYINSTRUMENT_AVAILABLE = True
+except ImportError:
+    PYINSTRUMENT_AVAILABLE = False
+    Profiler = None
+
+try:
+    from scalene import scalene_profiler
+    SCALENE_AVAILABLE = True
+except ImportError:
+    SCALENE_AVAILABLE = False
+    scalene_profiler = None
+
+
+OutputFormat = Literal["text", "html", "json", "speedscope"]
+
+
+class Profile:
+    """
+    现代化性能分析器
+
+    基于 pyinstrument 的采样式分析，低开销，支持异步代码。
+
+    Example:
+        >>> with Profile("任务名称") as p:
+        ...     time.sleep(0.1)
+        ...     do_something()
+
+        # 查看 HTML 报告
+        >>> p.open_in_browser()
+
+        # 保存报告
+        >>> p.save("report.html")
+    """
+
+    def __init__(
+        self,
+        name: str = "",
+        *,
+        interval: float = 0.001,  # 采样间隔（秒）
+        async_mode: str = "enabled",  # enabled, disabled, strict
+        show: bool = True,  # 退出时是否自动打印
+        show_all: bool = False,  # 显示所有帧（包括库代码）
+        timeline: bool = False,  # 时间线模式
+        output: OutputFormat = "text",  # 输出格式
+    ):
+        if not PYINSTRUMENT_AVAILABLE:
+            raise ImportError(
+                "pyinstrument 未安装，请运行: pip install pyinstrument"
+            )
+
+        self.name = name
+        self.show = show
+        self.show_all = show_all
+        self.output = output
+        self._profiler = Profiler(interval=interval, async_mode=async_mode)
+        self._timeline = timeline
+
+    def __enter__(self):
+        self._profiler.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._profiler.stop()
+        if self.show:
+            self.print()
+        return False
+
+    def start(self):
+        """手动启动分析"""
+        self._profiler.start()
+        return self
+
+    def stop(self):
+        """手动停止分析"""
+        self._profiler.stop()
+        return self
+
+    def print(self, **kwargs):
+        """打印分析报告到终端"""
+        if self.name:
+            print(f"\n{'='*20} {self.name} {'='*20}")
+        print(self._profiler.output_text(
+            unicode=True,
+            color=True,
+            show_all=self.show_all,
+            timeline=self._timeline,
+            **kwargs
+        ))
+
+    def to_html(self) -> str:
+        """生成 HTML 报告"""
+        return self._profiler.output_html()
+
+    def to_text(self, **kwargs) -> str:
+        """生成文本报告"""
+        return self._profiler.output_text(
+            unicode=True,
+            show_all=self.show_all,
+            timeline=self._timeline,
+            **kwargs
+        )
+
+    def to_json(self) -> str:
+        """生成 JSON 报告（用于程序化分析）"""
+        import json
+        return json.dumps(self._profiler.last_session.frame_records, indent=2)
+
+    def save(self, path: str):
+        """
+        保存报告到文件
+
+        根据文件扩展名自动选择格式：
+        - .html -> HTML 交互式报告
+        - .txt -> 文本报告
+        - .json -> JSON 数据
+        """
+        path = Path(path)
+        suffix = path.suffix.lower()
+
+        if suffix == ".html":
+            content = self.to_html()
+        elif suffix == ".json":
+            content = self.to_json()
+        else:
+            content = self.to_text()
+
+        path.write_text(content, encoding="utf-8")
+        print(f"报告已保存: {path}")
+
+    def open_in_browser(self):
+        """在浏览器中打开交互式 HTML 报告"""
+        self._profiler.open_in_browser(timeline=self._timeline)
+
+    @property
+    def session(self):
+        """获取原始 session 对象用于高级操作"""
+        return self._profiler.last_session
+
+
+def profile(
+    func=None,
+    *,
+    show: bool = True,
+    show_all: bool = False,
+    save_to: Optional[str] = None,
+):
+    """
+    函数装饰器 - 分析函数性能
+
+    Example:
+        @profile
+        def slow_function():
+            ...
+
+        @profile(show=False, save_to="report.html")
+        def another_function():
+            ...
+    """
+    def decorator(fn):
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            with Profile(fn.__name__, show=show, show_all=show_all) as p:
+                result = fn(*args, **kwargs)
+            if save_to:
+                p.save(save_to)
+            return result
+
+        @functools.wraps(fn)
+        async def async_wrapper(*args, **kwargs):
+            with Profile(fn.__name__, show=show, show_all=show_all) as p:
+                result = await fn(*args, **kwargs)
+            if save_to:
+                p.save(save_to)
+            return result
+
+        import asyncio
+        if asyncio.iscoroutinefunction(fn):
+            return async_wrapper
+        return wrapper
+
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+class ScaleneProfile:
+    """
+    全面性能分析器 - 基于 Scalene
+
+    支持 CPU + 内存 + GPU 分析，自动检测内存泄漏。
+
+    Example:
+        # 基本使用
+        with ScaleneProfile("数据处理"):
+            process_data()
+
+        # 含内存分析
+        with ScaleneProfile("内存密集任务", memory=True, gpu=True):
+            train_model()
+
+        # 生成 HTML 报告
+        with ScaleneProfile("分析", output="report.html"):
+            heavy_work()
+
+    Note:
+        Scalene 使用采样分析，对于运行时间 < 1 秒的代码可能采样不足。
+        建议用于分析耗时较长的代码块。
+    """
+
+    def __init__(
+        self,
+        name: str = "",
+        *,
+        cpu: bool = True,
+        memory: bool = False,
+        gpu: bool = False,
+        output: Optional[str] = None,  # HTML 报告路径
+        reduced_profile: bool = False,  # 仅显示有性能问题的行
+    ):
+        if not SCALENE_AVAILABLE:
+            raise ImportError(
+                "scalene 未安装，请运行: pip install scalene"
+            )
+
+        self.name = name
+        self.cpu = cpu
+        self.memory = memory
+        self.gpu = gpu
+        self.output = output
+        self.reduced_profile = reduced_profile
+
+    def __enter__(self):
+        if self.name:
+            print(f"\n🔬 Scalene 分析开始: {self.name}")
+        scalene_profiler.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        scalene_profiler.stop()
+        if self.name:
+            print(f"✅ Scalene 分析完成: {self.name}")
+        return False
+
+    def start(self):
+        """手动启动分析"""
+        scalene_profiler.start()
+        return self
+
+    def stop(self):
+        """手动停止分析"""
+        scalene_profiler.stop()
+        return self
+
+
+def scalene_run(
+    script: str,
+    *args,
+    cpu: bool = True,
+    memory: bool = True,
+    gpu: bool = False,
+    output: Optional[str] = None,
+    reduced: bool = False,
+    **kwargs,
+) -> subprocess.CompletedProcess:
+    """
+    使用 Scalene 运行 Python 脚本（推荐方式）
+
+    Scalene 的完整功能需要从命令行启动，此函数封装了命令行调用。
+
+    Example:
+        # 分析脚本
+        scalene_run("train.py", "--epochs", "10", output="report.html")
+
+        # 分析模块
+        scalene_run("-m", "pytest", "tests/", memory=True)
+
+    Args:
+        script: Python 脚本路径或 -m 模块名
+        *args: 传递给脚本的参数
+        cpu: 是否分析 CPU（默认 True）
+        memory: 是否分析内存（默认 True）
+        gpu: 是否分析 GPU
+        output: HTML 报告输出路径
+        reduced: 仅显示有问题的行
+
+    Returns:
+        subprocess.CompletedProcess 对象
+    """
+    cmd = [sys.executable, "-m", "scalene"]
+
+    if not cpu:
+        cmd.append("--cpu-only")
+    if memory:
+        cmd.append("--memory")
+    if gpu:
+        cmd.append("--gpu")
+    if reduced:
+        cmd.append("--reduced-profile")
+    if output:
+        cmd.extend(["--html", "--outfile", output])
+
+    cmd.append("---")  # 分隔 Scalene 参数和脚本参数
+    cmd.append(script)
+    cmd.extend(args)
+
+    print(f"🚀 运行: {' '.join(cmd)}")
+    return subprocess.run(cmd, **kwargs)
+
+
+def scalene_profile(
+    func=None,
+    *,
+    memory: bool = False,
+    gpu: bool = False,
+):
+    """
+    Scalene 函数装饰器
+
+    注意：需要用 `scalene` 命令启动脚本才能生效。
+
+    Example:
+        @scalene_profile(memory=True)
+        def process_data():
+            ...
+
+        # 运行时使用: scalene script.py
+    """
+    def decorator(fn):
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            with ScaleneProfile(fn.__name__, memory=memory, gpu=gpu):
+                return fn(*args, **kwargs)
+        return wrapper
+
+    if func is not None:
+        return decorator(func)
+    return decorator

maque/performance/_stat_memory.py

@@ -0,0 +1,51 @@
+import time
+import psutil
+
+_scale = {'KB': 1024, 'MB': 1024 * 1024, 'GB': 1024 * 1024 * 1024, }
+
+
+def get_virtual_memory(unit='GB'):
+    memory_tuple = psutil.virtual_memory()
+    _key = unit.upper()
+    memory_dict = {
+        "total": (memory_tuple[0] / _scale[_key], _key),
+        "available": (memory_tuple[1] / _scale[_key], _key),
+        "percent": (memory_tuple[2], '%'),
+        "used": (memory_tuple[3] / _scale[_key], _key),
+        "free": (memory_tuple[4] / _scale[_key], _key),
+    }
+    return memory_dict
+
+
+def elapsed_since(start):
+    return time.strftime("%H:%M:%S", time.gmtime(time.time() - start))
+
+
+def get_process_memory(unit='MB'):
+    return psutil.Process().memory_info().rss / _scale[unit.upper()]
+
+
+def _profile(func):
+    """
+    注意，这个装饰器只能得到函数推出后的内存增量，而要函数内的内存增量不应用此装饰器
+    若要方便调试对函数的使用，请使用 memory_profiler这个包
+    from memory_profiler import profile
+    @profile
+    def func():
+        ...
+    """
+
+    def wrapper(*args, **kwargs):
+        _key = "MB"
+        mem_before = get_process_memory(_key)
+        start = time.time()
+        result = func(*args, **kwargs)
+        elapsed_time = elapsed_since(start)
+        mem_after = get_process_memory(_key)
+        print(
+            f"{func.__name__}: memory before: {mem_before:.2f} {_key}, after: {mem_after:.2f} {_key}, consumed: {mem_after - mem_before:.2f} {_key}; exec time: {elapsed_time}"
+        )
+        return result
+
+    return wrapper
+

maque/pipelines/__init__.py

@@ -0,0 +1,15 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+Pipeline 模块 - 提供端到端的便利流水线
+
+这些流水线是可选的便利层，封装常见的工作流模式。
+用户仍可以使用底层 API (embedding, retriever, clustering) 进行完全自定义。
+"""
+
+from .clustering import ClusteringPipeline
+
+__all__ = [
+    "ClusteringPipeline",
+]

maque/pipelines/clustering.py

@@ -0,0 +1,252 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+聚类分析流水线
+
+提供向量化 + 聚类的端到端便利接口。
+这是可选的高层封装，用户仍可使用底层 API 进行完全自定义。
+"""
+
+from pathlib import Path
+from typing import List, Optional, Union
+
+import numpy as np
+from loguru import logger
+
+try:
+    import chromadb
+except ImportError:
+    chromadb = None
+
+from ..embedding.base import BaseEmbedding
+from ..retriever import ChromaRetriever, Document
+from ..clustering import ClusterAnalyzer, ClusterResult
+
+
+class ClusteringPipeline:
+    """
+    向量化 + 聚类的端到端流水线
+
+    封装常见的工作流：数据加载 → 向量化 → 存储 → 聚类分析
+
+    这是可选的便利层，用户仍可使用底层组件：
+    - TextEmbedding: 向量化
+    - ChromaRetriever: 向量存储与检索
+    - ClusterAnalyzer: 聚类分析
+
+    Example:
+        >>> # 基础用法
+        >>> from maque.embedding import TextEmbedding
+        >>> from maque.pipelines import ClusteringPipeline
+        >>> from maque.clustering import ClusterAnalyzer
+        >>> from maque.retriever import Document
+        >>>
+        >>> # 创建流水线
+        >>> pipeline = ClusteringPipeline(
+        ...     embedding=TextEmbedding(base_url="http://localhost:8000", model="jina-v3"),
+        ...     persist_dir="./chroma_db",
+        ...     collection_name="my_data",
+        ... )
+        >>>
+        >>> # 向量化文档
+        >>> docs = [Document.text(content=text, id=f"doc_{i}") for i, text in enumerate(texts)]
+        >>> pipeline.build_vectors(docs, batch_size=32, skip_existing=True)
+        >>>
+        >>> # 聚类分析
+        >>> analyzer = ClusterAnalyzer(algorithm="hdbscan", min_cluster_size=15)
+        >>> result = pipeline.analyze(analyzer, output_dir="./results", name="my_cluster")
+
+        >>> # 高级用法：直接访问底层组件
+        >>> pipeline.retriever.search("查询文本", top_k=10)  # 使用检索功能
+        >>> pipeline.embedding.embed(["文本1", "文本2"])     # 直接向量化
+    """
+
+    def __init__(
+        self,
+        embedding: BaseEmbedding,
+        persist_dir: Union[str, Path] = None,
+        collection_name: str = "default",
+        distance_metric: str = "cosine",
+        retriever=None,
+    ):
+        """
+        初始化聚类流水线
+
+        Args:
+            embedding: Embedding 实例（TextEmbedding 或 MultiModalEmbedding）
+            persist_dir: ChromaDB 持久化目录（当 retriever=None 时必须提供）
+            collection_name: 集合名称（当 retriever=None 时使用）
+            distance_metric: 距离度量方式 (cosine/l2/ip)（当 retriever=None 时使用）
+            retriever: 可选，直接传入 Retriever 实例（ChromaRetriever 或 MilvusRetriever）
+                      传入后将忽略 persist_dir, collection_name, distance_metric 参数
+        """
+        self.embedding = embedding
+
+        # 支持传入 retriever 或自动创建 ChromaRetriever
+        if retriever is not None:
+            self.retriever = retriever
+            self.persist_dir = getattr(retriever, 'persist_dir', None)
+            self.collection_name = getattr(retriever, 'collection_name', 'unknown')
+            logger.info(f"初始化 ClusteringPipeline: {self.collection_name} (外部 retriever)")
+        else:
+            if persist_dir is None:
+                raise ValueError("persist_dir 必须提供（当 retriever=None 时）")
+            self.persist_dir = Path(persist_dir)
+            self.collection_name = collection_name
+            # 创建 ChromaRetriever（底层组件仍可访问）
+            self.retriever = ChromaRetriever(
+                embedding=embedding,
+                persist_dir=str(persist_dir),
+                collection_name=collection_name,
+                distance_metric=distance_metric,
+            )
+            logger.info(f"初始化 ClusteringPipeline: {collection_name} @ {persist_dir}")
+
+    def build_vectors(
+        self,
+        documents: List[Document],
+        batch_size: int = 32,
+        skip_existing: bool = True,
+        show_progress: bool = True,
+    ) -> int:
+        """
+        向量化并存储文档
+
+        Args:
+            documents: 文档列表
+            batch_size: 批处理大小
+            skip_existing: 是否跳过已存在的文档（增量更新）
+            show_progress: 是否显示进度条
+
+        Returns:
+            实际插入的文档数量
+
+        Example:
+            >>> from maque.retriever import Document
+            >>> docs = [Document.text(content=f"text {i}", id=f"doc_{i}") for i in range(100)]
+            >>> count = pipeline.build_vectors(docs, skip_existing=True)
+            >>> print(f"插入 {count} 个文档")
+        """
+        logger.info(f"开始向量化 {len(documents)} 个文档...")
+
+        count = self.retriever.upsert_batch(
+            documents=documents,
+            batch_size=batch_size,
+            skip_existing=skip_existing,
+            show_progress=show_progress,
+        )
+
+        logger.info(f"完成向量化，当前集合共 {self.retriever.count()} 个文档")
+        return count
+
+    def analyze(
+        self,
+        analyzer: ClusterAnalyzer,
+        output_dir: Optional[Union[str, Path]] = None,
+        name: Optional[str] = None,
+        sample_size: Optional[int] = None,
+        where: Optional[dict] = None,
+        where_document: Optional[dict] = None,
+    ) -> ClusterResult:
+        """
+        从 ChromaDB 加载并执行聚类分析
+
+        Args:
+            analyzer: ClusterAnalyzer 实例（预配置算法参数）
+            output_dir: 输出目录（保存 JSON 和可视化图片）
+            name: 结果文件名前缀（默认使用 collection_name）
+            sample_size: 限制加载的样本数量（用于大规模数据）
+            where: 元数据过滤条件
+            where_document: 文档内容过滤条件
+
+        Returns:
+            ClusterResult: 聚类分析结果
+
+        Example:
+            >>> # HDBSCAN 自动聚类
+            >>> analyzer = ClusterAnalyzer(algorithm="hdbscan", min_cluster_size=15)
+            >>> result = pipeline.analyze(analyzer, output_dir="./results")
+
+            >>> # K-Means 指定簇数
+            >>> analyzer = ClusterAnalyzer(algorithm="kmeans", n_clusters=20)
+            >>> result = pipeline.analyze(analyzer, output_dir="./results", sample_size=10000)
+
+            >>> # 过滤特定类别数据
+            >>> result = pipeline.analyze(
+            ...     analyzer,
+            ...     output_dir="./results",
+            ...     where={"category": "tech"},
+            ... )
+        """
+        logger.info("开始聚类分析...")
+
+        # 使用 ClusterAnalyzer 的 analyze_chroma 方法
+        result = analyzer.analyze_chroma(
+            persist_dir=self.persist_dir,
+            collection_name=self.collection_name,
+            output_dir=output_dir,
+            name=name or self.collection_name,
+            sample_size=sample_size,
+            where=where,
+            where_document=where_document,
+        )
+
+        logger.info("聚类分析完成")
+        return result
+
+    def load_vectors(
+        self,
+        sample_size: Optional[int] = None,
+        where: Optional[dict] = None,
+    ) -> tuple[List[str], np.ndarray, List[str]]:
+        """
+        从 ChromaDB 加载向量数据（供高级用户自定义分析）
+
+        Args:
+            sample_size: 限制加载的样本数量
+            where: 元数据过滤条件
+
+        Returns:
+            (ids, embeddings, documents) 元组
+
+        Example:
+            >>> ids, embeddings, docs = pipeline.load_vectors(sample_size=1000)
+            >>> print(f"加载 {len(embeddings)} 个向量，维度={embeddings.shape[1]}")
+            >>> # 用户可自定义分析逻辑
+            >>> from sklearn.cluster import KMeans
+            >>> kmeans = KMeans(n_clusters=10)
+            >>> labels = kmeans.fit_predict(embeddings)
+        """
+        if chromadb is None:
+            raise ImportError("需要安装 chromadb: pip install chromadb")
+
+        client = chromadb.PersistentClient(path=str(self.persist_dir))
+        collection = client.get_collection(self.collection_name)
+
+        query_kwargs = {"include": ["embeddings", "documents", "metadatas"]}
+        if sample_size is not None:
+            query_kwargs["limit"] = sample_size
+        if where is not None:
+            query_kwargs["where"] = where
+
+        results = collection.get(**query_kwargs)
+
+        ids = results["ids"]
+        embeddings = np.array(results["embeddings"])
+        documents = results["documents"]
+
+        logger.info(f"加载 {len(embeddings)} 个向量, 维度={embeddings.shape[1]}")
+        return ids, embeddings, documents
+
+    def count(self) -> int:
+        """返回集合中的文档数量"""
+        return self.retriever.count()
+
+    def __repr__(self) -> str:
+        return (
+            f"ClusteringPipeline("
+            f"collection={self.collection_name!r}, "
+            f"count={self.count()}, "
+            f"persist_dir={str(self.persist_dir)!r})"
+        )