kalfa 3.0.0.dev0__py3-none-any.whl

kalfa/core/make.py

@@ -0,0 +1,121 @@
+from . import state, wrap
+from .errors import BuildError, KalfaError
+
+
+def as_list(value):
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return [value]
+    return list(value)
+
+
+def build_model(node, module, keys, ctx):
+    name = node.name
+    data_inputs = getattr(ctx.data, "inputs", None) if ctx.data is not None else None
+    inputs = as_list(keys.get("inputs")) or list(data_inputs or ("x",))
+    outputs = as_list(keys.get("outputs")) or [name]
+    init = keys.get("init")
+    if init is not None:
+        if not callable(init):
+            raise BuildError(f"{node.path}: init must be a call such as {{use: normal_init}}, got {init!r}")
+        init(module)
+    weights = keys.get("weights")
+    if weights is not None:
+        load_weights(node, module, weights)
+    freeze = keys.get("freeze", False)
+    entry = wrap.ModelEntry(name, module, inputs, outputs, freeze=freeze, weights=weights, init=init)
+    if "freeze" in keys:
+        entry.set_freeze(freeze)
+    ema = keys.get("ema")
+    if ema is not None and ema is not False:
+        entry.ema = ema_copy(node, module, ema)
+    entry.path = node.path
+    return entry
+
+
+def ema_copy(node, module, ema):
+    if isinstance(ema, dict):
+        decay = ema.get("decay", 0.999)
+        every = ema.get("every", 1)
+    elif ema is True:
+        decay, every = 0.999, 1
+    else:
+        decay, every = ema, 1
+    try:
+        return wrap.EmaCopy(module, float(decay), int(every))
+    except (TypeError, ValueError) as exc:
+        raise BuildError(f"{node.path}: ema takes a decay such as 0.999 or {{decay, every}}, got {ema!r}") from exc
+
+
+def load_weights(node, module, weights):
+    spec = {"run": weights} if isinstance(weights, str) else weights
+    if not isinstance(spec, dict) or not spec.get("run"):
+        raise BuildError(f"{node.path}: weights takes {{run, model, which}}, got {weights!r}")
+    model = spec.get("model") or node.name
+    try:
+        path = state.checkpoint_path(spec["run"], spec.get("which"))
+        checkpoint = state.read_checkpoint(path)
+    except KalfaError:
+        raise
+    except Exception as exc:
+        raise BuildError(f"{node.path}: weights could not be read from {spec['run']}: {exc}") from exc
+    nodes = checkpoint.get("nodes", {}) if isinstance(checkpoint, dict) else {}
+    key = f"models.{model}"
+    if key not in nodes:
+        raise BuildError(f"{node.path}: {path} holds no {key}; it holds {sorted(nodes)}")
+    try:
+        module.load_state_dict(nodes[key])
+    except Exception as exc:
+        raise BuildError(f"{node.path}: the weights of {key} in {path} do not fit {node.use}: {exc}") from exc
+
+
+def build_optimizer(node, kwargs, keys, signature, ctx):
+    over = groups_of(node, keys.get("over"))
+    entry = wrap.OptimizerEntry(node.name, node.target, kwargs, over, minimizes=keys.get("minimizes"),
+                                schedule=keys.get("schedule"))
+    entry.path = node.path
+    link_emas(entry, over, ctx)
+    try:
+        if callable(getattr(entry, "bind", None)):
+            entry.bind(ctx)
+        else:
+            entry.build()
+    except KalfaError:
+        raise
+    except Exception as exc:
+        raise BuildError(f"{node.path}: {node.use} could not be built: {type(exc).__name__}: {exc}") from exc
+    return entry
+
+
+def covers(model, module):
+    return any(module is part for part in model.module.modules())
+
+
+def link_emas(entry, over, ctx):
+    emas = getattr(entry, "emas", None)
+    if not isinstance(emas, list):
+        return
+    modules = [item.of if isinstance(item, wrap.Group) else item for item in over]
+    for model in ctx.models.values():
+        if model.ema is None or any(model.ema is known for known in emas):
+            continue
+        if any(covers(model, module) for module in modules):
+            emas.append(model.ema)
+
+
+def groups_of(node, over):
+    if over is None:
+        raise BuildError(f"{node.path}: over names the modules to optimize")
+    items = over if isinstance(over, list) else [over]
+    groups = []
+    for item in items:
+        if isinstance(item, dict):
+            rest = dict(item)
+            of = rest.pop("of", None)
+            if of is None:
+                raise BuildError(f"{node.path}: a group in over needs of: '@models.<name>'")
+            groups.append(wrap.Group(of, rest.pop("match", None), **rest))
+        else:
+            groups.append(item)
+    return groups

