gap-finder 0.1.0__py3-none-any.whl

gap_finder/__init__.py

@@ -0,0 +1,31 @@
+from typing import TYPE_CHECKING
+
+from gap_finder._auth import MissingAPIKeyError
+from gap_finder.model_types import (
+    EmbeddingModelInput,
+    LanguageModelInput,
+    SentenceTransformerLike,
+    TextGenerationPipelineLike,
+    TopicEmbeddingModelInput,
+)
+
+if TYPE_CHECKING:
+    from gap_finder.main import GapFinder
+
+__all__ = [
+    "EmbeddingModelInput",
+    "GapFinder",
+    "LanguageModelInput",
+    "MissingAPIKeyError",
+    "SentenceTransformerLike",
+    "TextGenerationPipelineLike",
+    "TopicEmbeddingModelInput",
+]
+
+
+def __getattr__(name: str):
+    if name == "GapFinder":
+        from gap_finder.main import GapFinder
+
+        return GapFinder
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

gap_finder/_auth.py

@@ -0,0 +1,18 @@
+import os
+
+
+class MissingAPIKeyError(RuntimeError):
+    """Raised when no API key was provided or configured."""
+
+
+def resolve_api_key(
+    explicit: str | None = None,
+    env_var: str = "GAP_FINDER_API_KEY",
+) -> str:
+    key = explicit or os.getenv(env_var)
+    if not key:
+        raise MissingAPIKeyError(
+            f"No API key found. Pass api_key=... or set the environment "
+            f"variable {env_var}."
+        )
+    return key

gap_finder/embedding_model/__init__.py

@@ -0,0 +1,40 @@
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+from gap_finder.embedding_model.client import OpenAIEmbeddingClient
+
+if TYPE_CHECKING:
+    from gap_finder.embedding_model.embedding_model import EmbeddingModel
+    from gap_finder.embedding_model.factory import create_embedding_model
+    from gap_finder.embedding_model.local_embedding_model import LocalEmbeddingModel
+    from gap_finder.embedding_model.openai_embedding_model import OpenAIEmbeddingModel
+
+__all__ = [
+    "EmbeddingModel",
+    "LocalEmbeddingModel",
+    "OpenAIEmbeddingClient",
+    "OpenAIEmbeddingModel",
+    "create_embedding_model",
+]
+
+
+def __getattr__(name: str):
+    if name == "EmbeddingModel":
+        from gap_finder.embedding_model.embedding_model import EmbeddingModel
+
+        return EmbeddingModel
+    if name == "create_embedding_model":
+        from gap_finder.embedding_model.factory import create_embedding_model
+
+        return create_embedding_model
+    if name == "LocalEmbeddingModel":
+        from gap_finder.embedding_model.local_embedding_model import LocalEmbeddingModel
+
+        return LocalEmbeddingModel
+    if name == "OpenAIEmbeddingModel":
+        from gap_finder.embedding_model.openai_embedding_model import OpenAIEmbeddingModel
+
+        return OpenAIEmbeddingModel
+    if name == "factory":
+        return import_module("gap_finder.embedding_model.factory")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

gap_finder/embedding_model/client.py

@@ -0,0 +1,24 @@
+from openai import OpenAI
+
+from gap_finder._auth import resolve_api_key
+
+
+class OpenAIEmbeddingClient:
+    """An OpenAI client paired with the embedding model it should call."""
+
+    def __init__(
+        self,
+        model_name: str,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        client: OpenAI | None = None,
+    ) -> None:
+        self.model_name = model_name
+        self.client = (
+            client
+            if client is not None
+            else OpenAI(
+                base_url=base_url,
+                api_key=resolve_api_key(api_key),
+            )
+        )

gap_finder/embedding_model/embedding_model.py

@@ -0,0 +1,37 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+
+
+class EmbeddingModel(ABC):
+    """Common interface for local and API-backed text embedding models."""
+
+    @abstractmethod
+    def encode(
+        self,
+        texts: list[str],
+        show_progress_bar: bool = False,
+        batch_size: int = 32,
+    ) -> np.ndarray:
+        """Encode text into an embedding matrix."""
+        raise NotImplementedError
+
+    def embed(self, documents: list[str], verbose: bool = False) -> np.ndarray:
+        """Embed documents using BERTopic-compatible arguments."""
+        return self.encode(
+            documents,
+            show_progress_bar=verbose,
+            batch_size=32,
+        )
+
+    def activate(self) -> None:
+        """Make internally owned resources ready for use."""
+        return
+
+    def deactivate(self) -> None:
+        """Optionally offload internally owned resources between steps."""
+        return
+
+    def cleanup(self) -> None:
+        """Release internally owned resources after the pipeline finishes."""
+        return

