A command-line tool that detects, extracts log templates, and structuring logs to in-process database, leveraging template patterns generated by GPT-4.
from __future__ import annotations

import glob
import logging
import re
import sys
from typing import Generator

import duckdb
from pypika import Column, Query, Schema, Table

from .util import Log, Memory


class Sink:
    def __init__(self, path: str, max_size: int = 512):
        self.path = path
        self.max_size = max_size
        self.buffer = []

    def __iter__(self):
        files = glob.glob(self.path)
        logging.info(f"detect log files: {files}")

        def f():
            for file in files:
                with open(file, "r") as f:
                    for line in f:
                        yield line[:-1][: self.max_size]

        def b():
            while len(self.buffer) != 0:
                line = self.buffer.pop()
                yield line

        file = f()

        while True:
            buf = b()
            yield from buf
            yield from file
            yield None
            if len(self.buffer) == 0:
                break

    def send(self, lines):
        self.buffer += lines


Memory().serialize(re.Pattern, lambda p: p.pattern)


class Match:
    def __init__(self, memory: Memory, sink: Sink):
        self.sink = sink
        self.patterns = memory.load("patterns", [])
        for offset in range(0, len(self.patterns)):
            self.patterns[offset] = re.compile(self.patterns[offset])

    def __iter__(self):
        for line in self.sink:
            if line is None:
                yield None
                continue

            for regex in self.patterns:
                match = regex.match(line)
                if match is None:
                    continue
                yield Log(regex.pattern, line, match.groups())
                break
            else:
                yield line

    def send(self, pattern: re.Pattern):
        self.patterns.append(pattern)


def collect(
    logs: Generator[Log, None, None],
    max_lines: int,
    db_file: str = ":default:",
) -> duckdb.DuckDBPyConnection:
    if max_lines <= 0:
        max_lines = sys.maxsize

    db = duckdb.connect(db_file)

    for line, _ in zip(logs, range(max_lines)):
        schema = Schema("information_schema")
        table = db.execute(
            Query.from_(schema.tables)
            .select("*")
            .where(schema.tables.table_name == line.pattern)
            .get_sql()
        ).fetchone()
        if table is None:
            columns = [Column(f"c{id}", "string") for id in range(0, len(line.groups))]
            db.sql(
                Query.create_table(line.pattern.replace('"', '""'))
                .columns(*columns)
                .get_sql()
            )

        table = Table(line.pattern.replace('"', '""'))
        db.sql(Query.into(table).insert(*line.groups).get_sql())

    return db

bakalog/__init__.py
ethe/bakalog
import logging
import os
from glob import glob

import click
import openai
from IPython import embed
from rich.logging import RichHandler

from bakalog import Match, Sink, collect
from bakalog.cluster import Cluster
from bakalog.extract import extract
from bakalog.util import Memory, parse_size


@click.group()
def main():
    pass


@main.command(
    help="bakalog cache all extracted patterns to each files as default, clean the cache as needed."
)
def clean():
    files = glob(f"{Memory.PATH}/*")
    for f in files:
        os.remove(f)


@main.command()
@click.argument("file")
@click.option(
    "--gpt-base",
    default=openai.api_base,
    help="OpenAI API base.",
    show_default=True,
)
@click.option(
    "--max-lines",
    default=4096,
    help="Max log lines would be parsed, set 0 to parese all logs.",
    type=int,
    show_default=True,
)
@click.option(
    "--buf-size",
    default="2MB",
    help="Number of logs to cluster detection.",
    show_default=True,
)
@click.option(
    "--max-len",
    default=512,
    help="Max length of each log, rest of log would be dropped.",
    type=int,
    show_default=True,
)
@click.option(
    "--threshold",
    default=0.85,
    help="Threshold of logs clustering.",
    type=float,
    show_default=True,
)
def run(file, gpt_base, max_lines, buf_size, max_len, threshold):
    if "OPENAI_API_KEY" not in os.environ:
        logging.error(
            "the tool relies on GPT4, please set env: `OPENAI_API_KEY` as OpenAI API key."
        )
        return

    logging.basicConfig(
        level="INFO",
        format="%(message)s",
        datefmt="[%X]",
        handlers=[RichHandler()],
    )
    buf_size = parse_size(buf_size)

    with Memory().current(file):
        f = Sink(file, max_size=max_len)
        m = Match(Memory(), f)
        c = Cluster(f, m, buf_size=buf_size, threshold=threshold)
        e = extract(
            c,
            m,
            api_base=gpt_base,
            model="gpt-4",
            temperature=0,
        )
        result = collect(e, max_lines)

        embed(header="use variable `result` to get the result")


