hamtaa-texttools 0.1.43__py3-none-any.whl

texttools/handlers/categorizer/categorizer.py

@@ -0,0 +1,61 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+
+from elasticsearch import Elasticsearch, helpers
+
+
+class ResultHandler(ABC):
+    """
+    Abstract base class for all result handlers.
+    Implement the handle() method to define custom handling logic.
+    """
+
+    @abstractmethod
+    def handle(self, results: dict[str, Enum]) -> None:
+        """
+        Process the categorization results.
+
+        Args:
+            results (dict[str, Enum]): A dictionary mapping text (or IDs) to categories.
+        """
+        pass
+
+
+class NoOpResultHandler(ResultHandler):
+    """
+    A result handler that does nothing.
+    Useful as a default when no other handler is provided.
+    """
+
+    def handle(self, results) -> None:
+        pass
+
+
+class PrintResultHandler(ResultHandler):
+    """
+    A simple handler that prints results to the console.
+    Useful for debugging or local tests.
+    """
+
+    def handle(self, results) -> None:
+        for key, value in results.items():
+            print(f"Text ID: {key}, Category: {value.name}")
+
+
+class SaveToElasticResultHandler(ResultHandler):
+    """
+    A simple handler that saves results to an elastic index.
+    """
+
+    def __init__(self, es_client: Elasticsearch, index_name: str):
+        self.es_client = es_client
+        self.index_name = index_name
+
+    def handle(self, results):
+        documents = [
+            {"TextID": key, "Category": value.name} for key, value in results.items()
+        ]
+
+        actions = [{"_index": self.index_name, "_source": doc} for doc in documents]
+
+        helpers.bulk(self.es_client, actions)

texttools/handlers/handlers.py

@@ -0,0 +1,88 @@
+import json
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class ResultHandler(ABC):
+    """
+    Abstract base class for all result handlers.
+    Implement the handle() method to define custom handling logic.
+    """
+
+    @abstractmethod
+    def handle(self, results: dict[str, Enum]) -> None:
+        """
+        Process the categorization results.
+
+        Args:
+            results (dict[str, Enum]): A dictionary mapping text (or IDs) to categories.
+        """
+        pass
+
+
+class NoOpResultHandler(ResultHandler):
+    """
+    A result handler that does nothing.
+    Useful as a default when no other handler is provided.
+    """
+
+    def handle(self, results: dict[str, Enum]) -> None:
+        pass
+
+
+class PrintResultHandler(ResultHandler):
+    """
+    A simple handler that prints results to the console.
+    Useful for debugging or local tests.
+    """
+
+    def handle(self, results: dict[str, Enum]) -> None:
+        for key, value in results.items():
+            print(f"Text ID: {key}, Category: {value.name}")
+
+
+class SaveToFileResultHandler(ResultHandler):
+    """
+    A handler that saves each question + result pair to a CSV-like file,
+    serializing whatever the result object is.
+    """
+
+    def __init__(self, file_path: str):
+        self.file_path = file_path
+
+    def handle(self, results: dict[str, Any]) -> None:
+        """
+        Expects `results` to be a dict with at least:
+          - "question": the original input text
+          - "result":   the classification output (bool, BaseModel, dict, str, etc.)
+
+        Appends one line per call:
+            question_text,serialized_result
+        """
+
+        # Helper to turn anything into a JSON/string
+        def serialize(val: Any) -> str:
+            if isinstance(val, BaseModel):
+                return val.model_dump_json()
+            try:
+                return json.dumps(val)
+            except (TypeError, ValueError):
+                return str(val)
+
+        q = results.get("question", "")
+        r = results.get("result", results)
+        line = f"{q},{serialize(r)}\n"
+
+        with open(self.file_path, "a", encoding="utf-8") as f:
+            f.write(line)
+
+
+# You can add more handlers here as needed:
+# - ElasticSearchResultHandler
+# - DatabaseResultHandler
+# - KafkaResultHandler
+# - NATSResultHandler
+# etc.

