poemai-utils 0.0.0__tar.gz

poemai-utils-0.0.0/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2023 Markus Emmenegger, poemAI GmbH, Horw, Switzerland
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

poemai-utils-0.0.0/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.1
+Name: poemai-utils
+Version: 0.0.0
+Summary: Add a short description here!
+Home-page: https://github.com/pyscaffold/pyscaffold/
+Author: Markus Emmenegger
+Author-email: markus.emmenegger@poemai.ch
+License: MIT
+Project-URL: Documentation, https://pyscaffold.org/
+Platform: any
+Classifier: Development Status :: 4 - Beta
+Classifier: Programming Language :: Python
+Description-Content-Type: text/markdown; charset=UTF-8
+Provides-Extra: testing
+License-File: LICENSE.txt
+
+# poemai-utils
+
+This package is a collection of utilities for AI projects.

poemai-utils-0.0.0/README.md

@@ -0,0 +1,3 @@
+# poemai-utils
+
+This package is a collection of utilities for AI projects.

poemai-utils-0.0.0/pyproject.toml

@@ -0,0 +1,12 @@
+[build-system]
+# AVOID CHANGING REQUIRES: IT WILL BE UPDATED BY PYSCAFFOLD!
+requires = ["setuptools>=46.1.0", "setuptools_scm[toml]>=5"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools_scm]
+# For smarter version schemes and other configuration options,
+# check out https://github.com/pypa/setuptools_scm
+version_scheme = "no-guess-dev"
+[tool.black]
+line-length = 88
+target-version = ['py311']

poemai-utils-0.0.0/setup.cfg

@@ -0,0 +1,74 @@
+[metadata]
+name = poemai-utils
+description = Add a short description here!
+author = Markus Emmenegger
+author_email = markus.emmenegger@poemai.ch
+license = MIT
+license_files = LICENSE.txt
+long_description = file: README.md
+long_description_content_type = text/markdown; charset=UTF-8
+url = https://github.com/pyscaffold/pyscaffold/
+project_urls = 
+	Documentation = https://pyscaffold.org/
+platforms = any
+classifiers = 
+	Development Status :: 4 - Beta
+	Programming Language :: Python
+
+[options]
+zip_safe = False
+packages = find_namespace:
+include_package_data = True
+package_dir = 
+	=src
+install_requires = 
+	importlib-metadata; python_version<"3.8"
+	sqlitedict
+	httpx
+	numpy
+
+[options.packages.find]
+where = src
+exclude = 
+	tests
+
+[options.extras_require]
+testing = 
+	setuptools
+	pytest
+	pytest-cov
+	transformers
+
+[options.entry_points]
+
+[tool:pytest]
+norecursedirs = 
+	dist
+	build
+	.tox
+testpaths = tests
+
+[devpi:upload]
+no_vcs = 1
+formats = bdist_wheel
+
+[flake8]
+max_line_length = 88
+extend_ignore = E203, W503
+exclude = 
+	.tox
+	build
+	dist
+	.eggs
+	docs/conf.py
+
+[pyscaffold]
+version = 4.4.1
+package = poemai_utils
+
+[options.package_data]
+
+[egg_info]
+tag_build = 
+tag_date = 0
+

poemai-utils-0.0.0/setup.py

@@ -0,0 +1,21 @@
+"""
+    Setup file for poemai-utils.
+    Use setup.cfg to configure your project.
+
+    This file was generated with PyScaffold 4.4.1.
+    PyScaffold helps you to put up the scaffold of your new Python project.
+    Learn more under: https://pyscaffold.org/
+"""
+from setuptools import setup
+
+if __name__ == "__main__":
+    try:
+        setup(use_scm_version={"version_scheme": "no-guess-dev"})
+    except:  # noqa
+        print(
+            "\n\nAn error occurred while building the project, "
+            "please ensure you have the most updated version of setuptools, "
+            "setuptools_scm and wheel with:\n"
+            "   pip install -U setuptools setuptools_scm wheel\n\n"
+        )
+        raise