if __name__ == "__main__":
    main()

bakalog/__main__.py
ethe/bakalog
from __future__ import annotations

import os
import contextlib
import logging
from typing import Generator, List

import numpy
import torch

from . import Log, Match, Sink


@contextlib.contextmanager
def pool(model):
    pool = model.start_multi_process_pool()
    try:
        yield pool
    finally:
        model.stop_multi_process_pool(pool)


class Community:
    texts: List[str]
    embeddings: List[torch.Tensor]

    def __init__(self):
        self.texts = []
        self.embeddings = []

    def append(self, text: str, embedding: torch.Tensor):
        self.texts.append(text)
        self.embeddings.append(embedding)

    @classmethod
    def from_list(cls, texts: List[str], embeddings: List[torch.Tensor]):
        self = cls()
        self.texts = texts
        self.embeddings = embeddings
        return self

    def __iter__(self):
        for i in zip(self.texts, self.embeddings):
            yield i


def sample(community: torch.Tensor) -> List[int]:
    from sentence_transformers.util import cos_sim

    sims = cos_sim(community, community)

    min = (0, 1)
    for offset, sim in enumerate(sims):
        mean = sim.mean()
        if mean < min[1]:
            min = (offset, mean)

    sort = sorted(
        [(offset, sim) for offset, sim in enumerate(sims[min[0]].numpy())],
        key=lambda i: i[1],
    )

    samples = [community[min[0]], community[sort[0][0]]]

    sims = cos_sim(community, torch.from_numpy(numpy.array(samples)))

    min1 = (0, 1)
    for offset, sim in enumerate(sims):
        mean = sim.mean()
        if mean < min1[1]:
            min1 = (offset, mean)

    return [min[0], sort[0][0], min1[0]]


_path = os.path.dirname(__file__)


class Cluster:
    def __init__(
        self,
        sink: Sink,
        match: Match,
        model=os.path.join(_path, "../all-MiniLM-L6-v2"),
        buf_size=8 * 1024 * 1024,
        threshold=0.7,
        min_community_size=3,
    ):
        from sentence_transformers import SentenceTransformer as Embedder

        self.buf_size = buf_size
        self.threshold = threshold
        self.min_community_size = min_community_size
        self.sink = sink
        self.match = match
        self.model = Embedder(model)
        self.buffer = []
        self.size = 0

    def __iter__(self):
        import torch
        from .util import community_detection

        with pool(self.model) as p:
            for line in self.match:
                if isinstance(line, Log):
                    yield line
                    continue

                if line is not None:
                    self.buffer.append(line)
                    self.size += len(line)

                if self.size < self.buf_size and line is not None:
                    continue

                if self.size == 0:
                    continue

                logging.info(
                    f"embedding {format(self.size / 1024, '.2f')}KB / {len(self.buffer)} logs, it might take a while."
                )
                embeddings = torch.from_numpy(
                    self.model.encode_multi_process(self.buffer, p)
                )

                if len(embeddings) < 3:
                    continue

                logging.info("analyze log communities, it might take a while.")
                clusters = community_detection(
                    embeddings,
                    min_community_size=self.min_community_size,
                    threshold=self.threshold,
                )
                logging.info(f"get {len(clusters)} log communities.")

                if len(clusters) > 0:
                    yield from self._sample(clusters, embeddings)
                    self._recycle()
                else:
                    logging.warning(
                        f"no cluster is detected, maybe you should decrease the threshold."
                    )

    def _sample(self, clusters, embeddings):
        for cluster, _ in zip(clusters, range(0, 3)):
            vecs = [embeddings[i] for i in cluster]
            ids = sample(torch.from_numpy(numpy.array(vecs)))
            samples = Community.from_list(
                [self.buffer[cluster[id]] for id in ids],
                [embeddings[cluster[id]] for id in ids],
            )
            # sample_clusters.append(samples)

            logging.info(f"yield samples {samples.texts}")
            yield samples.texts

    def _recycle(self):
        send, self.buffer = self.buffer, []
        self.size = 0
        self.sink.send(send)

