<commit_before>import torch
import torch.nn as nn
import torch.optim as optim
import time
import sys
import os
import uuid
from threading import Thread
from termcolor import colored
import math
import random
import tarfile
import tempfile
import shutil

from .distributions import Empirical
from . import state, util, __version__, TraceMode, InferenceEngine, Optimizer
from .nn import ObserveEmbedding, SampleEmbedding, Batch, InferenceNetwork
from .remote import ModelServer
from .analytics import save_report


class Model(nn.Module):
    def __init__(self, name='Unnamed pyprob model'):
        super().__init__()
        self.name = name
        self._inference_network = None
        self._trace_cache_path = None
        self._trace_cache = []
        self._trace_cache_discarded_file_names = []

    def forward(self):
        raise NotImplementedError()

    def _prior_trace_generator(self, trace_mode=TraceMode.DEFAULT, inference_network=None, metropolis_hastings_trace=None, *args, **kwargs):
        while True:
            if trace_mode == TraceMode.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK:
                self._inference_network.new_trace(util.pack_observes_to_variable(kwargs['observation']).unsqueeze(0))
            state.begin_trace(self.forward, trace_mode, inference_network, metropolis_hastings_trace)
            res = self.forward(*args, **kwargs)
            trace = state.end_trace(res)
            yield trace

    def _prior_sample_generator(self, *args, **kwargs):
        while True:
            yield self.forward(*args, **kwargs)

    def _prior_traces(self, num_traces=10, trace_mode=TraceMode.DEFAULT, inference_network=None, map_func=None, *args, **kwargs):
        generator = self._prior_trace_generator(trace_mode, inference_network, *args, **kwargs)
        ret = []
        time_start = time.time()
        if ((trace_mode != TraceMode.DEFAULT) and (util.verbosity > 1)) or (util.verbosity > 2):
            len_str_num_traces = len(str(num_traces))
            print('Time spent  | Time remain.| Progress             | {} | Traces/sec'.format('Trace'.ljust(len_str_num_traces * 2 + 1)))
            prev_duration = 0
        for i in range(num_traces):
            if ((trace_mode != TraceMode.DEFAULT) and (util.verbosity > 1)) or (util.verbosity > 2):
                duration = time.time() - time_start
                if (duration - prev_duration > util._print_refresh_rate) or (i == num_traces - 1):
                    prev_duration = duration
                    traces_per_second = (i + 1) / duration
                    print('{} | {} | {} | {}/{} | {:,.2f}       '.format(util.days_hours_mins_secs_str(duration), util.days_hours_mins_secs_str((num_traces - i) / traces_per_second), util.progress_bar(i+1, num_traces), str(i+1).rjust(len_str_num_traces), num_traces, traces_per_second), end='\r')
                    sys.stdout.flush()
            trace = next(generator)
            if map_func is not None:
                ret.append(map_func(trace))
            else:
                ret.append(trace)
        if ((trace_mode != TraceMode.DEFAULT) and (util.verbosity > 1)) or (util.verbosity > 2):
            print()
        return ret

    def prior_sample(self, *args, **kwargs):
        generator = self._prior_sample_generator(*args, **kwargs)
        return next(generator)

    def prior_distribution(self, num_traces=1000, *args, **kwargs):
        generator = self._prior_sample_generator(*args, **kwargs)
        ret = []
        time_start = time.time()
        if util.verbosity > 1:
            len_str_num_traces = len(str(num_traces))
            print('Time spent  | Time remain.| Progress             | {} | Traces/sec'.format('Trace'.ljust(len_str_num_traces * 2 + 1)))
        prev_duration = 0
        for i in range(num_traces):
            if util.verbosity > 1:
                duration = time.time() - time_start
                if (duration - prev_duration > util._print_refresh_rate) or (i == num_traces - 1):
                    prev_duration = duration
                    traces_per_second = (i + 1) / duration
                    print('{} | {} | {} | {}/{} | {:,.2f}       '.format(util.days_hours_mins_secs_str(duration), util.days_hours_mins_secs_str((num_traces - i) / traces_per_second), util.progress_bar(i+1, num_traces), str(i+1).rjust(len_str_num_traces), num_traces, traces_per_second), end='\r')
                    sys.stdout.flush()
            ret.append(next(generator))
        if util.verbosity > 1:
            print()
        return Empirical(ret, name='Prior, num_traces={}'.format(num_traces))

    def posterior_distribution(self, num_traces=1000, inference_engine=InferenceEngine.IMPORTANCE_SAMPLING, burn_in=None, *args, **kwargs):
        if (inference_engine == InferenceEngine.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK) and (self._inference_network is None):
            raise RuntimeError('Cannot run inference with inference network because there is none available. Use learn_inference_network first.')
        if burn_in is not None:
            if burn_in >= num_traces:
                raise ValueError('burn_in must be less than num_traces')
        else:
            # Default burn_in
            burn_in = int(min(num_traces / 10, 1000))

        if inference_engine == InferenceEngine.IMPORTANCE_SAMPLING:
            ret = self._prior_traces(num_traces, trace_mode=TraceMode.IMPORTANCE_SAMPLING_WITH_PRIOR, inference_network=None, map_func=lambda trace: (trace.log_importance_weight, trace.result), *args, **kwargs)
            ret = [list(t) for t in zip(*ret)]
            log_weights = ret[0]
            results = ret[1]
            name = 'Posterior, importance sampling (with prior), num_traces={}'.format(num_traces)
        elif inference_engine == InferenceEngine.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK:
            self._inference_network.eval()
            ret = self._prior_traces(num_traces, trace_mode=TraceMode.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK, inference_network=self._inference_network, map_func=lambda trace: (trace.log_importance_weight, trace.result), *args, **kwargs)
            ret = [list(t) for t in zip(*ret)]
            log_weights = ret[0]
            results = ret[1]
            name = 'Posterior, importance sampling (with learned proposal), num_traces={}'.format(num_traces)
        else:  # inference_engine == InferenceEngine.LIGHTWEIGHT_METROPOLIS_HASTINGS
            results = []
            current_trace = next(self._prior_trace_generator(trace_mode=TraceMode.LIGHTWEIGHT_METROPOLIS_HASTINGS, *args, **kwargs))
            time_start = time.time()
            traces_accepted = 1
            samples_reused = 1
            samples_all = 1
            if util.verbosity > 1:
                len_str_num_traces = len(str(num_traces))
                print('Time spent  | Time remain.| Progress             | {} | Accepted|Smp reuse| Traces/sec'.format('Trace'.ljust(len_str_num_traces * 2 + 1)))
                prev_duration = 0
            for i in range(num_traces):
                if util.verbosity > 1:
                    duration = time.time() - time_start
                    if (duration - prev_duration > util._print_refresh_rate) or (i == num_traces - 1):
                        prev_duration = duration
                        traces_per_second = (i + 1) / duration
                        print('{} | {} | {} | {}/{} | {} | {} | {:,.2f}       '.format(util.days_hours_mins_secs_str(duration), util.days_hours_mins_secs_str((num_traces - i) / traces_per_second), util.progress_bar(i+1, num_traces), str(i+1).rjust(len_str_num_traces), num_traces, '{:,.2f}%'.format(100 * (traces_accepted / (i + 1))).rjust(7), '{:,.2f}%'.format(100 * samples_reused / samples_all).rjust(7), traces_per_second), end='\r')
                        sys.stdout.flush()
                candidate_trace = next(self._prior_trace_generator(trace_mode=TraceMode.LIGHTWEIGHT_METROPOLIS_HASTINGS, metropolis_hastings_trace=current_trace, *args, **kwargs))<commit_msg>Minor<commit_after>import torch
