<commit_before>import ast
import hashlib
import inspect
import os.path

from collections import OrderedDict
from functools import partial


def _parse_python_str_constants(source):
    constants = {}
    parsed = ast.parse(source)
    nv = ast.NodeVisitor()
    nv.visit_Module = lambda n: nv.generic_visit(n)
    def visit_Assign(n):
        if len(n.targets) != 1:
            return

        target = n.targets[0]
        if not isinstance(target, ast.Name):
            return

        if not isinstance(n.value, ast.Str):
            return

        constants[target.id] = n.value.s

    nv.visit_Assign = visit_Assign
    nv.visit(parsed)
    return constants

def _reify(resolution):
    init_rest = resolution.get('init_rest', [])
    loaded_fn = resolution.get('loaded_fn', None)

    init = {}
    if init_rest:
        init = eval("dict(%s" % init_rest[0])

    loaded = loaded_fn()
    is_loaded_a_class = False
    
    if callable(loaded):
        sig = inspect.signature(loaded)
        needed_params = OrderedDict(sig.parameters.items())
        # Only allow unspecified leading_value to be given.
        for k in init.keys():
            del needed_params[k]
        members = []
        if isinstance(loaded, type):
            members = inspect.getmembers(loaded, inspect.isfunction)
            is_loaded_a_class = True
        model_fn = partial(loaded, **init)
    else:
        members = inspect.getmembers(loaded, inspect.ismethod)
        needed_params = OrderedDict()
        model_fn = lambda: loaded

    members = [(k, v) for (k, v) in members if not k.startswith('_')]

    if resolution['can_refresh'] and is_loaded_a_class:
        models = []
        orig_model_fn = model_fn
        def wrapped_model_fn(*a, **kw):
            model = orig_model_fn(*a, **kw)
            # TODO(adamb) Should actually use weak references to avoid a memory leak!
            models.append(model)
            return model
        model_fn = wrapped_model_fn

        def refresh_model():
            reloaded = loaded_fn()
            for model in models:
                model.__class__ = reloaded
                print("Replacing __class__", model, id(reloaded), reloaded)
                if hasattr(model, '__tfi_refresh__'):
                    model.__tfi_refresh__()

        resolution['refresh_fn'] = refresh_model

    resolution['model_fn_needed_params'] = needed_params
    resolution['model_members'] = members
    resolution['model_fn'] = model_fn
    
    return resolution

def resolve_exported(model_class_from_path_fn, leading_value):
    return _reify({<commit_msg>Track models to reload in a WeakSet to save RAM<commit_after>import ast
import hashlib
import inspect
import os.path
import weakref

from collections import OrderedDict
from functools import partial


def _parse_python_str_constants(source):
    constants = {}
    parsed = ast.parse(source)
    nv = ast.NodeVisitor()
    nv.visit_Module = lambda n: nv.generic_visit(n)
    def visit_Assign(n):
        if len(n.targets) != 1:
            return

        target = n.targets[0]
        if not isinstance(target, ast.Name):
            return

        if not isinstance(n.value, ast.Str):
            return

        constants[target.id] = n.value.s

    nv.visit_Assign = visit_Assign
    nv.visit(parsed)
    return constants

def _reify(resolution):
    init_rest = resolution.get('init_rest', [])
    loaded_fn = resolution.get('loaded_fn', None)

    init = {}
    if init_rest:
        init = eval("dict(%s" % init_rest[0])

    loaded = loaded_fn()
    is_loaded_a_class = False
    
    if callable(loaded):
        sig = inspect.signature(loaded)
        needed_params = OrderedDict(sig.parameters.items())
        # Only allow unspecified leading_value to be given.
        for k in init.keys():
            del needed_params[k]
        members = []
        if isinstance(loaded, type):
            members = inspect.getmembers(loaded, inspect.isfunction)
            is_loaded_a_class = True
        model_fn = partial(loaded, **init)
    else:
        members = inspect.getmembers(loaded, inspect.ismethod)
        needed_params = OrderedDict()
        model_fn = lambda: loaded

    members = [(k, v) for (k, v) in members if not k.startswith('_')]

    if resolution['can_refresh'] and is_loaded_a_class:
        models = weakref.WeakSet()
        orig_model_fn = model_fn

        def _wrapped_model_fn(*a, **kw):
            model = orig_model_fn(*a, **kw)
            # TODO(adamb) Should actually use weak references to avoid a memory leak!

            models.add(model)
            return model
        model_fn = _wrapped_model_fn

        def refresh_model():
            reloaded = loaded_fn()
            for model in models:
                previous_class = model.__class__
                try:
                    model.__class__ = reloaded
                    print("Replacing __class__", model, id(reloaded), reloaded)
                    if hasattr(model, '__tfi_refresh__'):
                        model.__tfi_refresh__()
                except:
                    model.__class__ = previous_class
                    raise

        resolution['refresh_fn'] = refresh_model

    resolution['model_fn_needed_params'] = needed_params
    resolution['model_members'] = members
    resolution['model_fn'] = model_fn
    
    return resolution

def resolve_exported(model_class_from_path_fn, leading_value):
    return _reify({
{'fraction_non_alphanumeric': 0.09225972120322817, 'fraction_numerical': 0.0011005135730007337, 'mean_word_length': 2.3048484848484847, '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': '5596979', 'n_tokens_mistral': 1599, 'n_tokens_neox': 1548, 'n_words': 384}