bakalog/cluster.py
ethe/bakalog
from __future__ import annotations

import json
import logging
import re
from typing import Generator

import openai

from . import Match
from .cluster import Cluster
from .util import Log


PROMPT = r"""
| Characters                  | Meaning                                                      |
| :-------------------------- | :----------------------------------------------------------- |
| `[xyz][a-c]`                | A character class. Matches any one of the enclosed characters. You can specify a range of characters by using a hyphen, but if the hyphen appears as the first or last character enclosed in the square brackets, it is taken as a literal hyphen to be included in the character class as a normal character.For example, `[abcd]` is the same as `[a-d]`. They match the "b" in "brisket", and the "a" or the "c" in "arch", but not the "-" (hyphen) in "non-profit".For example, `[abcd-]` and `[-abcd]` match the "b" in "brisket", the "a" or the "c" in "arch", and the "-" (hyphen) in "non-profit".For example, `[\w-]` is the same as `[A-Za-z0-9_-]`. They both match any of the characters in "no_reply@example-server.com" except for the "@" and the ".". |
| `[^xyz][^a-c]`              | A negated or complemented character class. That is, it matches anything that is not enclosed in the square brackets. You can specify a range of characters by using a hyphen, but if the hyphen appears as the first or last character enclosed in the square brackets, it is taken as a literal hyphen to be included in the character class as a normal character. For example, `[^abc]` is the same as `[^a-c]`. They initially match "o" in "bacon" and "h" in "chop".**Note:** The ^ character may also indicate the [beginning of input](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Guide/Regular_expressions/Assertions). |
| `.`                         | Has one of the following meanings:Matches any single character *except* line terminators: `\n`, `\r`, `\u2028` or `\u2029`. For example, `/.y/` matches "my" and "ay", but not "yes", in "yes make my day".Inside a character class, the dot loses its special meaning and matches a literal dot.Note that the `m` multiline flag doesn't change the dot behavior. So to match a pattern across multiple lines, the character class `[^]` can be used — it will match any character including newlines.The `s` "dotAll" flag allows the dot to also match line terminators. |
| `\d`                        | Matches any digit (Arabic numeral). Equivalent to `[0-9]`. For example, `/\d/` or `/[0-9]/` matches "2" in "B2 is the suite number". |
| `\D`                        | Matches any character that is not a digit (Arabic numeral). Equivalent to `[^0-9]`. For example, `/\D/` or `/[^0-9]/` matches "B" in "B2 is the suite number". |
| `\w`                        | Matches any alphanumeric character from the basic Latin alphabet, including the underscore. Equivalent to `[A-Za-z0-9_]`. For example, `/\w/` matches "a" in "apple", "5" in "$5.28", and "3" in "3D". |
| `\W`                        | Matches any character that is not a word character from the basic Latin alphabet. Equivalent to `[^A-Za-z0-9_]`. For example, `/\W/` or `/[^A-Za-z0-9_]/` matches "%" in "50%". |
| `\s`                        | Matches a single white space character, including space, tab, form feed, line feed, and other Unicode spaces. Equivalent to `[ \f\n\r\t\v\u00a0\u1680\u2000-\u200a\u2028\u2029\u202f\u205f\u3000\ufeff]`. For example, `/\s\w*/` matches " bar" in "foo bar". |
| `\S`                        | Matches a single character other than white space. Equivalent to `[^ \f\n\r\t\v\u00a0\u1680\u2000-\u200a\u2028\u2029\u202f\u205f\u3000\ufeff]`. For example, `/\S\w*/` matches "foo" in "foo bar". |
| `\t`                        | Matches a horizontal tab.                                    |
| `\r`                        | Matches a carriage return.                                   |
| `\n`                        | Matches a linefeed.                                          |
| `\v`                        | Matches a vertical tab.                                      |
| `\f`                        | Matches a form-feed.                                         |
| `[\b]`                      | Matches a backspace. If you're looking for the word-boundary character (`\b`), see [Boundaries](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Guide/Regular_expressions/Assertions). |
| `\0`                        | Matches a NUL character. Do not follow this with another digit. |
| `\c*X*`                     | Matches a control character using [caret notation](https://en.wikipedia.org/wiki/Caret_notation), where "X" is a letter from A–Z (corresponding to code points `U+0001`*–*`U+001F`). For example, `/\cM/` matches "\r" in "\r\n". |
| `\x*hh*`                    | Matches the character with the code `*hh*` (two hexadecimal digits). |
| `\u*hhhh*`                  | Matches a UTF-16 code-unit with the value `*hhhh*` (four hexadecimal digits). |
| `\u*{hhhh}* or *\u{hhhhh}*` | (Only when the `u` flag is set.) Matches the character with the Unicode value `U+*hhhh*` or `U+*hhhhh*` (hexadecimal digits). |
| `\`                         | Indicates that the following character should be treated specially, or "escaped". It behaves one of two ways.For characters that are usually treated literally, indicates that the next character is special and not to be interpreted literally. For example, `/b/` matches the character "b". By placing a backslash in front of "b", that is by using `/\b/`, the character becomes special to mean match a word boundary.For characters that are usually treated specially, indicates that the next character is not special and should be interpreted literally. For example, "*" is a special character that means 0 or more occurrences of the preceding character should be matched; for example, `/a*/` means match 0 or more "a"s. To match `*` literally, precede it with a backslash; for example, `/a\*/` matches "a*".Note that some characters like `:`, `-`, `@`, etc. neither have a special meaning when escaped nor when unescaped. Escape sequences like `\:`, `\-`, `\@` will be equivalent to their literal, unescaped character equivalents in regular expressions. However, in regular expressions with the [unicode flag](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Guide/Regular_expressions#advanced_searching_with_flags), these will cause an *invalid identity escape* error. This is done to ensure backward compatibility with existing code that uses new escape sequences like `\p` or `\k`.**Note:** To match this character literally, escape it with itself. In other words to search for `\` use `/\\/`. |
| `*x*|*y*`                   | **Disjunction:** Matches either "x" or "y". Each component, separated by a pipe (`|`), is called an *alternative*. For example, `/green|red/` matches "green" in "green apple" and "red" in "red apple".**Note:** A disjunction is another way to specify "a set of choices", but it's not a character class. Disjunctions are not atoms — you need to use a [group](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Guide/Regular_expressions/Groups_and_backreferences) to make it part of a bigger pattern. `[abc]` is functionally equivalent to `(?:a|b|c)`. |

| Characters | Meaning                                                      |
| :--------- | :----------------------------------------------------------- |
| `^`        | Matches the beginning of input. If the multiline flag is set to true, also matches immediately after a line break character. For example, `/^A/` does not match the "A" in "an A", but does match the first "A" in "An A".**Note:** This character has a different meaning when it appears at the start of a [character class](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Guide/Regular_expressions/Character_classes). |
| `$`        | Matches the end of input. If the multiline flag is set to true, also matches immediately before a line break character. For example, `/t$/` does not match the "t" in "eater", but does match it in "eat". |
| `\b`       | Matches a word boundary. This is the position where a word character is not followed or preceded by another word-character, such as between a letter and a space. Note that a matched word boundary is not included in the match. In other words, the length of a matched word boundary is zero.Examples:`/\bm/` matches the "m" in "moon".`/oo\b/` does not match the "oo" in "moon", because "oo" is followed by "n" which is a word character.`/oon\b/` matches the "oon" in "moon", because "oon" is the end of the string, thus not followed by a word character.`/\w\b\w/` will never match anything, because a word character can never be followed by both a non-word and a word character.To match a backspace character (`[\b]`), see [Character Classes](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Guide/Regular_expressions/Character_classes). |
| `\B`       | Matches a non-word boundary. This is a position where the previous and next character are of the same type: Either both must be words, or both must be non-words, for example between two letters or between two spaces. The beginning and end of a string are considered non-words. Same as the matched word boundary, the matched non-word boundary is also not included in the match. For example, `/\Bon/` matches "on" in "at noon", and `/ye\B/` matches "ye" in "possibly yesterday". |

| Characters | Meaning                                                      |
| :--------- | :----------------------------------------------------------- |
| `x(?=y)`   | **Lookahead assertion:** Matches "x" only if "x" is followed by "y". For example, /`Jack(?=Sprat)/` matches "Jack" only if it is followed by "Sprat". `/Jack(?=Sprat|Frost)/` matches "Jack" only if it is followed by "Sprat" or "Frost". However, neither "Sprat" nor "Frost" is part of the match results. |
| `x(?!y)`   | **Negative lookahead assertion:** Matches "x" only if "x" is not followed by "y". For example, `/\d+(?!\.)/` matches a number only if it is not followed by a decimal point. `/\d+(?!\.)/.exec('3.141')` matches "141" but not "3". |
| `(?<=y)x`  | **Lookbehind assertion:** Matches "x" only if "x" is preceded by "y". For example, `/(?<=Jack)Sprat/` matches "Sprat" only if it is preceded by "Jack". `/(?<=Jack|Tom)Sprat/` matches "Sprat" only if it is preceded by "Jack" or "Tom". However, neither "Jack" nor "Tom" is part of the match results. |
| `(?<!y)x`  | **Negative lookbehind assertion:** Matches "x" only if "x" is not preceded by "y". For example, `/(?<!-)\d+/` matches a number only if it is not preceded by a minus sign. `/(?<!-)\d+/.exec('3')` matches "3". `/(?<!-)\d+/.exec('-3')` match is not found because the number is preceded by the minus sign. |

| Characters   | Meaning                                                      |
| :----------- | :----------------------------------------------------------- |
| `(*x*)`      | **Capturing group:** Matches `*x*` and remembers the match. For example, `/(foo)/` matches and remembers "foo" in "foo bar".A regular expression may have multiple capturing groups. In results, matches to capturing groups typically in an array whose members are in the same order as the left parentheses in the capturing group. This is usually just the order of the capturing groups themselves. This becomes important when capturing groups are nested. Matches are accessed using the index of the result's elements (`[1], …, [n]`) or from the predefined `RegExp` object's properties (`$1, …, $9`).Capturing groups have a performance penalty. If you don't need the matched substring to be recalled, prefer non-capturing parentheses (see below).`String.prototype.match()` won't return groups if the `/.../g` flag is set. However, you can still use `String.prototype.matchAll()` to get all matches. |
| `(?<Name>x)` | **Named capturing group:** Matches "x" and stores it on the groups property of the returned matches under the name specified by `<Name>`. The angle brackets (`<` and `>`) are required for group name.For example, to extract the United States area code from a phone number, we could use `/\((?<area>\d\d\d)\)/`. The resulting number would appear under `matches.groups.area`. |
| `(?:*x*)`    | **Non-capturing group:** Matches "x" but does not remember the match. The matched substring cannot be recalled from the resulting array's elements (`[1], …, [n]`) or from the predefined `RegExp` object's properties (`$1, …, $9`). |
| `\*n*`       | Where "n" is a positive integer. A back reference to the last substring matching the n parenthetical in the regular expression (counting left parentheses). For example, `/apple(,)\sorange\1/` matches "apple, orange," in "apple, orange, cherry, peach". |
| \k<Name>     | A back reference to the last substring matching the **Named capture group** specified by `<Name>`.For example, `/(?<title>\w+), yes \k<title>/` matches "Sir, yes Sir" in "Do you copy? Sir, yes Sir!".**Note:** `\k` is used literally here to indicate the beginning of a back reference to a Named capture group. |

| Characters                                               | Meaning                                                      |
| :------------------------------------------------------- | :----------------------------------------------------------- |
| `*x**`                                                   | Matches the preceding item "x" 0 or more times. For example, `/bo*/` matches "boooo" in "A ghost booooed" and "b" in "A bird warbled", but nothing in "A goat grunted". |
| `*x*+`                                                   | Matches the preceding item "x" 1 or more times. Equivalent to `{1,}`. For example, `/a+/` matches the "a" in "candy" and all the "a"'s in "caaaaaaandy". |
| `*x*?`                                                   | Matches the preceding item "x" 0 or 1 times. For example, `/e?le?/` matches the "el" in "angel" and the "le" in "angle."If used immediately after any of the quantifiers `*`, `+`, `?`, or `{}`, makes the quantifier non-greedy (matching the minimum number of times), as opposed to the default, which is greedy (matching the maximum number of times). |
| `*x*{*n*}`                                               | Where "n" is a positive integer, matches exactly "n" occurrences of the preceding item "x". For example, `/a{2}/` doesn't match the "a" in "candy", but it matches all of the "a"'s in "caandy", and the first two "a"'s in "caaandy". |
| `*x*{*n*,}`                                              | Where "n" is a positive integer, matches at least "n" occurrences of the preceding item "x". For example, `/a{2,}/` doesn't match the "a" in "candy", but matches all of the a's in "caandy" and in "caaaaaaandy". |
| `*x*{*n*,*m*}`                                           | Where "n" is 0 or a positive integer, "m" is a positive integer, and `*m* > *n*`, matches at least "n" and at most "m" occurrences of the preceding item "x". For example, `/a{1,3}/` matches nothing in "cndy", the "a" in "candy", the two "a"'s in "caandy", and the first three "a"'s in "caaaaaaandy". Notice that when matching "caaaaaaandy", the match is "aaa", even though the original string had more "a"s in it. |
| `*x**?` `*x*+?` `*x*??` `*x*{n}?` `*x*{n,}?` `*x*{n,m}?` | By default quantifiers like `*` and `+` are "greedy", meaning that they try to match as much of the string as possible. The `?` character after the quantifier makes the quantifier "non-greedy": meaning that it will stop as soon as it finds a match. For example, given a string like "some <foo> <bar> new </bar> </foo> thing":`/<.*>/` will match "<foo> <bar> new </bar> </foo>"`/<.*?>/` will match "<foo>" |
"""