poemai-utils-0.0.0/src/poemai_utils/__init__.py

@@ -0,0 +1,16 @@
+import sys
+
+if sys.version_info[:2] >= (3, 8):
+    # TODO: Import directly (no need for conditional) when `python_requires = >= 3.8`
+    from importlib.metadata import PackageNotFoundError, version  # pragma: no cover
+else:
+    from importlib_metadata import PackageNotFoundError, version  # pragma: no cover
+
+try:
+    # Change here if project is renamed and does not equal the package name
+    dist_name = "poemai-utils"
+    __version__ = version(dist_name)
+except PackageNotFoundError:  # pragma: no cover
+    __version__ = "unknown"
+finally:
+    del version, PackageNotFoundError

poemai-utils-0.0.0/src/poemai_utils/basic_types_utils.py

@@ -0,0 +1,52 @@
+def linebreak(text, max_length=80):
+    """Add linebreaks to text to make it more readable"""
+
+    out_lines = []
+    for line in text.split("\n"):
+        if len(line) <= max_length:
+            out_lines.append(line)
+        else:
+            # split line into words
+            words = line.split(" ")
+            out_line = ""
+            for word in words:
+                if len(out_line) + len(word) > max_length:
+                    out_lines.append(out_line)
+                    out_line = ""
+                out_line += word + " "
+            out_lines.append(out_line)
+    return "\n".join(out_lines)
+
+
+def short_display(text, max_lenght=100):
+    if len(text) > max_lenght:
+        return text[:max_lenght] + "..."
+    return text
+
+
+def remove_none_from_dict(
+    dictionary, remove_empty_lists=False, remove_empty_dicts=False
+):
+    retval = {}
+
+    for key, value in list(dictionary.items()):
+        if value is None:
+            continue
+        elif isinstance(value, dict):
+            new_value = remove_none_from_dict(value)
+            if not new_value and remove_empty_dicts:
+                continue
+            else:
+                retval[key] = new_value
+        elif isinstance(value, list):
+            if len(value) == 0 and remove_empty_lists:
+                continue
+            else:
+                retval[key] = [
+                    remove_none_from_dict(item) if isinstance(item, dict) else item
+                    for item in value
+                ]
+        else:
+            retval[key] = value
+
+    return retval

poemai-utils-0.0.0/src/poemai_utils/configuration_api.py

@@ -0,0 +1,2 @@
+class ConfigurationAPI:
+    pass

poemai-utils-0.0.0/src/poemai_utils/embeddings/cross_similarity_analyzer.py