texttools/tools/__init__.py

@@ -0,0 +1,33 @@
+from .categorizer import EmbeddingCategorizer, GemmaCategorizer, LLMCategorizer
+from .keyword_extractor import GemmaKeywordExtractor
+from .ner import GemmaNERExtractor
+from .question_detector import GemmaQuestionDetector, LLMQuestionDetector
+from .question_generator import GemmaQuestionGenerator
+from .reranker import GemmaReranker, GemmaScorer, GemmaSorter
+from .rewriter import GemmaQuestionRewriter, RewriteMode
+from .merger import GemmaQuestionMerger, MergingMode
+from .subject_to_question import GemmaQuestionGeneratorFromSubject
+from .summarizer import GemmaSummarizer, LLMSummarizer
+from .translator import GemmaTranslator
+
+__all__ = [
+    "EmbeddingCategorizer",
+    "GemmaCategorizer",
+    "LLMCategorizer",
+    "GemmaTranslator",
+    "GemmaSummarizer",
+    "LLMSummarizer",
+    "GemmaNERExtractor",
+    "GemmaQuestionDetector",
+    "LLMQuestionDetector",
+    "GemmaQuestionGenerator",
+    "GemmaScorer",
+    "GemmaSorter",
+    "GemmaReranker",
+    "GemmaQuestionRewriter",
+    "RewriteMode",
+    "GemmaKeywordExtractor",
+    "GemmaQuestionGeneratorFromSubject",
+    "GemmaQuestionMerger",
+    "MergingMode",
+]

texttools/tools/categorizer/__init__.py

@@ -0,0 +1,2 @@
+from .encoder_model import *
+from .llm import *

texttools/tools/categorizer/encoder_model/__init__.py

@@ -0,0 +1 @@
+from .encoder_vectorizer import EmbeddingCategorizer

texttools/tools/categorizer/encoder_model/encoder_vectorizer.py

@@ -0,0 +1,51 @@
+from enum import Enum
+from typing import Any, Optional
+
+import numpy as np
+
+from texttools.base import BaseCategorizer
+from texttools.handlers import ResultHandler
+
+
+class EmbeddingCategorizer(BaseCategorizer):
+    """
+    Uses pre-stored embeddings on each Enum member.
+    """
+
+    def __init__(
+        self,
+        categories: Enum,
+        embedding_model: Any,
+        handlers: Optional[list[ResultHandler]] = None,
+    ):
+        """
+        :param categories: your Enum class, whose members have `.embeddings`
+        :param embedding_model: something with `.encode(text: str) -> list[float]`
+        """
+        super().__init__(categories, handlers)
+        self.embedding_model = embedding_model
+
+    def categorize(self, text: str) -> Enum:
+        # 1. Preprocess
+        text = self.preprocess(text)
+
+        # 2. Encode the text
+        vec = np.array(self.embedding_model.encode(text), dtype=float)
+
+        # 3. Find best category
+        best_cat = None
+        best_score = -1.0
+
+        for cat in self.categories:
+            for proto in cat.embeddings:
+                score = self._cosine_similarity(vec, proto)
+                if score > best_score:
+                    best_score = score
+                    best_cat = cat
+
+        self._dispatch({text: best_cat})
+        return best_cat
+
+    @staticmethod
+    def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
+        return float(a.dot(b) / (np.linalg.norm(a) * np.linalg.norm(b)))

texttools/tools/categorizer/llm/__init__.py

@@ -0,0 +1,2 @@
+from .openai_categorizer import LLMCategorizer
+from .gemma_categorizer import GemmaCategorizer

texttools/tools/categorizer/llm/gemma_categorizer.py