def extract(
    cluster: Cluster,
    match: Match,
    api_base: str = openai.api_base,
    model: str = "gpt-4",
    temperature: float = 0,
    max_tokens: int = 512,
) -> Generator[Log, None, None]:
    openai.api_base = api_base

    def compile(pattern: str) -> re.Pattern:
        return re.compile(pattern)

    kwargs = {
        "model": model,
        "functions": [
            {
                "name": "compile",
                "description": "Compile a regular expression pattern into a regular expression object.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "pattern": {
                            "type": "string",
                            "description": "regular expression",
                        },
                    },
                    "required": ["pattern"],
                },
            }
        ],
        "temperature": temperature,
        "max_tokens": max_tokens,
    }

    def get_messages(samples):
        samples = "\n".join(samples)

        return [
            {
                "role": "system",
                "content": "You are a senior regular expression developer."
                "Create a regular expression which could match, group and extract the log template of several logs below."
                "The exactly same part between logs must be the part of template."
                "Pattern should start with `^` and end with `$`",
            },
            {
                "role": "user",
                "content": samples,
            },
        ]

    for message in cluster:
        if isinstance(message, Log):
            yield message
            continue

        kwargs["messages"] = get_messages(message)

        logging.info(f"thinking about the template of log: {message}...")
        completion = openai.ChatCompletion.create(**kwargs)
        assert isinstance(completion, dict)
        logging.info(f"got {completion['choices'][0]['message']}.")

        try:
            pattern = compile(
                json.loads(
                    completion["choices"][0]["message"]["function_call"]["arguments"]
                )["pattern"]
            )
            logging.info(
                f"extractd regex: {pattern.pattern}",
            )
            match.send(pattern)
            continue
        except Exception as e:
            logging.warning(
                f"faild to compile regex: {completion}, error: {e}, guessing..."
            )

        try:
            pattern = re.compile(
                json.loads(completion["choices"][0]["message"]["content"])["pattern"]
            )
            logging.info(
                f"guessing success: {pattern.pattern}",
            )
            match.send(pattern)
        except Exception as e:
            logging.error(f"failed to compile regex: {completion}, error: {e}")