@@ -0,0 +1,121 @@
+import numpy as np
+from poemai_utils.embeddings.self_similarity_analyzer import SelfSimilarityAnalyzer
+
+
+class CrossSimilarityAnalyzer:
+    """A class to analyze the similarity between two lists of embeddings."""
+
+    def __init__(
+        self,
+        embedding_cache=None,
+        similarity_analyzer_1=None,
+        similarity_analyzer_2=None,
+        embedder=None,
+    ):
+        if similarity_analyzer_1 is None:
+            self.analyzer_1 = SelfSimilarityAnalyzer(
+                embedder,
+                embedding_cache=embedding_cache,
+            )
+        else:
+            self.analyzer_1 = similarity_analyzer_1
+        if similarity_analyzer_2 is None:
+            self.analyzer_2 = SelfSimilarityAnalyzer(
+                embedder, embedding_cache=embedding_cache
+            )
+        else:
+            self.analyzer_2 = similarity_analyzer_2
+
+        if (
+            self.analyzer_1.use_cosine_similarity
+            != self.analyzer_2.use_cosine_similarity
+        ):
+            raise ValueError("Both analyzers must have the same use_cosine_similarity")
+
+        self.use_cosine_similarity = self.analyzer_1.use_cosine_similarity
+
+    def add_text_for_1(self, text, metadata):
+        self.analyzer_1.add_text(text, metadata)
+
+    def add_embedding_for_1(self, text, embedding, metadata):
+        self.analyzer_1.add_embedding(text, embedding, metadata)
+
+    def add_text_for_2(self, text, metadata):
+        self.analyzer_2.add_text(text, metadata)
+
+    def add_embedding_for_2(self, text, embedding, metadata):
+        self.analyzer_2.add_embedding(text, embedding, metadata)
+
+    @staticmethod
+    def normalize_rows(x):
+        return x / np.linalg.norm(x, ord=2, axis=1, keepdims=True)
+
+    def cross_similarity_matrix(self):
+        m1 = self.analyzer_1.embedding_store.embedding_matrix
+        m2 = self.analyzer_2.embedding_store.embedding_matrix
+
+        if self.use_cosine_similarity:
+            m1 = self.normalize_rows(m1)
+            m2 = self.normalize_rows(m2)
+
+        return np.matmul(
+            m1,
+            m2.T,
+        )
+
+    def find_top_k_similar_pairs(self, k=5):
+        similarity_matrix = self.cross_similarity_matrix()
+
+        top_k_idxs = np.argsort(similarity_matrix, axis=None)[-k:]
+        rows, cols = np.unravel_index(top_k_idxs, similarity_matrix.shape)
+        values = similarity_matrix[rows, cols]
+
+        retval = []
+        for row, col, value in zip(rows, cols, values):
+            row_metadata = self.analyzer_1.metadata.get(row)
+            col_metadata = self.analyzer_2.metadata.get(col)
+
+            retval.append(
+                (
+                    row_metadata,
+                    col_metadata,
+                    value,
+                )
+            )
+
+        return retval
+
+    def map_1_to_2(self, k=1):
+        similarity_matrix = self.cross_similarity_matrix()
+        best = np.argsort(similarity_matrix, axis=1)[:, ::-1][:, :k]
+
+        retval = []
+        for i, row in enumerate(best):
+            row_metadata = self.analyzer_1.metadata.get(i)
+            col_metadata = [self.analyzer_2.metadata.get(b) for b in row]
+
+            retval.append(
+                (
+                    row_metadata,
+                    col_metadata,
+                )
+            )
+        return retval
+
+    def map_1_to_2_with_scores(self, k=1):
+        similarity_matrix = self.cross_similarity_matrix()
+        best = np.argsort(similarity_matrix, axis=1)[:, ::-1][:, :k]
+
+        retval = []
+        for i, row in enumerate(best):
+            row_metadata = self.analyzer_1.metadata.get(i)
+            col_metadata = [self.analyzer_2.metadata.get(b) for b in row]
+
+            col_result = list(zip(col_metadata, similarity_matrix[i, row]))
+            retval.append(
+                (
+                    row_metadata,
+                    col_result,
+                )
+            )
+        return retval

poemai-utils-0.0.0/src/poemai_utils/embeddings/embedder_base.py

@@ -0,0 +1,3 @@
+class EbedderBase:
+    def __init__(self, use_cosine_similarity=False):
+        self.use_cosine_similarity = use_cosine_similarity

poemai-utils-0.0.0/src/poemai_utils/embeddings/embedding_cache.py

@@ -0,0 +1,34 @@
+import logging
+from hashlib import sha256
+
+import numpy as np
+from sqlitedict import SqliteDict
+
+_logger = logging.getLogger(__name__)
+
+
+class EmbeddingCache:
+    def __init__(self, cache_file):
+        self.cache = SqliteDict(cache_file, autocommit=True)
+
+    def get(self, text, model_name):
+        text_hash = self.text_hash(text)
+        key = self.calc_key(model_name, text_hash)
+        short_key = (model_name, text_hash[:6])
+        if key in self.cache:
+            _logger.debug(f"Cache hit for {short_key}")
+            return np.array(self.cache[key])
+        return None
+
+    def put(self, text, embedding, model_name):
+        text_hash = self.text_hash(text)
+        key = self.calc_key(model_name, text_hash)
+        self.cache[key] = embedding.tolist()
+
+    def text_hash(self, text):
+        # calculate sha256
+        return sha256(text.encode("utf-8")).hexdigest()
+
+    @staticmethod
+    def calc_key(model_name, text_hash):
+        return model_name + "_" + text_hash