kalfa/core/plan.py

@@ -0,0 +1,156 @@
+from pathlib import Path
+
+from .build import fixed_of, mutable_paths, node_signature
+from .context import train_of
+from .errors import RunError
+from .record import version as kalfa_version
+from .refs import jsonable, ref_json, ref_parts
+
+
+FORMAT = 1
+SET_PREFIX = {"train": "train", "valid": "val", "test": "test"}
+DEFAULT_ON = ("train", "valid", "test")
+
+
+def plan_of(tree, version=None):
+    settings = tree.settings or {}
+    nodes = {}
+    for path, node in tree.nodes.items():
+        signature = node_signature(node.target, node.slot) if node.target is not None else None
+        checked = signature is not None and not signature.accepts_extra
+        nodes[path] = {
+            "slot": node.slot,
+            "use": node.use,
+            "params": jsonable(node.params),
+            "keys": jsonable(node.keys),
+            "fixed": list(fixed_of(node)),
+            "signature": "checked" if checked else "open",
+        }
+    return {
+        "format": FORMAT,
+        "kalfa": version if version is not None else kalfa_version(),
+        "seed": jsonable(settings.get("seed")),
+        "device": jsonable(settings.get("device")),
+        "record": jsonable(settings.get("record")),
+        "nodes": nodes,
+        "data": data_of(tree),
+        "train": jsonable(settings.get("train") or {}),
+        "rules": rules_of(tree),
+        "series": series_of(tree),
+        "mutable": mutable_paths(tree),
+    }
+
+
+def data_of(tree):
+    node = tree.nodes.get("data")
+    if node is None:
+        return {"inputs": [], "exposes": []}
+    target = node.target
+    inputs = getattr(target, "inputs", None) if target is not None else None
+    exposes = getattr(target, "exposes", None) if target is not None else None
+    out = {"inputs": list(inputs) if inputs else ["x"], "exposes": list(exposes or ())}
+    fields = fields_of(tree, node)
+    if fields is not None:
+        out["fields"] = fields
+    return out
+
+
+def fields_of(tree, node):
+    file = node.params.get("file")
+    if not isinstance(file, str) or not file:
+        return None
+    candidates = [Path(file)]
+    for directory in getattr(tree.config, "directories", None) or ():
+        candidates.append(Path(directory) / file)
+    for candidate in candidates:
+        if candidate.is_file():
+            try:
+                return read_header(candidate)
+            except Exception:
+                return None
+    return None
+
+
+def read_header(path):
+    suffix = path.suffix.lower()
+    if suffix in (".parquet", ".pq"):
+        import pyarrow.parquet as parquet
+
+        schema = parquet.read_schema(str(path))
+    elif suffix in (".csv", ".tsv"):
+        import pyarrow.csv as csv
+
+        options = csv.ParseOptions(delimiter="\t") if suffix == ".tsv" else None
+        reader = csv.open_csv(str(path), parse_options=options)
+        try:
+            schema = reader.schema
+        finally:
+            close = getattr(reader, "close", None)
+            if callable(close):
+                close()
+    else:
+        return None
+    return {name: str(schema.field(name).type) for name in schema.names}
+
+
+def call_of(tree, path, value):
+    if isinstance(value, str):
+        return {"use": value}
+    node = tree.nodes.get(path)
+    if node is None or not isinstance(value, dict) or "use" not in value:
+        return jsonable(value)
+    return {"use": value["use"], **jsonable(node.keys), **jsonable(node.params)}
+
+
+def patch_value(value):
+    if ref_parts(value) is None:
+        return jsonable(value)
+    try:
+        return ref_json(value)
+    except RunError:
+        return jsonable(value)
+
+
+def rules_of(tree):
+    out = []
+    for spec in tree.rules:
+        out.append({
+            "id": spec.name,
+            "when": call_of(tree, f"rules.{spec.name}.when", spec.when),
+            "after": spec.after,
+            "once": bool(spec.once),
+            "patch": [{"path": str(path), "value": patch_value(value)} for path, value in spec.set.items()],
+            "stop": bool(spec.stop),
+            "run": call_of(tree, f"rules.{spec.name}.run", spec.run),
+        })
+    return out
+
+
+def sets_of(node):
+    on = node.keys.get("on")
+    if on is None:
+        on = getattr(node.target, "on", None) if node.target is not None else None
+    if on is None:
+        on = DEFAULT_ON
+    if isinstance(on, str):
+        on = (on,)
+    return [SET_PREFIX.get(str(name), str(name)) for name in on]
+
+
+def series_of(tree):
+    series = []
+    for node in tree.nodes.values():
+        if node.slot == "loss":
+            series.append(f"train/{node.name}")
+            series.extend(f"{prefix}/{node.name}" for prefix in sets_of(node))
+        elif node.slot == "metric":
+            series.extend(f"{prefix}/{node.name}" for prefix in sets_of(node))
+    for node in tree.nodes.values():
+        if node.slot == "optimizer":
+            series.extend((f"lr/{node.name}", f"minimizes/{node.name}"))
+    track = train_of(tree).get("track")
+    if isinstance(track, str):
+        track = [track]
+    for path in track or ():
+        series.append(f"track/{path}")
+    return list(dict.fromkeys(series))

