<commit_before>from bayesianpy.network import Builder as builder
import pandas as pd
import bayesianpy.network
from bayesianpy.jni import *
import bayesianpy.distributed as dk
import logging
from typing import List, Tuple

class Template:
    def __init__(self, discrete=pd.DataFrame(), continuous=pd.DataFrame(), label:str=None):
        self._discrete = discrete
        self._continuous = continuous
        self._label = label

    def create(self, network_factory: bayesianpy.network.NetworkFactory):
        return network_factory.create()

    def get_label(self):
        return self._label


class Tpl(Template):
    def __init__(self):
        super().__init__()

    def create(self, network_factory: bayesianpy.network.NetworkFactory):
        super().create(network_factory)


class NaiveBayes(Template):
    def __init__(self, parent_node: str, logger, discrete=pd.DataFrame(), continuous=pd.DataFrame(), discrete_states={}):
        super().__init__(discrete=discrete, continuous=continuous)
        self._discrete_states = discrete_states
        self._logger = logger
        self._parent_node = parent_node

    def create(self, network_factory):
        network = network_factory.create()

        if not dk.empty(self._continuous):
            for c_name in self._continuous.columns:
                c = builder.create_continuous_variable(network, c_name)

        if dk.empty(self._discrete):
            for d_name in self._discrete.columns:
                if d_name in self._discrete_states:
                    states = self._discrete_states[d_name]
                else:
                    states = dk.compute(self._discrete[d_name].dropna().unique()).tolist()

                try:
                    c = builder.create_discrete_variable(network, self._discrete, d_name, states)
                except BaseException as e:
                    self._logger.warn(e)

        parent_node = builder.try_get_node(network, self._parent_node)
        if parent_node is None:
            raise ValueError("Parent node: {} not recognised".format(self._parent_node))

        for node in network.getNodes():
            if node == parent_node:
                continue
            builder.create_link(network, parent_node, node)

        return network

class MixtureNaiveBayes(Template):

    def __init__(self, logger, discrete=pd.DataFrame(), continuous=pd.DataFrame(), latent_states=10, discrete_states={}, latent_variable_name='Cluster'):
        super().__init__(discrete=discrete, continuous=continuous)
        self._latent_states = latent_states
        self._discrete_states = discrete_states
        self._logger = logger
        self._latent_variable_name = latent_variable_name

    def create(self, network_factory):
        network = network_factory.create()
        cluster = builder.try_get_node(network, "Cluster")
        if cluster is None:
            cluster = builder.create_cluster_variable(network, self._latent_states, variable_name=self._latent_variable_name)

        if not dk.empty(self._continuous):
            for c_name in self._continuous.columns:
                c = builder.create_continuous_variable(network, c_name)
                try:
                    builder.create_link(network, cluster, c)
                except ValueError as e:
                    self._logger.warn(e)

        if not dk.empty(self._discrete):
            for d_name in self._discrete.columns:
                if d_name in self._discrete_states:
                    states = self._discrete_states[str(d_name)]
                else:
                    states = dk.compute(self._discrete[str(d_name)].dropna().unique()).tolist()

                try:
                    c = builder.create_discrete_variable(network, self._discrete, str(d_name), states)
                <commit_msg>Slightly better logging<commit_after>from bayesianpy.network import Builder as builder
import pandas as pd
import bayesianpy.network
from bayesianpy.jni import *
import bayesianpy.distributed as dk
import logging
from typing import List, Tuple

class Template:
    def __init__(self, discrete=pd.DataFrame(), continuous=pd.DataFrame(), label:str=None):
        self._discrete = discrete
        self._continuous = continuous
        self._label = label

    def create(self, network_factory: bayesianpy.network.NetworkFactory):
        return network_factory.create()

    def get_label(self):
        return self._label


class Tpl(Template):
    def __init__(self):
        super().__init__()

    def create(self, network_factory: bayesianpy.network.NetworkFactory):
        super().create(network_factory)


class NaiveBayes(Template):
    def __init__(self, parent_node: str, logger, discrete=pd.DataFrame(), continuous=pd.DataFrame(), discrete_states={}):
        super().__init__(discrete=discrete, continuous=continuous)
        self._discrete_states = discrete_states
        self._logger = logger
        self._parent_node = parent_node

    def create(self, network_factory):
        network = network_factory.create()

        if not dk.empty(self._continuous):
            for c_name in self._continuous.columns:
                c = builder.create_continuous_variable(network, c_name)

        if dk.empty(self._discrete):
            for d_name in self._discrete.columns:
                if d_name in self._discrete_states:
                    states = self._discrete_states[d_name]
                else:
                    states = dk.compute(self._discrete[d_name].dropna().unique()).tolist()

                try:
                    c = builder.create_discrete_variable(network, self._discrete, d_name, states)
                except BaseException as e:
                    self._logger.warn(e)

        parent_node = builder.try_get_node(network, self._parent_node)
        if parent_node is None:
            raise ValueError("Parent node: {} not recognised".format(self._parent_node))

        for node in network.getNodes():
            if node == parent_node:
                continue
            builder.create_link(network, parent_node, node)

        return network

class MixtureNaiveBayes(Template):

    def __init__(self, logger, discrete=pd.DataFrame(), continuous=pd.DataFrame(), latent_states=10, discrete_states={}, latent_variable_name='Cluster'):
        super().__init__(discrete=discrete, continuous=continuous)
        self._latent_states = latent_states
        self._discrete_states = discrete_states
        self._logger = logger
        self._latent_variable_name = latent_variable_name

    def create(self, network_factory):
        network = network_factory.create()
        cluster = builder.try_get_node(network, "Cluster")
        if cluster is None:
            cluster = builder.create_cluster_variable(network, self._latent_states, variable_name=self._latent_variable_name)

        if not dk.empty(self._continuous):
            for c_name in self._continuous.columns:
                self._logger.info("Pre-processing {} column".format(c_name))
                c = builder.create_continuous_variable(network, c_name)
                try:
                    builder.create_link(network, cluster, c)
                except ValueError as e:
                    self._logger.warn(e)

        if not dk.empty(self._discrete):
            for d_name in self._discrete.columns:
                if d_name in self._discrete_states:
                    states = self._discrete_states[str(d_name)]
                else:
                    states = dk.compute(self._discrete[str(d_name)].dropna().unique()).tolist()

                try:
                    c = builder.create_discrete_variable(network, self._discrete, str(d_name), states)
                
{'fraction_non_alphanumeric': 0.08967215239082545, 'fraction_numerical': 0.0005183361409874304, 'mean_word_length': 2.4859981933152664, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '773449', 'n_tokens_mistral': 2139, 'n_tokens_neox': 2029, 'n_words': 444}