import torch.nn as nn
import torch.optim as optim
import time
import sys
import os
import uuid
from threading import Thread
from termcolor import colored
import math
import random
import tarfile
import tempfile
import shutil

from .distributions import Empirical
from . import state, util, __version__, TraceMode, InferenceEngine, Optimizer
from .nn import ObserveEmbedding, SampleEmbedding, Batch, InferenceNetwork
from .remote import ModelServer
from .analytics import save_report


class Model(nn.Module):
    def __init__(self, name='Unnamed pyprob model'):
        super().__init__()
        self.name = name
        self._inference_network = None
        self._trace_cache_path = None
        self._trace_cache = []
        self._trace_cache_discarded_file_names = []

    def forward(self):
        raise NotImplementedError()

    def _prior_trace_generator(self, trace_mode=TraceMode.DEFAULT, inference_network=None, metropolis_hastings_trace=None, *args, **kwargs):
        while True:
            if trace_mode == TraceMode.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK:
                self._inference_network.new_trace(util.pack_observes_to_variable(kwargs['observation']).unsqueeze(0))
            state.begin_trace(self.forward, trace_mode, inference_network, metropolis_hastings_trace)
            res = self.forward(*args, **kwargs)
            trace = state.end_trace(res)
            yield trace

    def _prior_sample_generator(self, *args, **kwargs):
        while True:
            yield self.forward(*args, **kwargs)

    def _prior_traces(self, num_traces=10, trace_mode=TraceMode.DEFAULT, inference_network=None, map_func=None, *args, **kwargs):
        generator = self._prior_trace_generator(trace_mode, inference_network, *args, **kwargs)
        ret = []
        time_start = time.time()
        if ((trace_mode != TraceMode.DEFAULT) and (util.verbosity > 1)) or (util.verbosity > 2):
            len_str_num_traces = len(str(num_traces))
            print('Time spent  | Time remain.| Progress             | {} | Traces/sec'.format('Trace'.ljust(len_str_num_traces * 2 + 1)))
            prev_duration = 0
        for i in range(num_traces):
            if ((trace_mode != TraceMode.DEFAULT) and (util.verbosity > 1)) or (util.verbosity > 2):
                duration = time.time() - time_start
                if (duration - prev_duration > util._print_refresh_rate) or (i == num_traces - 1):
                    prev_duration = duration
                    traces_per_second = (i + 1) / duration
                    print('{} | {} | {} | {}/{} | {:,.2f}       '.format(util.days_hours_mins_secs_str(duration), util.days_hours_mins_secs_str((num_traces - i) / traces_per_second), util.progress_bar(i+1, num_traces), str(i+1).rjust(len_str_num_traces), num_traces, traces_per_second), end='\r')
                    sys.stdout.flush()
            trace = next(generator)
            if map_func is not None:
                ret.append(map_func(trace))
            else:
                ret.append(trace)
        if ((trace_mode != TraceMode.DEFAULT) and (util.verbosity > 1)) or (util.verbosity > 2):
            print()
        return ret

    def prior_sample(self, *args, **kwargs):
        generator = self._prior_sample_generator(*args, **kwargs)
        return next(generator)

    def prior_distribution(self, num_traces=1000, *args, **kwargs):
        generator = self._prior_sample_generator(*args, **kwargs)
        ret = []
        time_start = time.time()
        if util.verbosity > 1:
            len_str_num_traces = len(str(num_traces))
            print('Time spent  | Time remain.| Progress             | {} | Traces/sec'.format('Trace'.ljust(len_str_num_traces * 2 + 1)))
        prev_duration = 0
        for i in range(num_traces):
            if util.verbosity > 1:
                duration = time.time() - time_start
                if (duration - prev_duration > util._print_refresh_rate) or (i == num_traces - 1):
                    prev_duration = duration
                    traces_per_second = (i + 1) / duration
                    print('{} | {} | {} | {}/{} | {:,.2f}       '.format(util.days_hours_mins_secs_str(duration), util.days_hours_mins_secs_str((num_traces - i) / traces_per_second), util.progress_bar(i+1, num_traces), str(i+1).rjust(len_str_num_traces), num_traces, traces_per_second), end='\r')
                    sys.stdout.flush()
            ret.append(next(generator))
        if util.verbosity > 1:
            print()
        return Empirical(ret, name='Prior, num_traces={}'.format(num_traces))

    def posterior_distribution(self, num_traces=1000, inference_engine=InferenceEngine.IMPORTANCE_SAMPLING, burn_in=None, *args, **kwargs):
        if (inference_engine == InferenceEngine.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK) and (self._inference_network is None):
            raise RuntimeError('Cannot run inference with inference network because there is none available. Use learn_inference_network first.')
        if burn_in is not None:
            if burn_in >= num_traces:
                raise ValueError('burn_in must be less than num_traces')
        else:
            # Default burn_in
            burn_in = int(min(num_traces / 10, 1000))

        if inference_engine == InferenceEngine.IMPORTANCE_SAMPLING:
            ret = self._prior_traces(num_traces, trace_mode=TraceMode.IMPORTANCE_SAMPLING_WITH_PRIOR, inference_network=None, map_func=lambda trace: (trace.log_importance_weight, trace.result), *args, **kwargs)
            ret = [list(t) for t in zip(*ret)]
            log_weights = ret[0]
            results = ret[1]
            name = 'Posterior, importance sampling (with prior), num_traces={}'.format(num_traces)
        elif inference_engine == InferenceEngine.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK:
            self._inference_network.eval()
            ret = self._prior_traces(num_traces, trace_mode=TraceMode.IMPORTANCE_SAMPLING_WITH_INFERENCE_NETWORK, inference_network=self._inference_network, map_func=lambda trace: (trace.log_importance_weight, trace.result), *args, **kwargs)
            ret = [list(t) for t in zip(*ret)]
            log_weights = ret[0]
            results = ret[1]
            name = 'Posterior, importance sampling (with learned proposal, training_traces={}), num_traces={}'.format(self._inference_network._total_train_traces, num_traces)
        else:  # inference_engine == InferenceEngine.LIGHTWEIGHT_METROPOLIS_HASTINGS
            results = []
            current_trace = next(self._prior_trace_generator(trace_mode=TraceMode.LIGHTWEIGHT_METROPOLIS_HASTINGS, *args, **kwargs))
            time_start = time.time()
            traces_accepted = 1
            samples_reused = 1
            samples_all = 1
            if util.verbosity > 1:
                len_str_num_traces = len(str(num_traces))
                print('Time spent  | Time remain.| Progress             | {} | Accepted|Smp reuse| Traces/sec'.format('Trace'.ljust(len_str_num_traces * 2 + 1)))
                prev_duration = 0
            for i in range(num_traces):
                if util.verbosity > 1:
                    duration = time.time() - time_start
                    if (duration - prev_duration > util._print_refresh_rate) or (i == num_traces - 1):
                        prev_duration = duration
                        traces_per_second = (i + 1) / duration
                        print('{} | {} | {} | {}/{} | {} | {} | {:,.2f}       '.format(util.days_hours_mins_secs_str(duration), util.days_hours_mins_secs_str((num_traces - i) / traces_per_second), util.progress_bar(i+1, num_traces), str(i+1).rjust(len_str_num_traces), num_traces, '{:,.2f}%'.format(100 * (traces_accepted / (i + 1))).rjust(7), '{:,.2f}%'.format(100 * samples_reused / samples_all).rjust(7), traces_per_second), end='\r')
                        sys.stdout.flush()
                candidate_trace = next(self._prior_trace_generator(trace_mode=TraceMode.LIGHTWEIGHT_METROPOLIS_HASTINGS, metropolis_hastings_trace=current_trace, *args, **kwargs))
{'fraction_non_alphanumeric': 0.1109042388133805, 'fraction_numerical': 0.008688636504685657, 'mean_word_length': 2.7257803468208093, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 33, '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': '297777', 'n_tokens_mistral': 5058, 'n_tokens_neox': 4817, 'n_words': 989}