gap_finder/embedding_model/factory.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from gap_finder.embedding_model.client import OpenAIEmbeddingClient
+from gap_finder.embedding_model.embedding_model import EmbeddingModel
+from gap_finder.model_types import EmbeddingModelInput
+
+
+def create_embedding_model(
+    embedding_model: EmbeddingModelInput,
+) -> EmbeddingModel:
+    """Create an embedding wrapper from a model name or configured client."""
+    if isinstance(embedding_model, EmbeddingModel):
+        return embedding_model
+
+    if isinstance(embedding_model, OpenAIEmbeddingClient):
+        from gap_finder.embedding_model.openai_embedding_model import OpenAIEmbeddingModel
+
+        return OpenAIEmbeddingModel(client=embedding_model)
+    if isinstance(embedding_model, str):
+        from gap_finder.embedding_model.local_embedding_model import LocalEmbeddingModel
+
+        return LocalEmbeddingModel(client=embedding_model)
+
+    from sentence_transformers import SentenceTransformer
+
+    if isinstance(embedding_model, SentenceTransformer):
+        from gap_finder.embedding_model.local_embedding_model import LocalEmbeddingModel
+
+        return LocalEmbeddingModel(client=embedding_model)
+    raise TypeError(
+        "embedding_model must be an EmbeddingModel, model name, "
+        "SentenceTransformer, or "
+        f"OpenAIEmbeddingClient; got {type(embedding_model).__name__}."
+    )

gap_finder/embedding_model/local_embedding_model.py

@@ -0,0 +1,92 @@
+import gc
+from typing import Any
+
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+from gap_finder.embedding_model.embedding_model import EmbeddingModel
+
+
+class LocalEmbeddingModel(EmbeddingModel):
+    """Embedding model backed by a local SentenceTransformer."""
+
+    def __init__(
+        self,
+        client: str | SentenceTransformer,
+    ) -> None:
+        super().__init__()
+        self._owns_client = isinstance(client, str)
+        self.model_name = (
+            client
+            if isinstance(client, str)
+            else getattr(client, "model_name_or_path", None)
+        )
+        self.client: SentenceTransformer | None = (
+            None if isinstance(client, str) else client
+        )
+        self._target_device: Any = None
+        self._active = not self._owns_client
+        if self.client is not None:
+            self._target_device = getattr(self.client, "device", None)
+
+    def activate(self) -> None:
+        """Load or reactivate an internally owned SentenceTransformer."""
+        if not self._owns_client or self._active:
+            return
+        if self.client is None:
+            if not isinstance(self.model_name, str):
+                raise ValueError("Local embedding model has no model name.")
+            self.client = SentenceTransformer(self.model_name)
+            self._target_device = getattr(self.client, "device", None)
+        elif self._target_device is not None and hasattr(self.client, "to"):
+            self.client.to(self._target_device)
+        self._active = True
+
+    def _require_client(self) -> SentenceTransformer:
+        self.activate()
+        if self.client is None:
+            raise RuntimeError("Local embedding model could not be activated.")
+        return self.client
+
+    def encode(
+        self,
+        texts: list[str],
+        show_progress_bar: bool = False,
+        batch_size: int = 32,
+    ) -> np.ndarray:
+        if batch_size < 1:
+            raise ValueError("batch_size must be at least 1.")
+
+        embeddings = self._require_client().encode(
+            texts,
+            show_progress_bar=show_progress_bar,
+            batch_size=batch_size,
+        )
+        return np.asarray(embeddings, dtype=np.float32)
+
+    def deactivate(self) -> None:
+        """Offload an internally owned model while retaining it for reuse."""
+        if not self._owns_client or self.client is None or not self._active:
+            return
+
+        try:
+            import torch
+        except ImportError:
+            gc.collect()
+            return
+
+        if hasattr(self.client, "to"):
+            self.client.to(torch.device("cpu"))
+        self._active = False
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    def cleanup(self) -> None:
+        """Release an internally owned model after the pipeline finishes."""
+        if not self._owns_client or self.client is None:
+            return
+        self.deactivate()
+        self.client = None
+        self._active = False
+        gc.collect()