kalfa/core/predict.py

@@ -0,0 +1,290 @@
+import fnmatch
+
+import numpy
+import pandas
+import torch
+
+from .errors import KalfaError, RunError
+from .record import atomic
+
+
+WIDE = 32
+
+
+def batch_size(batch):
+    if isinstance(batch, dict):
+        values = batch.values()
+    elif isinstance(batch, (list, tuple)):
+        values = batch
+    else:
+        values = [batch]
+    for value in values:
+        shape = getattr(value, "shape", None)
+        if shape is not None and len(shape):
+            return int(shape[0])
+    return 1
+
+
+def enter(ctx, batch):
+    ctx.batches += 1
+    return batch if ctx.device is None else ctx.device.move(batch)
+
+
+def empty(loader):
+    try:
+        return len(loader) == 0
+    except TypeError:
+        return False
+
+
+def as_matrix(value):
+    if isinstance(value, torch.Tensor):
+        value = value.detach().cpu().numpy()
+    value = numpy.asarray(value)
+    if value.ndim == 0:
+        value = value.reshape(1)
+    return value.reshape(len(value), -1)
+
+
+def spread(label, values):
+    values = numpy.asarray(values)
+    if values.ndim <= 1:
+        return {label: values}
+    if values.shape[1] == 1:
+        return {label: values[:, 0]}
+    return {f"{label}_{position}": values[:, position] for position in range(values.shape[1])}
+
+
+def target_names(data, batch=None):
+    names = [str(name) for name in getattr(data, "target_names", None) or ()]
+    return names if batch is None else [name for name in names if name in batch]
+
+
+def is_label(matrix):
+    return bool(numpy.issubdtype(matrix.dtype, numpy.integer) or matrix.dtype == bool)
+
+
+def labelled(data, name, truths):
+    truth = truths.get(name)
+    if truth is not None:
+        return is_label(truth)
+    asks = getattr(data, "is_label", None)
+    return bool(asks(name)) if callable(asks) else False
+
+
+def rescales(data, name):
+    asks = getattr(data, "rescales", None)
+    return bool(asks(name)) if callable(asks) else False
+
+
+def inverse(data, name, matrix):
+    transform = getattr(data, "inverse", None)
+    if not callable(transform):
+        return matrix
+    values = matrix[:, 0] if matrix.ndim == 2 and matrix.shape[1] == 1 else matrix
+    return numpy.asarray(transform(name, values))
+
+
+def decoded(data, name, indexes):
+    decode = getattr(data, "decode", None)
+    if not callable(decode):
+        return None
+    return numpy.asarray(decode(name, indexes))
+
+
+def pairings(ctx, outputs, targets):
+    pairs = {}
+    for entry in [*ctx.losses.values(), *ctx.metrics.values()]:
+        pred, target = getattr(entry, "pred", None), getattr(entry, "target", None)
+        if pred in outputs and isinstance(target, str):
+            pairs.setdefault(pred, target)
+    for name in outputs:
+        if name in targets:
+            pairs.setdefault(name, name)
+    if len(targets) == 1:
+        for name in outputs:
+            pairs.setdefault(name, targets[0])
+    return pairs
+
+
+def rows_of(data, set_name, count, loader=None, own=True):
+    rows = getattr(getattr(loader, "dataset", None), "rows", None)
+    if rows is not None and len(rows) == count:
+        return numpy.asarray(rows)
+    rows = getattr(data, "rows", None)
+    if own and callable(rows):
+        try:
+            found = numpy.asarray(rows(set_name))
+        except KalfaError:
+            found = None
+        if found is not None and len(found) == count:
+            return found
+    return numpy.arange(count)
+
+
+def aligned(frame, rows, count):
+    if frame is None or len(frame) == count:
+        return frame
+    try:
+        picked = frame.loc[list(rows)]
+    except KeyError:
+        picked = None
+    if picked is None or len(picked) != count:
+        raise RunError(f"{count} predicted rows cannot be aligned with the {len(frame)} rows of the frame; the "
+                       f"frame needs a unique index")
+    return picked
+
+
+def spectator_columns(data, set_name, count, frame=None):
+    if frame is not None:
+        patterns = [str(pattern) for pattern in getattr(data, "spectator_patterns", None) or ()]
+        return {column: frame[column].to_numpy() for column in frame.columns
+                if any(fnmatch.fnmatchcase(str(column), pattern) for pattern in patterns)}
+    extra = data.spectators(set_name) if callable(getattr(data, "spectators", None)) else None
+    if extra is None:
+        return {}
+    if isinstance(extra, pandas.DataFrame):
+        extra = {column: extra[column].to_numpy() for column in extra.columns}
+    return {name: numpy.asarray(values) for name, values in dict(extra).items() if len(values) == count}
+
+
+def collect_outputs(ctx, loader, entry=None):
+    data = ctx.data
+    names = None
+    guesses, truths = {}, {}
+    with torch.no_grad():
+        for batch in loader:
+            batch = enter(ctx, batch)
+            if names is None:
+                names = target_names(data, batch)
+            outputs = entry.call(batch) if entry is not None else ctx.step.predict(batch, ctx)
+            for name, value in outputs.items():
+                guesses.setdefault(name, []).append(as_matrix(value))
+            for name in names:
+                truths.setdefault(name, []).append(as_matrix(batch[name]))
+    guesses = {name: numpy.concatenate(pieces) for name, pieces in guesses.items()}
+    truths = {name: numpy.concatenate(pieces) for name, pieces in truths.items()}
+    return guesses, truths
+
+
+def spread_truth(columns, data, name, matrix):
+    columns.update(spread(name, inverse(data, name, matrix) if rescales(data, name) else matrix))
+
+
+def truth_columns(columns, data, truths, frame):
+    if frame is not None:
+        for name in target_names(data):
+            matrix = truths.get(name)
+            if matrix is not None and matrix.shape[1] > 1 and not is_label(matrix):
+                spread_truth(columns, data, name, matrix)
+            elif name in frame.columns:
+                columns[name] = frame[name].to_numpy()
+        return
+    for name, matrix in truths.items():
+        if is_label(matrix):
+            indexes = matrix[:, 0] if matrix.shape[1] == 1 else matrix.argmax(axis=1)
+            labels = decoded(data, name, indexes)
+            columns[name] = indexes if labels is None else labels
+        else:
+            spread_truth(columns, data, name, matrix)
+
+
+def softmax(matrix):
+    shifted = matrix.astype(numpy.float64) - matrix.max(axis=1, keepdims=True)
+    weights = numpy.exp(shifted)
+    return weights / weights.sum(axis=1, keepdims=True)
+
+
+def sigmoid(values):
+    return 0.5 * (1.0 + numpy.tanh(0.5 * values.astype(numpy.float64)))
+
+
+def class_names(data, target, count):
+    indexes = numpy.arange(count)
+    try:
+        labels = decoded(data, target, indexes)
+    except (KalfaError, IndexError, KeyError, ValueError):
+        labels = None
+    if labels is None or len(labels) != count:
+        return [str(index) for index in indexes]
+    return [str(label) for label in labels]
+
+
+def class_scores(data, target, width):
+    if width == 1:
+        return True
+    try:
+        labels = decoded(data, target, numpy.arange(width))
+    except (KalfaError, IndexError, KeyError, ValueError):
+        return False
+    return labels is None or len(labels) == width
+
+
+def probability_columns(columns, data, name, target, matrix):
+    if matrix.shape[1] > 1:
+        probabilities = softmax(matrix)
+        for position, label in enumerate(class_names(data, target, matrix.shape[1])):
+            columns[f"prob_{name}_{label}"] = probabilities[:, position]
+    else:
+        positive = class_names(data, target, 2)[1]
+        columns[f"prob_{name}_{positive}"] = sigmoid(matrix[:, 0])
+
+
+def guess_columns(columns, ctx, guesses, truths):
+    data = ctx.data
+    targets = target_names(data) or list(truths)
+    pairs = pairings(ctx, list(guesses), targets)
+    for name, matrix in guesses.items():
+        target = pairs.get(name)
+        if target is not None and target not in targets:
+            target = None
+        if target is not None and labelled(data, target, truths) and not class_scores(data, target, matrix.shape[1]):
+            target = None
+        if target is None:
+            if matrix.shape[1] > WIDE:
+                continue
+            columns.update(spread(f"pred_{name}", matrix))
+        elif labelled(data, target, truths):
+            indexes = matrix.argmax(axis=1) if matrix.shape[1] > 1 else (matrix[:, 0] > 0).astype(numpy.int64)
+            labels = decoded(data, target, indexes)
+            columns[f"pred_{name}"] = indexes if labels is None else labels
+            if labels is not None:
+                columns[f"raw_{name}"] = indexes
+            probability_columns(columns, data, name, target, matrix)
+        elif rescales(data, target):
+            columns.update(spread(f"pred_{name}", inverse(data, target, matrix)))
+            columns.update(spread(f"raw_{name}", matrix))
+        else:
+            columns.update(spread(f"pred_{name}", matrix))
+
+
+def predictions_table(ctx, set_name="test", loader=None, frame=None, entry=None):
+    own = loader is None
+    if own:
+        loader = ctx.loaders().get(set_name)
+    if loader is None or empty(loader):
+        return pandas.DataFrame()
+    ctx.set_modes(False)
+    ctx.set_name = set_name
+    data = ctx.data
+    guesses, truths = collect_outputs(ctx, loader, entry)
+    if not guesses:
+        return pandas.DataFrame()
+    count = len(next(iter(guesses.values())))
+    rows = rows_of(data, set_name, count, loader, own)
+    frame = aligned(frame, rows, count)
+    columns = {"row": rows}
+    truth_columns(columns, data, truths, frame)
+    guess_columns(columns, ctx, guesses, truths)
+    for name, values in spectator_columns(data, set_name, count, frame).items():
+        columns.setdefault(name, values)
+    return pandas.DataFrame(columns)
+
+
+def write_predictions(ctx, set_name):
+    table = predictions_table(ctx, set_name)
+    if ctx.record is not None and len(table):
+        name = "predictions.parquet" if set_name == "test" else f"predictions_{set_name}.parquet"
+        with atomic(ctx.record.path(name)) as temporary:
+            table.to_parquet(temporary, index=False)
+    return table

