sirchmunk 0.0.0__py3-none-any.whl → 0.0.1.post1__py3-none-any.whl

sirchmunk/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .search import AgenticSearch
+from .version import __version__
+
+__all__ = [
+    "__version__",
+    "AgenticSearch",
+]

sirchmunk/base.py

@@ -0,0 +1,17 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from abc import ABC, abstractmethod
+
+
+class BaseSearch(ABC):
+    """
+    Abstract base class for search implementations.
+    """
+
+    def __init__(self, *args, **kwargs): ...
+
+    @abstractmethod
+    def search(self, *args, **kwargs):
+        """
+        Perform a search based on the given query.
+        """
+        raise NotImplementedError("search method not implemented")

sirchmunk/insight/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""
+Data Insights Module to empower agents with corpus analysis capabilities.
+"""

sirchmunk/insight/text_insights.py

@@ -0,0 +1,292 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import re
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+from modelscope import snapshot_download
+
+from sirchmunk.llm.openai_chat import OpenAIChat
+from sirchmunk.llm.prompts import (
+    SNAPSHOT_KEYWORDS_EXTRACTION,
+    SNAPSHOT_TOC_EXTRACTION,
+)
+
+
+class KeyPhraseExtractor:
+    """
+    Key phrase extraction using sentence-transformers and scikit-learn.
+    """
+
+    def __init__(
+        self,
+        model_id: str = "AI-ModelScope/all-MiniLM-L6-v2",
+        device: Union[str, None] = None,
+        stop_words: Optional[List[str]] = None,
+        ngram_range: Tuple[int, int] = (1, 2),
+    ):
+        """
+        Key phrase extraction using sentence-transformers and scikit-learn.
+
+        Args:
+            model_id (str): model id from ModelScope.
+            device (str): device to load the model on ('cpu' or 'cuda' or 'mps').
+            stop_words (Optional[List[str]]): list of stop words to filter out.
+            ngram_range (Tuple[int, int]): n-gram range for candidate generation.
+        """
+        from sentence_transformers import SentenceTransformer
+
+        model_dir: str = snapshot_download(
+            model_id=model_id,
+            ignore_patterns=[
+                "openvino/*",
+                "onnx/*",
+                "pytorch_model.bin",
+                "rust_model.ot",
+                "tf_model.h5",
+            ],
+        )
+        self.model = SentenceTransformer(model_dir, device=device)
+
+        self.ngram_range = ngram_range
+        self.stop_words = set(stop_words) if stop_words else set()
+        if not self.stop_words:
+            self.stop_words = {
+                # Chinese stop words
+                "的",
+                "了",
+                "在",
+                "是",
+                "我",
+                "有",
+                "和",
+                "就",
+                "不",
+                "人",
+                "都",
+                "一",
+                "一个",
+                "上",
+                # English stop words
+                "the",
+                "a",
+                "an",
+                "and",
+                "or",
+                "but",
+                "in",
+                "on",
+                "at",
+                "to",
+                "for",
+                "of",
+                "with",
+                "by",
+                "is",
+                "are",
+                "was",
+                "were",
+                "be",
+                "been",
+                "have",
+                "has",
+                "had",
+                "do",
+                "does",
+            }
+
+    @staticmethod
+    def _preprocess(text: str) -> str:
+        """
+        Keep only Chinese, English, digits, and spaces; convert to lowercase.
+        """
+        text = re.sub(r"[^\w\s\u4e00-\u9fa5]", " ", text.lower())
+        return " ".join(text.split())
+
+    def _generate_candidates(self, docs: List[str]) -> List[str]:
+        """
+        Generate candidate key phrases using CountVectorizer.
+        """
+        from sklearn.feature_extraction.text import CountVectorizer
+
+        n_docs = len(docs)
+        if n_docs == 0:
+            return []
+
+        vectorizer_params = {
+            "ngram_range": self.ngram_range,
+            "stop_words": list(self.stop_words),
+            "token_pattern": r"(?u)\b[\w\u4e00-\u9fa5]+\b",
+            "lowercase": True,
+            "min_df": 1,
+        }
+        if n_docs > 3:
+            vectorizer_params["max_df"] = 0.85
+
+        try:
+            vectorizer = CountVectorizer(**vectorizer_params)
+            vectorizer.fit(docs)
+        except ValueError as e:
+            if "max_df corresponds to" in str(e):
+                # fallback
+                vectorizer_params.pop("max_df", None)
+                vectorizer = CountVectorizer(**vectorizer_params)
+                vectorizer.fit(docs)
+            else:
+                raise
+
+        candidates = vectorizer.get_feature_names_out().tolist()
+        candidates = [c.strip() for c in candidates if len(re.sub(r"\s+", "", c)) > 1]
+
+        return candidates
+
+    def extract(
+        self,
+        contents: List[str],
+        top_k: int = 20,
+        use_mmr: bool = True,
+        diversity: float = 0.5,
+        candidates: Optional[List[str]] = None,
+        batch_size: int = 32,
+    ) -> List[Tuple[str, float]]:
+        """
+        Extract key phrases from the given contents.
+
+        contents (List[str]): List of document strings.
+        top_k (int): Number of top key phrases to return.
+        use_mmr (bool): Whether to use MMR for diversity.
+        diversity (float): Diversity factor for MMR (0 to 1).
+        candidates (Optional[List[str]]): Predefined candidate phrases/keywords.
+        batch_size (int): Batch size for encoding.
+        """
+        if not contents:
+            return []
+
+        # Step 1: merge and preprocess documents
+        doc = " ".join(contents)
+        doc = self._preprocess(doc)
+
+        # Step 2: Generate candidates if not provided
+        if candidates is None:
+            candidates = self._generate_candidates([doc])
+        if not candidates:
+            return []
+
+        # Step 3: Encode document and candidates
+        sentences = [doc] + candidates
+        embeddings = self.model.encode(
+            sentences,
+            batch_size=batch_size,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+            normalize_embeddings=True,
+        )
+        doc_emb = embeddings[0]
+        cand_embs = embeddings[1:]
+
+        # Step 4: Compute similarities
+        similarities = (doc_emb @ cand_embs.T).flatten()  # shape: (n_candidates,)
+
+        if not use_mmr:
+            idx = np.argsort(similarities)[::-1][:top_k]
+            return [(candidates[i], float(similarities[i])) for i in idx]
+
+        # Step 5 (Optional): MMR diversity selection
+        selected_idx = [int(np.argmax(similarities))]
+        remaining_idx = list(set(range(len(candidates))) - set(selected_idx))
+
+        while len(selected_idx) < top_k and remaining_idx:
+            candidate_scores = []
+            for i in remaining_idx:
+                sim_to_doc = similarities[i]
+                sim_to_selected = np.max(cand_embs[selected_idx] @ cand_embs[i])
+                mmr_score = diversity * sim_to_doc - (1 - diversity) * sim_to_selected
+                candidate_scores.append((i, mmr_score))
+
+            # Select the candidate with the highest MMR score
+            next_i = max(candidate_scores, key=lambda x: x[1])[0]
+            selected_idx.append(next_i)
+            remaining_idx.remove(next_i)
+
+        return [(candidates[i], float(similarities[i])) for i in selected_idx]
+
+
+class TextInsights:
+    """
+    Text insights: key information extraction.
+    """
+
+    def __init__(self, llm: Optional[OpenAIChat] = None, **kwargs):
+        """
+        Initialize TextInsights with an optional LLM instance.
+
+        Args:
+            llm (Optional[OpenAIChat]): An instance of OpenAIChat.
+            **kwargs: Additional keyword arguments for KeyPhraseExtractor.
+        """
+        self._llm = llm
+
+        self._key_phrase_extractor = (
+            KeyPhraseExtractor(
+                model_id="AI-ModelScope/all-MiniLM-L6-v2",
+                device=None,
+                ngram_range=(1, 2),
+            )
+            if self._llm is None
+            else None
+        )
+
+        self._kwargs = kwargs
+
+    def extract_phrase(self, contents: List[str], max_num: int = 20) -> List[str]:
+        """
+        Extract key phrases from the given contents.
+
+        Args:
+            contents (List[str]): List of document strings.
+            max_num (int): Maximum number of key phrases to extract.
+
+        Returns:
+            List[str]: Extracted key phrases.
+        """
+
+        if self._llm is not None:
+            prompt = SNAPSHOT_KEYWORDS_EXTRACTION.format(
+                document_content="\n\n".join(contents), max_num=max_num
+            )
+            messages = [{"role": "user", "content": prompt}]
+            response: str = self._llm.chat(
+                messages=messages,
+                stream=True,
+            )
+
+            results = [
+                phrase.strip().lower()
+                for phrase in response.split(",")
+                if phrase.strip()
+            ]
+
+        else:
+            extracted_phrases = self._key_phrase_extractor.extract(
+                contents=contents, top_k=max_num, **self._kwargs
+            )
+            results = [phrase for phrase, _ in extracted_phrases]
+
+        return results
+
+    def extract_toc(self, contents: List[str]) -> str:
+        """
+        Extract the `Table of contents from input document.`
+        """
+
+        if self._llm is None:
+            return ""
+
+        prompt = SNAPSHOT_TOC_EXTRACTION.format(document_content="\n\n".join(contents))
+        messages = [{"role": "user", "content": prompt}]
+
+        response: str = self._llm.chat(
+            messages=messages,
+            stream=True,
+        )
+
+        return response

sirchmunk/learnings/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.