gap_finder/embedding_model/openai_embedding_model.py

@@ -0,0 +1,43 @@
+import numpy as np
+from tqdm import tqdm
+
+from gap_finder.embedding_model.client import OpenAIEmbeddingClient
+from gap_finder.embedding_model.embedding_model import EmbeddingModel
+
+
+class OpenAIEmbeddingModel(EmbeddingModel):
+    """Embedding model backed by an OpenAI-compatible embeddings API."""
+
+    def __init__(
+        self,
+        client: OpenAIEmbeddingClient,
+    ) -> None:
+        super().__init__()
+        self.model_name = client.model_name
+        self.client = client.client
+
+    def encode(
+        self,
+        texts: list[str],
+        show_progress_bar: bool = False,
+        batch_size: int = 32,
+    ) -> np.ndarray:
+        if batch_size < 1:
+            raise ValueError("batch_size must be at least 1.")
+        if not texts:
+            return np.empty((0, 0), dtype=np.float32)
+
+        embeddings: list[list[float]] = []
+        batch_starts = range(0, len(texts), batch_size)
+        for start in tqdm(batch_starts, disable=not show_progress_bar):
+            batch = texts[start:start + batch_size]
+            response = self.client.embeddings.create(
+                model=self.model_name,
+                input=batch,
+            )
+            response_data = list(response.data)
+            if all(hasattr(item, "index") for item in response_data):
+                response_data.sort(key=lambda item: item.index)
+            embeddings.extend(item.embedding for item in response_data)
+
+        return np.asarray(embeddings, dtype=np.float32)

gap_finder/io_helper.py

@@ -0,0 +1,71 @@
+"""I/O helpers for loading and saving all pipeline artifacts."""
+import logging
+from datetime import datetime, timezone
+from pathlib import Path
+from uuid import uuid4
+
+import pandas as pd
+
+from gap_finder.singleton import constants
+
+logger = logging.getLogger(__name__)
+
+
+def _save_dataframe(df: pd.DataFrame, path: Path):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    df.to_json(path, orient='records', force_ascii=False, indent=2)
+    logger.info(f"Saved {len(df)} entries to {path}")
+
+
+def _save_text(text: str, path: Path):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(text, encoding='utf-8')
+    logger.info(f"Saved report to {path}")
+
+
+def bertopic_model_path(base_path: str, id: str) -> Path:
+    return Path(base_path) / Path(id) / constants.DataFileNames.BERTOPIC_MODEL
+
+
+class FileHandler:
+    """Save the artifacts produced by one GapFinder run."""
+
+    def __init__(self, base_path: str | Path):
+        timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ")
+        self.run_path = Path(base_path) / f"{timestamp}_{uuid4().hex[:8]}"
+        self.run_path.mkdir(parents=True, exist_ok=False)
+        logger.info(f"Saving intermediate results to {self.run_path}")
+
+    def _path(self, filename: str) -> Path:
+        return self.run_path / filename
+
+    def save_questions(self, questions: pd.DataFrame):
+        _save_dataframe(
+            questions, self._path(constants.DataFileNames.QUESTIONS))
+
+    def save_questions_with_topics(self, questions: pd.DataFrame):
+        _save_dataframe(
+            questions,
+            self._path(constants.DataFileNames.QUESTIONS_WITH_TOPICS),
+        )
+
+    def save_questions_rankings(self, questions: pd.DataFrame):
+        _save_dataframe(
+            questions, self._path(constants.DataFileNames.QUESTIONS_RANKINGS))
+
+    def save_topic_rankings(self, topics: pd.DataFrame):
+        _save_dataframe(
+            topics, self._path(constants.DataFileNames.TOPICS_RANKINGS_JSON))
+
+    def save_topic_explanations(self, topics: pd.DataFrame):
+        _save_dataframe(
+            topics, self._path(constants.DataFileNames.TOPIC_EXPLANATIONS_JSON))
+
+    def save_topic_relevance_scores(self, topics: pd.DataFrame):
+        _save_dataframe(
+            topics,
+            self._path(constants.DataFileNames.TOPIC_RELEVANCE_SCORES_JSON),
+        )
+
+    def save_report(self, report: str):
+        _save_text(report, self._path(constants.DataFileNames.REPORT_MD))