kalfa/core/protocols.py

@@ -0,0 +1,122 @@
+from torch import nn
+
+
+class Objective(nn.Module):
+    on = ("train", "valid", "test")
+    every = 1
+    fixed = ()
+    name = ""
+    path = ""
+
+    def bind(self, ctx):
+        pass
+
+    def forward(self, outputs, batch, ctx):
+        raise NotImplementedError
+
+
+class ContextMetric:
+    on = ("train", "valid", "test")
+    every = 1
+    name = ""
+    path = ""
+
+    def bind(self, ctx):
+        pass
+
+    def reset(self):
+        raise NotImplementedError
+
+    def update(self, outputs, batch, ctx):
+        raise NotImplementedError
+
+    def compute(self):
+        raise NotImplementedError
+
+
+class Step:
+    order = None
+    amp = False
+    grad_clip = None
+    accumulate = 1
+    stepped = True
+    fixed = ()
+    persist = ()
+    name = "step"
+    path = "train.step"
+
+    def bind(self, ctx):
+        pass
+
+    def train_turn(self, ctx):
+        raise NotImplementedError
+
+    def train_batch(self, batch, ctx):
+        raise NotImplementedError
+
+    def eval_batch(self, batch, ctx):
+        raise NotImplementedError
+
+    def predict(self, batch, ctx):
+        raise NotImplementedError
+
+    def state_dict(self):
+        return {name: getattr(self, name) for name in self.persist}
+
+    def load_state_dict(self, state):
+        for name, value in (state or {}).items():
+            setattr(self, name, value)
+
+
+class Trigger:
+    persist = ()
+    fixed = ()
+    name = ""
+    path = ""
+
+    def bind(self, ctx):
+        pass
+
+    def fires(self, ctx):
+        raise NotImplementedError
+
+    def reset(self):
+        pass
+
+    def describe(self):
+        return type(self).__name__
+
+    def state_dict(self):
+        return {name: getattr(self, name) for name in self.persist}
+
+    def load_state_dict(self, state):
+        for name, value in (state or {}).items():
+            setattr(self, name, value)
+
+
+class Action:
+    name = ""
+    path = ""
+
+    def bind(self, ctx):
+        pass
+
+    def __call__(self, ctx):
+        raise NotImplementedError
+
+
+class Logger:
+    name = ""
+    path = ""
+
+    def bind(self, ctx):
+        pass
+
+    def turn(self, ctx, line):
+        pass
+
+    def update(self, ctx, line):
+        pass
+
+    def close(self, ctx):
+        pass