bakalog/extract.py
ethe/bakalog
import contextlib
import json
import os
import re
from dataclasses import dataclass
from typing import Any, Tuple


@dataclass
class Log:
    pattern: str
    content: str
    groups: Tuple[str, ...]


class SingletonMeta(type):
    _instances = {}

    def __call__(cls, *args, **kwargs):
        if cls not in cls._instances:
            instance = super().__call__(*args, **kwargs)
            cls._instances[cls] = instance
        return cls._instances[cls]


class Singleton(metaclass=SingletonMeta):
    pass


class Serializer(json.JSONEncoder, metaclass=SingletonMeta):
    def __init__(self, *args, **kwargs):
        self.serializers = {}
        super().__init__(*args, **kwargs)

    def default(self, o: Any) -> Any:
        if type(o) in self.serializers:
            return self.serializers[type(o)](o)
        return super().default(o)

    def register(self, type, encoder):
        self.serializers[type] = encoder


class Memory(Singleton):
    home = os.environ["HOME"]
    PATH = f"{home}/.bakalog"

    def __init__(
        self,
    ):
        if not os.path.exists(self.PATH):
            os.mkdir(self.PATH)

    @contextlib.contextmanager
    def current(self, file: str):
        with open(f"{self.PATH}/patterns", "a+") as f:
            f.seek(0)
            if f.read() == "":
                f.write(r"{}")
                f.seek(0)
            f.seek(0)
            all = json.loads(f.read())
            if file not in all:
                all[file] = {}
            self.memory = all[file]
            try:
                yield
            finally:
                f.seek(0)
                f.truncate()
                f.write(json.dumps(all, cls=Serializer))

    @classmethod
    def serialize(cls, type, serializer):
        Serializer().register(type, serializer)

    def load(self, field, init):
        if field not in self.memory:
            self.memory[field] = init
        return self.memory[field]