gap_finder/language_model/__init__.py

@@ -0,0 +1,39 @@
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+from gap_finder._auth import MissingAPIKeyError
+from gap_finder.language_model.client import ClientError, OpenAIClient
+from gap_finder.language_model.language_model import LanguageModel
+
+if TYPE_CHECKING:
+    from gap_finder.language_model.factory import create_language_model
+    from gap_finder.language_model.local_language_model import LocalLanguageModel
+    from gap_finder.language_model.openai_language_model import OpenAILanguageModel
+
+__all__ = [
+    "ClientError",
+    "LanguageModel",
+    "LocalLanguageModel",
+    "MissingAPIKeyError",
+    "OpenAIClient",
+    "OpenAILanguageModel",
+    "create_language_model",
+]
+
+
+def __getattr__(name: str):
+    if name == "create_language_model":
+        from gap_finder.language_model.factory import create_language_model
+
+        return create_language_model
+    if name == "LocalLanguageModel":
+        from gap_finder.language_model.local_language_model import LocalLanguageModel
+
+        return LocalLanguageModel
+    if name == "OpenAILanguageModel":
+        from gap_finder.language_model.openai_language_model import OpenAILanguageModel
+
+        return OpenAILanguageModel
+    if name == "factory":
+        return import_module("gap_finder.language_model.factory")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

gap_finder/language_model/client.py

@@ -0,0 +1,43 @@
+from openai import OpenAI
+
+from gap_finder._auth import resolve_api_key
+
+
+class ClientError(Exception):
+    """Error raised when the language-model client returns an invalid result."""
+
+    def __init__(self, *args, **context):
+        super().__init__(*args)
+        self.context = context
+
+
+class OpenAIClient:
+    """Configured OpenAI-compatible chat-completions client."""
+
+    def __init__(
+        self,
+        model_name: str,
+        base_url: str | None = None,
+        api_key: str | None = None,
+    ) -> None:
+        self.model_name = model_name
+        self.client = OpenAI(
+            base_url=base_url,
+            api_key=resolve_api_key(api_key),
+        )
+
+    def invoke(self, prompt: str, system_prompt: str | None = None) -> str | None:
+        """Send a chat request and return its response text."""
+        try:
+            messages = []
+            if system_prompt is not None:
+                messages.append({"role": "system", "content": system_prompt})
+            messages.append({"role": "user", "content": prompt})
+
+            response = self.client.chat.completions.create(
+                model=self.model_name,
+                messages=messages,
+            )
+            return response.choices[0].message.content
+        except Exception as error:
+            raise ClientError(*error.args) from error

gap_finder/language_model/factory.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from gap_finder.language_model import client as language_model_client
+from gap_finder.language_model.language_model import LanguageModel
+from gap_finder.model_types import LanguageModelInput
+
+
+def create_language_model(
+    language_model: LanguageModelInput,
+    base_url: str | None = None,
+    api_key: str | None = None,
+) -> LanguageModel:
+    """Create a language-model wrapper from a model name or client."""
+    if isinstance(language_model, LanguageModel):
+        return language_model
+
+    if isinstance(language_model, str):
+        from gap_finder.language_model.openai_language_model import OpenAILanguageModel
+
+        client = language_model_client.OpenAIClient(
+            model_name=language_model,
+            base_url=base_url,
+            api_key=api_key,
+        )
+        return OpenAILanguageModel(client=client)
+    if isinstance(language_model, language_model_client.OpenAIClient):
+        from gap_finder.language_model.openai_language_model import OpenAILanguageModel
+
+        return OpenAILanguageModel(client=language_model)
+
+    from transformers.pipelines.text_generation import TextGenerationPipeline
+
+    if isinstance(language_model, TextGenerationPipeline):
+        from gap_finder.language_model.local_language_model import LocalLanguageModel
+
+        return LocalLanguageModel(client=language_model)
+    raise TypeError(
+        "language_model must be a LanguageModel, model name, OpenAIClient, or "
+        f"TextGenerationPipeline; got {type(language_model).__name__}."
+    )