poemai-utils-0.0.0/src/poemai_utils/embeddings/embedding_store.py

@@ -0,0 +1,60 @@
+from hashlib import sha256
+
+import numpy as np
+
+
+class EmbeddingStore:
+    def __init__(self, embedding_cache=None, embedder=None):
+        self.embedder = embedder
+
+        self.texts = []
+        self.embedding_matrix = None
+        self.embedding_cache = embedding_cache
+
+    def _store_embedding(self, embedding):
+        if self.embedding_matrix is None:
+            self.embedding_matrix = np.expand_dims(embedding, 0)
+        else:
+            self.embedding_matrix = np.concatenate(
+                [self.embedding_matrix, [embedding]], 0
+            )
+
+    def add_text(self, text, is_query: bool = False):
+        self.texts.append(text)
+        self._store_embedding(self._create_embedding(text, is_query=is_query))
+        return len(self.texts) - 1
+
+    def add_embedding(self, text, embedding):
+        self.texts.append(text)
+        self._store_embedding(embedding)
+        return len(self.texts) - 1
+
+    def query_by_text(self, query, k=5):
+        if self.embedding_matrix is None:
+            return []
+        query_embedding = self._create_embedding(query)
+        scores = self.embedding_matrix.dot(query_embedding)
+        idxs = np.argsort(scores)[-k:][::-1]
+
+        return [(self.texts[i], scores[i], i) for i in idxs]
+
+    def query_by_embedding(self, query_embedding, k=5):
+        if self.embedding_matrix is None:
+            return []
+        scores = self.embedding_matrix.dot(query_embedding)
+        idxs = np.argsort(scores)[-k:][::-1]
+
+        return [(self.texts[i], scores[i], i) for i in idxs]
+
+    def _create_embedding(self, text, is_query):
+        if self.embedding_cache is not None:
+            embedding = self.embedding_cache.get(text, self.embedder.model_name)
+            if embedding is not None:
+                return embedding
+
+        embedding = self.embedder.calc_embedding(text, is_query=is_query)
+
+        if self.embedding_cache is not None:
+            self.embedding_cache.put(text, embedding, self.embedder.model_name)
+
+        return embedding

poemai-utils-0.0.0/src/poemai_utils/embeddings/openai_embedder.py

@@ -0,0 +1,16 @@
+import numpy as np
+import openai
+from poemai_utils.embeddings.embedder_base import EbedderBase
+
+
+class OpenAIEmbedder(EbedderBase):
+    def __init__(self, model_name="text-embedding-ada-002"):
+        super().__init__()
+
+        self.model_name = model_name
+
+    def calc_embedding(self, text, is_query: bool = False):
+        response = openai.Embedding.create(input=text, model=self.model_name)
+        embedding = response["data"][0]["embedding"]
+        embedding = np.array(embedding, dtype=np.float32)
+        return embedding

poemai-utils-0.0.0/src/poemai_utils/embeddings/self_similarity_analyzer.py