@@ -0,0 +1,169 @@
+from typing import Any, Literal, Optional
+
+from openai import OpenAI
+from pydantic import BaseModel
+
+from texttools.base.base_categorizer import BaseCategorizer
+from texttools.formatter.gemma3_formatter import Gemma3Formatter
+from texttools.handlers import ResultHandler
+
+
+class Output(BaseModel):
+    reason: str
+    main_tag: Literal[
+        "باورهای دینی",
+        "اخلاق اسلامی",
+        "احکام و فقه",
+        "تاریخ اسلام و شخصیت ها",
+        "منابع دینی",
+        "دین و جامعه/سیاست",
+        "عرفان و معنویت",
+        "هیچکدام",
+    ] = None
+
+
+class GemmaCategorizer(BaseCategorizer):
+    """
+    Categorizer for Gemma-style models. It requires a predefined Enum of categories
+    to choose from and returns an Enum member.
+    Outputs JSON with a single string field: {"category": "..."}.
+
+    Allows optional extra instructions via `prompt_template`.
+    """
+
+    def __init__(
+        self,
+        client: OpenAI,
+        *,
+        model: str,
+        output_structure: BaseModel = Output,
+        chat_formatter: Optional[Any] = None,
+        use_reason: bool = False,
+        temperature: float = 0.0,
+        prompt_template: Optional[str] = None,
+        handlers: Optional[list[ResultHandler]] = None,
+        **client_kwargs: Any,
+    ):
+        super().__init__(handlers=handlers)
+        self.client = client
+        self.model = model
+        self.temperature = temperature
+        self.client_kwargs = client_kwargs
+        self.output_structure = output_structure
+        self.chat_formatter = chat_formatter or Gemma3Formatter()
+
+        self.use_reason = use_reason
+        self.prompt_template = prompt_template
+
+    def _build_messages(
+        self, text: str, reason: Optional[str] = None
+    ) -> list[dict[str, str]]:
+        """
+        Builds the message list for the LLM API call for categorization.
+        """
+        clean_text = self.preprocess(text)
+
+        messages: list[dict[str, str]] = []
+
+        if self.prompt_template:
+            messages.append({"role": "user", "content": self.prompt_template})
+
+        if reason:
+            messages.append(
+                {"role": "user", "content": f"Based on this analysis: {reason}"}
+            )
+
+        messages.append(
+            {
+                "role": "user",
+                "content": """
+                تو یک متخصص علوم دینی هستی
+                من به عنوان کاربر یک متن به تو میدم و از تو میخوام که
+                اون متن رو در یکی از دسته بندی های زیر طبقه بندی کنی
+                
+        "باورهای دینی",
+        "اخلاق اسلامی",
+        "احکام و فقه",
+        "تاریخ اسلام و شخصیت ها",
+        "منابع دینی",
+        "دین و جامعه/سیاست",
+        "عرفان و معنویت",
+        "هیچکدام",
+
+                
+                در خروجی که از تو خواسته شده بخشی با عنوان reason وجود دارد
+                در اون بخش، دلیل انتخاب دسته بندی رو به صورت خلاصه بیاور
+
+
+                متنی که باید طبقه بندی کنی:
+                
+                
+                
+                """,
+            }
+        )
+        messages.append({"role": "user", "content": clean_text})
+        restructured = self.chat_formatter.format(messages=messages)
+
+        return restructured
+
+    def _reason(self, text: str) -> str:
+        """
+        Internal reasoning step to help the model analyze the text for categorization.
+        """
+        messages = [
+            {
+                "role": "user",
+                "content": """
+                هدف ما طبقه بندی متن هست
+                متن رو بخون و ایده اصلی و آنالیزی کوتاه از اون رو ارائه بده
+                
+                بسیار خلاصه باشه خروجی تو
+                نهایتا 20 کلمه
+                    """,
+            },
+            {
+                "role": "user",
+                "content": f"""
+                    {text}
+                    """,
+            },
+        ]
+
+        restrucruted = self.chat_formatter.format(messages=messages)
+
+        resp = self.client.chat.completions.create(
+            model=self.model,
+            messages=restrucruted,
+            temperature=self.temperature,
+            **self.client_kwargs,
+        )
+
+        reason_summary = resp.choices[0].message.content.strip()
+        return reason_summary
+
+    def categorize(self, text: str):
+        """
+        Categorizes `text` by selecting an appropriate member from the predefined Enum.
+        Optionally uses an internal reasoning step for better accuracy.
+        """
+        reason_summary = None
+        if self.use_reason:
+            reason_summary = self._reason(text)
+
+        messages = self._build_messages(text, reason_summary)
+        completion = self.client.beta.chat.completions.parse(
+            model=self.model,
+            messages=messages,
+            response_format=Output,
+            temperature=self.temperature,
+            extra_body=dict(guided_decoding_backend="auto"),
+            **self.client_kwargs,
+        )
+        message = completion.choices[0].message
+
+        category_name = message.parsed.main_tag
+
+        # dispatch and return - Note: _dispatch expects dict
+        self._dispatch(results={"main_tag": category_name})
+        return {"main_tag": category_name}