gap_finder/language_model/language_model.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+
+
+class LanguageModel(ABC):
+    """Abstract interface for prompt-based language models."""
+
+    @abstractmethod
+    def invoke(self, prompt: str, system_prompt: str | None = None) -> str:
+        """Send a prompt and an optional system prompt to the model."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def invoke_batch(
+        self,
+        prompts: list[str],
+        system_prompt: str | None = None,
+    ) -> list[str]:
+        """Send multiple prompts and an optional system prompt to the model."""
+        raise NotImplementedError
+
+    def activate(self) -> None:
+        """Make internally owned resources ready for use."""
+        return
+
+    def deactivate(self) -> None:
+        """Optionally offload internally owned resources between steps."""
+        return
+
+    def cleanup(self) -> None:
+        """Release internally owned resources after the pipeline finishes."""
+        return

gap_finder/language_model/local_language_model.py

@@ -0,0 +1,69 @@
+from typing import Any
+
+from transformers.pipelines.text_generation import TextGenerationPipeline
+
+from gap_finder.language_model.language_model import LanguageModel
+
+
+class LocalLanguageModel(LanguageModel):
+    """Language model backed by a Transformers text-generation pipeline."""
+
+    def __init__(
+        self,
+        client: TextGenerationPipeline,
+    ) -> None:
+        if getattr(client, "task", None) != "text-generation":
+            raise ValueError(
+                "language_model pipeline must have task 'text-generation'."
+            )
+        self.client = client
+
+    def invoke(self, prompt: str, system_prompt: str | None = None) -> str:
+        text = self._format_prompt(prompt, system_prompt)
+        return self._extract_text(self._generate(text))
+
+    def invoke_batch(
+        self,
+        prompts: list[str],
+        system_prompt: str | None = None,
+    ) -> list[str]:
+        if not prompts:
+            return []
+        texts = [self._format_prompt(prompt, system_prompt) for prompt in prompts]
+        results = self._generate(texts)
+        return [self._extract_text(result) for result in results]
+
+    def _format_prompt(self, prompt: str, system_prompt: str | None = None) -> str:
+        messages = []
+        if system_prompt is not None:
+            messages.append({"role": "system", "content": system_prompt})
+        messages.append({"role": "user", "content": prompt})
+        formatted_prompt = self.client.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        if not isinstance(formatted_prompt, str):
+            raise TypeError(
+                "Tokenizer chat template must return text when tokenize=False."
+            )
+        return formatted_prompt
+
+    def _generate(self, text: str | list[str]) -> Any:
+        return self.client(
+            text,
+            do_sample=False,
+            return_full_text=False,
+        )
+
+    @staticmethod
+    def _extract_text(result: Any) -> str:
+        while isinstance(result, list):
+            if not result:
+                raise ValueError("Language model returned no generations.")
+            result = result[0]
+        if not isinstance(result, dict) or not isinstance(
+            result.get("generated_text"), str
+        ):
+            raise ValueError("Language model returned an invalid generation.")
+        return result["generated_text"]

gap_finder/language_model/openai_language_model.py

@@ -0,0 +1,45 @@
+from gap_finder.language_model.client import ClientError, OpenAIClient
+from gap_finder.language_model.language_model import LanguageModel
+
+
+class OpenAILanguageModel(LanguageModel):
+    """Language model backed by a configured OpenAI-compatible client."""
+
+    def __init__(self, client: OpenAIClient) -> None:
+        self.client = client
+
+    def invoke(self, prompt: str, system_prompt: str | None = None) -> str:
+        response = self.client.invoke(prompt=prompt, system_prompt=system_prompt)
+        return self._require_response(response, prompt, system_prompt)
+
+    def invoke_batch(
+        self,
+        prompts: list[str],
+        system_prompt: str | None = None,
+    ) -> list[str]:
+        return [
+            self._require_response(
+                self.client.invoke(prompt=prompt, system_prompt=system_prompt),
+                prompt,
+                system_prompt,
+                index=index,
+            )
+            for index, prompt in enumerate(prompts)
+        ]
+
+    def _require_response(
+        self,
+        response: str | None,
+        prompt: str,
+        system_prompt: str | None,
+        **context,
+    ) -> str:
+        if response is None:
+            raise ClientError(
+                "OpenAIClient returned None response.",
+                prompt=prompt,
+                system_prompt=system_prompt,
+                client=self.client,
+                **context,
+            )
+        return response