@@ -0,0 +1,70 @@
+import numpy as np
+from poemai_utils.embeddings.embedding_store import EmbeddingStore
+
+
+class SelfSimilarityAnalyzer:
+    def __init__(self, embedder, embedding_cache=None):
+        if embedder is None:
+            raise ValueError("embedder must be provided")
+        self.embedding_store = EmbeddingStore(embedding_cache, embedder=embedder)
+        self.use_cosine_similarity = embedder.use_cosine_similarity
+        self.metadata = {}
+
+    def add_text(self, text, metadata, is_query=False):
+        if text is None:
+            return
+        text_id = self.embedding_store.add_text(text, is_query=is_query)
+        self.metadata[text_id] = metadata
+
+    def add_embedding(self, text, embedding, metadata):
+        text_id = self.embedding_store.add_embedding(text, embedding)
+        self.metadata[text_id] = metadata
+
+    def query(self, query, k=5):
+        top_k = self.embedding_store.query_by_text(query, k)
+        return [
+            (text_id, score, self.metadata.get(text_id)) for _, score, text_id in top_k
+        ]
+
+    @staticmethod
+    def normalize_rows(x):
+        return x / np.linalg.norm(x, ord=2, axis=1, keepdims=True)
+
+    def similarity_matrix(self):
+        m1 = self.embedding_store.embedding_matrix
+        if self.use_cosine_similarity:
+            m1 = self.normalize_rows(m1)
+
+        return np.matmul(
+            m1,
+            m1.T,
+        )
+
+    def find_top_k_similar_pairs(self, k=5):
+        similarity_matrix = self.similarity_matrix()
+
+        similarity_matrix = np.tril(similarity_matrix)
+        np.fill_diagonal(similarity_matrix, 0)
+
+        # find the top k most similar sections
+        top_k_idxs = np.argsort(similarity_matrix, axis=None)[-k:]
+        rows, cols = np.unravel_index(top_k_idxs, similarity_matrix.shape)
+        values = similarity_matrix[rows, cols]
+
+        retval = []
+        for row, col, value in zip(rows, cols, values):
+            row_metadata = self.metadata.get(row)
+            col_metadata = self.metadata.get(col)
+
+            retval.append(
+                (
+                    row_metadata,
+                    col_metadata,
+                    value,
+                )
+            )
+
+        return retval
+
+    def metadata(self, text_id):
+        return self.metadata.get(text_id)

poemai-utils-0.0.0/src/poemai_utils/embeddings/sentence_transformer_embedder.py

@@ -0,0 +1,46 @@
+from enum import Enum
+
+from poemai_utils.embeddings.embedder_base import EbedderBase
+from poemai_utils.enum_utils import add_enum_attrs, add_enum_repr
+from sentence_transformers import SentenceTransformer
+
+
+class SentenceTransformerEmbeddingModel(Enum):
+    LABSE = "sentence-transformers/LaBSE"
+    DISTILUSE = "distiluse-base-multilingual-cased-v1"
+    BI_ELECTRA_GERMAN = "svalabs/bi-electra-ms-marco-german-uncased"
+    DISTILBERT = "msmarco-distilbert-base-tas-b"
+
+
+add_enum_attrs(
+    {
+        SentenceTransformerEmbeddingModel.LABSE: {
+            "use_cosine_similarity": False,
+        },
+        SentenceTransformerEmbeddingModel.DISTILUSE: {
+            "use_cosine_similarity": False,
+        },
+        SentenceTransformerEmbeddingModel.BI_ELECTRA_GERMAN: {
+            "use_cosine_similarity": True,
+        },
+        SentenceTransformerEmbeddingModel.DISTILBERT: {
+            "use_cosine_similarity": False,
+        },
+    }
+)
+add_enum_repr(SentenceTransformerEmbeddingModel)
+
+
+class SentenceTransformerEmbedder(EbedderBase):
+    # msmarco-distilbert-base-tas-b
+    # distiluse-base-multilingual-cased-v1
+    # sentence-transformers/LaBSE'
+    # svalabs/bi-electra-ms-marco-german-uncased : use_cosine_similarity=True
+    def __init__(self, model_id, use_cosine_similarity=False):
+        super().__init__(use_cosine_similarity=model_id.use_cosine_similarity)
+
+        self.model_name = model_id.value
+        self.model = SentenceTransformer(model_id.value)
+
+    def calc_embedding(self, text, is_query: bool = False):
+        return self.model.encode(text, show_progress_bar=False)