texttools/tools/categorizer/llm/openai_categorizer.py

@@ -0,0 +1,80 @@
+from enum import Enum
+from typing import Any, Optional
+
+from openai import OpenAI
+from pydantic import BaseModel, create_model
+
+from texttools.base import BaseCategorizer
+from texttools.handlers import NoOpResultHandler
+
+
+class LLMCategorizer(BaseCategorizer):
+    """
+    LLM-based categorizer using OpenAI's client.responses.parse
+    for Structured Outputs (Pydantic models).
+    """
+
+    def __init__(
+        self,
+        client: OpenAI,
+        categories: Enum,
+        *,
+        model: str,
+        temperature: float = 0.0,
+        prompt_template: str = None,
+        handlers: Optional[list[NoOpResultHandler]] = None,
+        **client_kwargs: Any,
+    ):
+        """
+        :param client: an instantiated OpenAI client
+        :param categories: an Enum class of allowed categories
+        :param model: the model name (e.g. "gpt-4o-2024-08-06")
+        :param temperature: sampling temperature
+        :param prompt_template: override default prompt instructions
+        :param handlers: list of handler instances to process the output
+        :param client_kwargs: any other OpenAI kwargs (e.g. `max_tokens`, `top_p`, etc.)
+        """
+        super().__init__(categories, handlers=handlers)
+        self.client = client
+        self.model = model
+        self.temperature = temperature
+        self.client_kwargs = client_kwargs
+
+        self.prompt_template = prompt_template or (
+            "You are a text classifier. Choose exactly one category from the list."
+        )
+
+        self._OutputModel = create_model(
+            "CategorizationOutput",
+            category=(self.categories, ...),
+        )
+
+    def _build_messages(self, text: str) -> list[dict[str, str]]:
+        """
+        Builds the message list for the OpenAI API based on the input text.
+        """
+        clean = self.preprocess(text)
+        return [
+            {"role": "system", "content": self.prompt_template},
+            {"role": "user", "content": clean},
+        ]
+
+    def categorize(self, text: str) -> Enum:
+        """
+        Categorizes the input text using OpenAI API and processes it using handlers.
+        """
+        msgs = self._build_messages(text)
+
+        resp = self.client.responses.parse(
+            model=self.model,
+            input=msgs,
+            text_format=self._OutputModel,
+            temperature=self.temperature,
+            **self.client_kwargs,
+        )
+
+        output: BaseModel = resp.output_parsed
+
+        self._dispatch({"text": text, "category": output.category})
+
+        return output.category

texttools/tools/keyword_extractor/__init__.py

@@ -0,0 +1 @@
+from texttools.tools.keyword_extractor.gemma_extractor import GemmaKeywordExtractor

texttools/tools/keyword_extractor/gemma_extractor.py