def community_detection(
    embeddings, threshold=0.75, min_community_size=3, batch_size=1024
):
    """
    Function for Fast Community Detection
    Finds in the embeddings all communities, i.e. embeddings that are close (closer than threshold).
    Returns only communities that are larger than min_community_size. The communities are returned
    in decreasing order. The first element in each list is the central point in the community.
    """

    import torch
    from sentence_transformers.util import cos_sim

    threshold = torch.tensor(threshold, device=embeddings.device)

    extracted_communities = []

    # Maximum size for community
    min_community_size = min(min_community_size, len(embeddings))
    sort_max_size = min(max(2 * min_community_size, 50), len(embeddings))

    for start_idx in range(0, len(embeddings), batch_size):
        # Compute cosine similarity scores
        cos_scores = cos_sim(embeddings[start_idx : start_idx + batch_size], embeddings)

        # Minimum size for a community
        top_k_values, _ = cos_scores.topk(k=min_community_size, largest=True)

        # Filter for rows >= min_threshold
        for i in range(len(top_k_values)):
            if top_k_values[i][-1] >= threshold:
                new_cluster = []

                # Only check top k most similar entries
                top_val_large, top_idx_large = cos_scores[i].topk(
                    k=sort_max_size, largest=True
                )

                # Check if we need to increase sort_max_size
                while top_val_large[-1] > threshold:
                    sort_max_size = min(2 * sort_max_size, len(embeddings))
                    top_val_large, top_idx_large = cos_scores[i].topk(
                        k=sort_max_size, largest=True
                    )
                    if min(2 * sort_max_size, len(embeddings)) <= sort_max_size:
                        break

                for idx, val in zip(top_idx_large.tolist(), top_val_large):
                    if val < threshold:
                        break

                    new_cluster.append(idx)

                extracted_communities.append(new_cluster)

        del cos_scores

    # Largest cluster first
    extracted_communities = sorted(
        extracted_communities, key=lambda x: len(x), reverse=True
    )

    # Step 2) Remove overlapping communities
    unique_communities = []
    extracted_ids = set()

    for _, community in enumerate(extracted_communities):
        community = sorted(community)
        non_overlapped_community = []
        for idx in community:
            if idx not in extracted_ids:
                non_overlapped_community.append(idx)

        if len(non_overlapped_community) >= min_community_size:
            unique_communities.append(non_overlapped_community)
            extracted_ids.update(non_overlapped_community)

    unique_communities = sorted(unique_communities, key=lambda x: len(x), reverse=True)

    return unique_communities


units = {"B": 1, "KB": 2**10, "MB": 2**20, "GB": 2**30, "TB": 2**40}


def parse_size(size):
    size = size.upper()
    if not re.match(r" ", size):
        size = re.sub(r"([KMGT]?B)", r" \1", size)
    number, unit = [string.strip() for string in size.split()]
    return int(float(number) * units[unit])

bakalog/util.py
ethe/bakalog