poemai-utils-0.0.0/src/poemai_utils/embeddings/sgpt_embedder.py

@@ -0,0 +1,85 @@
+import torch
+from poemai_utils.embeddings.embedder_base import EbedderBase
+from transformers import AutoModel, AutoTokenizer
+
+
+class SGPTEmbedder(EbedderBase):
+    def __init__(
+        self, model_name="Muennighoff/SGPT-125M-weightedmean-msmarco-specb-bitfit"
+    ):
+        super().__init__()
+
+        self.use_cosine_similarity = True
+        self.model_name = model_name
+        # Get our models - The package will take care of downloading the models automatically
+        # For best performance: Muennighoff/SGPT-5.8B-weightedmean-msmarco-specb-bitfit
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+
+        self.SPECB_QUE_BOS = self.tokenizer.encode("[", add_special_tokens=False)[0]
+        self.SPECB_QUE_EOS = self.tokenizer.encode("]", add_special_tokens=False)[0]
+
+        self.SPECB_DOC_BOS = self.tokenizer.encode("{", add_special_tokens=False)[0]
+        self.SPECB_DOC_EOS = self.tokenizer.encode("}", add_special_tokens=False)[0]
+
+        self.model = AutoModel.from_pretrained(self.model_name)
+        # Deactivate Dropout (There is no dropout in the above models so it makes no difference here but other SGPT models may have dropout)
+        self.model.eval()
+
+    def tokenize_with_specb(self, texts, is_query):
+        # Tokenize without padding
+        batch_tokens = self.tokenizer(texts, padding=False, truncation=True)
+        # Add special brackets & pay attention to them
+        for seq, att in zip(batch_tokens["input_ids"], batch_tokens["attention_mask"]):
+            if is_query:
+                seq.insert(0, self.SPECB_QUE_BOS)
+                seq.append(self.SPECB_QUE_EOS)
+            else:
+                seq.insert(0, self.SPECB_DOC_BOS)
+                seq.append(self.SPECB_DOC_EOS)
+            att.insert(0, 1)
+            att.append(1)
+        # Add padding
+        batch_tokens = self.tokenizer.pad(
+            batch_tokens, padding=True, return_tensors="pt"
+        )
+        return batch_tokens
+
+    def get_weightedmean_embedding(self, batch_tokens):
+        # Get the embeddings
+        with torch.no_grad():
+            # Get hidden state of shape [bs, seq_len, hid_dim]
+            last_hidden_state = self.model(
+                **batch_tokens, output_hidden_states=True, return_dict=True
+            ).last_hidden_state
+
+        # Get weights of shape [bs, seq_len, hid_dim]
+        weights = (
+            torch.arange(start=1, end=last_hidden_state.shape[1] + 1)
+            .unsqueeze(0)
+            .unsqueeze(-1)
+            .expand(last_hidden_state.size())
+            .float()
+            .to(last_hidden_state.device)
+        )
+
+        # Get attn mask of shape [bs, seq_len, hid_dim]
+        input_mask_expanded = (
+            batch_tokens["attention_mask"]
+            .unsqueeze(-1)
+            .expand(last_hidden_state.size())
+            .float()
+        )
+
+        # Perform weighted mean pooling across seq_len: bs, seq_len, hidden_dim -> bs, hidden_dim
+        sum_embeddings = torch.sum(
+            last_hidden_state * input_mask_expanded * weights, dim=1
+        )
+        sum_mask = torch.sum(input_mask_expanded * weights, dim=1)
+
+        embeddings = sum_embeddings / sum_mask
+
+        return embeddings
+
+    def calc_embedding(self, text, is_query: bool = False):
+        tokens = self.tokenize_with_specb([text], is_query=is_query)
+        return self.get_weightedmean_embedding(tokens)[0]