@@ -0,0 +1,138 @@
+from typing import Any, Optional
+
+from openai import OpenAI
+from pydantic import BaseModel
+
+from texttools.base.base_keyword_extractor import BaseKeywordExtractor
+from texttools.formatter import Gemma3Formatter
+
+
+class Output(BaseModel):
+    keywords: list
+
+
+class GemmaKeywordExtractor(BaseKeywordExtractor):
+    """
+    Keyword extractor for Gemma-style models with optional reasoning step.
+    Outputs JSON with a single array field: {"keywords": ["keyword1", "keyword2", ...]}.
+
+    Allows optional extra instructions via `prompt_template`.
+    """
+
+    def __init__(
+        self,
+        client: OpenAI,
+        *,
+        model: str,
+        use_reason: bool = False,
+        chat_formatter: Optional[Any] = None,
+        temperature: float = 0.0,
+        prompt_template: str = None,
+        handlers: list[Any] = None,
+        **client_kwargs: Any,
+    ):
+        super().__init__(handlers)
+        self.client = client
+        self.model = model
+        self.temperature = temperature
+        self.client_kwargs = client_kwargs
+        self.chat_formatter = chat_formatter or Gemma3Formatter()
+        self.use_reason = use_reason
+        self.prompt_template = prompt_template
+
+        self.output = Output
+
+    def _build_messages(
+        self, text: str, reason: Optional[str] = None
+    ) -> list[dict[str, str]]:
+        clean_text = self.preprocess(text)
+
+        messages: list[dict[str, str]] = []
+
+        if self.prompt_template:
+            messages.append({"role": "user", "content": self.prompt_template})
+
+        if reason:  # Include the reason if available
+            messages.append(
+                {"role": "user", "content": f"Based on this analysis: {reason}"}
+            )
+
+        messages.append(
+            {
+                "role": "user",
+                "content": "Extract the most relevant keywords from the following text. Provide them as a list of strings.",
+            }
+        )
+        messages.append({"role": "user", "content": clean_text})
+
+        # Ensure the schema is dumped as a valid JSON string
+        schema_instr = f"Respond only in JSON format: {self.output.model_dump_json()}"
+        messages.append({"role": "user", "content": schema_instr})
+
+        # Deprecated
+        # messages.append(
+        #     {"role": "assistant", "content": "{"}
+        # )  # Start with '{' to hint JSON
+
+        messages = self.chat_formatter.format(messages=messages)
+
+        return messages
+
+    def _reason(self, text: str) -> str:
+        """
+        Internal reasoning step to help the model identify potential keywords.
+        """
+        messages = [
+            {
+                "role": "user",
+                "content": """
+                    Analyze the following text to identify its main topics, concepts, and important terms.
+                    Provide a concise summary of your findings that will help in extracting relevant keywords.
+                    """,
+            },
+            {
+                "role": "user",
+                "content": f"""
+                    {text}
+                    """,
+            },
+        ]
+        messages = self.chat_formatter.format(messages=messages)
+
+        resp = self.client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            temperature=self.temperature,
+            **self.client_kwargs,
+        )
+
+        reason_summary = resp.choices[0].message.content.strip()
+        return reason_summary
+
+    def extract_keywords(self, text: str) -> list[str]:
+        """
+        Extracts keywords from `text`.
+        Optionally uses an internal reasoning step for better accuracy.
+        """
+        reason_summary = None
+        if self.use_reason:
+            reason_summary = self._reason(text)
+
+        messages = self._build_messages(text, reason_summary)
+
+        completion = self.client.beta.chat.completions.parse(
+            model=self.model,
+            messages=messages,
+            response_format=Output,
+            temperature=self.temperature,
+            extra_body=dict(guided_decoding_backend="auto"),
+            **self.client_kwargs,
+        )
+
+        message = completion.choices[0].message
+
+        keywords = message.parsed.keywords
+
+        # dispatch and return
+        self._dispatch({"original_text": text, "keywords": keywords})
+        return keywords

texttools/tools/merger/__init__.py

@@ -0,0 +1,2 @@
+from .gemma_question_merger import GemmaQuestionMerger
+from texttools.base.base_question_merger import MergingMode