<commit_before>from typing import Dict, List, Callable, Optional, Any, Union, Iterable, Set
from typing import Generic, Sequence, Tuple, TypeVar
import numpy
import contextlib
import srsly
from pathlib import Path
import copy

from .backends import NumpyOps, CupyOps, get_current_ops
from .optimizers import Optimizer  # noqa: F401
from .backends.mem import Memory
from .shims import Shim
from .util import copy_array, get_width, create_thread_local
from .types import Array


InT = TypeVar("InT")
OutT = TypeVar("OutT")


def create_init(initializers: Dict[str, Callable]) -> Callable:
    """Create an init function, given a dictionary of parameter initializers."""

    def init(
        model: Model, X: Optional[Array] = None, Y: Optional[Array] = None
    ) -> None:
        if X is not None:
            model.set_dim("nI", get_width(X))
        if Y is not None:
            model.set_dim("nO", get_width(Y))
        W = model.ops.alloc_f2d(model.get_dim("nO"), model.get_dim("nI"))
        b = model.ops.alloc_f1d(model.get_dim("nO"))
        if "W" in initializers:
            initializers["W"](W, inplace=True)
        if "b" in initializers:
            initializers["b"](b, inplace=True)
        model.set_param("W", W)
        model.set_param("b", b)

    return init


class Model(Generic[InT, OutT]):
    """Class for implementing Thinc models and layers."""

    global_id: int = 0
    _thread_local = create_thread_local({"operators": {}})

    name: str
    ops: Union[NumpyOps, CupyOps]
    id: int
    _func: Callable
    _init: Callable
    _mem: Memory
    _params: Dict[str, Optional[bool]]
    _dims: Dict[str, Optional[int]]
    _grads: Dict[str, Optional[bool]]
    _layers: List["Model"]
    _shims: List[Shim]
    _attrs: Dict[str, Any]

    # This "locks" the class, so we get an error if you try to assign to
    # an unexpected variable.
    __slots__ = [
        "name",
        "id",
        "ops",
        "_func",
        "_init",
        "_mem",
        "_params",
        "_dims",
        "_grads",
        "_attrs",
        "_refs",
        "_layers",
        "_shims",
    ]

    def __init__(
        self,
        name: str,
        forward: Callable,
        *,
        init: Callable = lambda *a, **k: None,
        dims: Dict[str, Optional[int]] = {},
        params: Dict[str, Optional[Array]] = {},
        grads: Dict[str, Optional[Array]] = {},
        layers: Sequence["Model"] = [],
        shims: List[Shim] = [],
        attrs: Dict[str, Any] = {},
        refs: Dict[str, Optional["Model"]] = {},
        ops: Optional[Union[NumpyOps, CupyOps]] = None,
    ):
        """Initialize a new model."""
        self.name = name
        # Assign to callable attrs: https://github.com/python/mypy/issues/2427
        setattr(self, "_func", forward)
        setattr(self, "_init", init)
        self.ops = ops if ops is not None else get_current_ops()
        self._mem = Memory(self.ops)
        self._dims = dict(dims)
        self._attrs = dict(attrs)
        self._refs = dict(refs)
        self._layers = list(layers)
        self._shims = list(shims)
        self.__class__.global_id += 1
        self.id = self.__class__.global_id
        self._params = {}
        self._grads = {}
        for name, value in params.items():
            self._params[name] = None
            if value is not None:
                self.set_param(name, value)
        for name, value in grads.items():
            self._grads[name] = None
            if value is not None:
                self.set_grad(name, value)

    @property
    def layers(self) -> List["Model"]:
        """A list of child layers of the model. You can append to it to add
        layers but not reassign it.
        """
        return self._layers

    @property
    def shims(self) -> List[Shim]:
        return self._shims

    @property
    def param_names(self) -> Tuple[str, ...]:
        """Get the names of registered parameter (including unset)."""
        return tuple(self._params.keys())

    @property
    def grad_names(self) -> Tuple[str, ...]:
        """Get the names of parameters with registered gradients (including unset)."""
        return tuple([name for name in self.param_names if self.has_grad(name)])

    @property
    def dim_names(self) -> Tuple[str, ...]:
        """Get the names of registered dimensions (including unset)."""
        return tuple(self._dims.keys())

    @property
    def attr_names(self) -> Tuple[str, ...]:
        """Get the names of attributes."""
        return tuple(self._attrs.keys())

    @property
    def ref_names(self) -> Tuple[str, ...]:
        """Get the names of registered node references (including unset)."""
        return tuple(self._refs.keys())

    @classmethod
    @contextlib.contextmanager
    def define_operators(cls, operators: Dict[str, Callable]):
        """Bind arbitrary binary functions to Python operators, for use in any
        `Model` instance. Can (and should) be used as a contextmanager.

        EXAMPLE:
            with Model.define_operators({">>": chain}):
                model = ReLu(512) >> ReLu(512) >> Softmax()
        """
        curr_operators = dict(cls._thread_local.operators)
        cls._thread_local.operators = dict(operators)
        yield
        cls._thread_local.operators = dict(curr_operators)

    def has_dim(self, name: str) -> Optional[bool]:
        """Check whether the model has a dimension of a given name. If the
        dimension is registered but the value is unset, returns None.
        """
        if name not in self._dims:
            return False
        elif self._dims[name] is not None:
            return True
        else:
            return None

    def get_dim(self, name: str) -> int:
        """Retrieve the value of a dimension of the given name."""
        if name not in self._dims:
            raise KeyError(f"Cannot get dimension '{name}' for model '{self.name}'")
        value = self._dims[name]
        if value is None:
            err = f"Cannot get dimension '{name}' for model '{self.name}': value unset"
            raise ValueError(err)
        else:
            return value

    def set_dim(self, name: str, value: int) -> None:
        """Set a value for a dimension."""
        if name not in self._dims:
            raise KeyError(f"Cannot set dimension '{name}' for model '{self.name}'.")
        old_value = self._dims[name]
        if old_value is not None and old_value != value:
            err = f"Attempt to change dimension '{name}' for model '{self.name}' from {old_value} to {value}"
            raise ValueError(err)
        self._dims[name] = value

    def has_param(self, name: str) -> Optional[bool]:
        """Check whether the model has a weights parameter of the given name.

        Returns None if the parameter is registered but currently unset.
        """
        if name not in self._params:
            return False
        elif self._params[name] is not None:
            return True
        else:
            return None

    def get_param(self, name: str) -> Array:
        """Retrieve a weights parameter by name."""
        if name not in self._params:
            raise KeyError(f"Unknown param: '{name}' for model '{self.name}'.")
        key = (self.id, name)
        if key not in self._mem:
            raise KeyError(
                f"Parameter '{name}' for model '{self.name}' has not been allocated yet."
            )
        return self._mem[key]

    def set_param(self, name: str, value: Optional[Array]) -> None:
        """Set a weights parameter's value."""
        if value is None:
            self._params[name] = None
        else:
            key = (self.id, name)
            if key not in self._mem:
                self._mem.add(key, value.shape)
            data = self._mem[(self.id, name)]
            try:
                copy_array(dst=data, src=value)
            except ValueError as e:
                raise ValueError(
                    f"Cannot set param '{name}' for model '{self.name}': {e}."
                )
            self._params[name] = True

    def inc_grad(self, name: str, value: Array) -> None:
        """Check whether the model has a gradient of the given name."""
        grad_name = f"d_{name}"
        key = (self.id, grad_name)
        param_key = (self.id, name)
        if key in self._mem:
            grad = self._mem[key]
        else:
            grad = self._mem.add_gradient(key, param_key)

        if grad.shape != value.shape:
            raise ValueError(
                f"Cannot add a value to the gradient of parameter '{name}' for model '{self.name}' "
                f"as the shapes should match, but found {grad.shape} for the original gradient, and "
                f"{value.shape} for the value that should be added."
            )

        grad += value
        self._grads[grad_name] = True

    def has_grad(self, name: str) -> Optional[bool]:
        """Check whether the model has a non-zero gradient for a parameter.
        Returns None if the gradient is allocated but currently 0.
        """
        grad_name = f"d_{name}"
        key = (self.id, grad_name)
        if key not in self._mem:
            return False
        elif not self._mem[key].any():
            return None
        else:
            return True

    def get_grad(self, name: str) -> Array:
        """Get a gradient from the model."""
        grad_name = f"d_{name}"
        key = (self.id, grad_name)
        if key not in self._mem:
            err = f"Gradient '{grad_name}' has not been allocated yet for model '{self.name}'"
            raise KeyError(err)
        return self._mem[key]

    def set_grad(self, name: str, value: Array) -> None:
        """Set a gradient value for the model."""
        grad_name = f"d_{name}"
        key = (self.id, grad_name)
        if key not in self._mem:
            self.inc_grad(name, value)
        else:
            data = self._mem[key]
            try:
                copy_array(dst=data, src=value)
            except ValueError as e:
                err = f"Cannot set grad '{grad_name}' for model '{self.name}': {e}"
                raise ValueError(err)

    def has_attr(self, name: str) -> bool:
        """Check whether the model has the given attribute."""
        return name in self._attrs

    def get_attr(self, name: str) -> Any:
        """Get the attribute. Raises KeyError if not present."""
        if name not in self._attrs:
            raise KeyError(f"Cannot get attribute '{name}' for model '{self.name}'.")
        return self._attrs[name]

    def set_attr(self, name: str, value: Any) -> None:
        """Set the attribute to the given value."""
        self._attrs[name] = value

    def has_ref(self, name: str) -> Optional[bool]:
        """Check whether the model has a reference of a given name. If the
        reference is registered but the value is unset, returns None.
        """
        if name not in self._refs:
            return False
        elif self._refs[name] is not None:
            return True
        else:
            return None

    def get_ref(self, name: str) -> "Model":
        """Retrieve the value of a reference of the given name."""
        if name not in self._refs:
            raise KeyError(f"Cannot get reference '{name} for model '{self.name}'.")
        value = self._refs[name]
        if value is None:
            err = f"Cannot get reference '{name}' for model '{self.name}': value unset."
            raise ValueError(err)
        else:
            return value

    def set_ref(self, name: str, value: Optional["Model"]) -> None:
        """Set a value for a reference."""
        if value is None:
            self._refs[name] = value
        elif value in self.walk():
            self._refs[name] = value
        else:
            raise ValueError("Cannot add reference to node not in tree.")

    def __call__(self, X: InT, is_train: bool = False) -> Tuple[OutT, Callable]:
        """Call the model's `forward` function, returning the output and a
        callback to compute the gradients via backpropagation."""
        return self._func(self, X, is_train=is_train)

    def initialize(self, X: Optional[InT] = None, Y: Optional[OutT] = None) -> "Model":
        """Finish initialization of the model, optionally providing a batch of
        example input and output data to perform shape inference."""
        if self._init is not None:
            self._init(self, X=X, Y=Y)
        return self

    def begin_update(self, X: InT) -> Tuple[OutT, Callable[[InT], OutT]]:
        """Run the model over a batch of data, returning the output and a
        callback to complete the backward pass. A tuple (Y, finish_update),
        where Y is a batch of output data, and finish_update is a callback that
        takes the gradient with respect to the output and an optimizer function,
        and returns the gradient with respect to the input.
        """
        return self._func(self, X, is_train=True)

    def predict(self, X: InT) -> OutT:
        """Call the model's `forward` function with `is_train=False`, and return
        only the output, instead of the `(output, callback)` tuple.
        """
        return self._func(self, X, is_train=False)[0]

    def finish_update(self, optimizer: Optimizer) -> None:
        """Update parameters with current gradients. The optimizer is called
        with each parameter and gradient of the model.<commit_msg>Fix model signature<commit_after>from typing import Dict, List, Callable, Optional, Any, Union, Iterable, Set
from typing import Generic, Sequence, Tuple, TypeVar
import numpy
import contextlib
import srsly
from pathlib import Path
import copy

from .backends import NumpyOps, CupyOps, get_current_ops
from .optimizers import Optimizer  # noqa: F401
from .backends.mem import Memory
from .shims import Shim
from .util import copy_array, get_width, create_thread_local
from .types import Array


InT = TypeVar("InT")
OutT = TypeVar("OutT")


def create_init(initializers: Dict[str, Callable]) -> Callable:
    """Create an init function, given a dictionary of parameter initializers."""

    def init(
        model: Model, X: Optional[Array] = None, Y: Optional[Array] = None
    ) -> None:
        if X is not None:
            model.set_dim("nI", get_width(X))
        if Y is not None:
            model.set_dim("nO", get_width(Y))
        W = model.ops.alloc_f2d(model.get_dim("nO"), model.get_dim("nI"))
        b = model.ops.alloc_f1d(model.get_dim("nO"))
        if "W" in initializers:
            initializers["W"](W, inplace=True)
        if "b" in initializers:
            initializers["b"](b, inplace=True)
        model.set_param("W", W)
        model.set_param("b", b)

    return init


class Model(Generic[InT, OutT]):
    """Class for implementing Thinc models and layers."""

    global_id: int = 0
    _thread_local = create_thread_local({"operators": {}})

    name: str
    ops: Union[NumpyOps, CupyOps]
    id: int
    _func: Callable
    _init: Callable
    _mem: Memory
    _params: Dict[str, Optional[bool]]
    _dims: Dict[str, Optional[int]]
    _grads: Dict[str, Optional[bool]]
    _layers: List["Model"]
    _shims: List[Shim]
    _attrs: Dict[str, Any]

    # This "locks" the class, so we get an error if you try to assign to
    # an unexpected variable.
    __slots__ = [
        "name",
        "id",
        "ops",
        "_func",
        "_init",
        "_mem",
        "_params",
        "_dims",
        "_grads",
        "_attrs",
        "_refs",
        "_layers",
        "_shims",
    ]

    def __init__(
        self,
        name: str,
        forward: Callable,
        *,
        init: Callable = lambda *a, **k: None,
        dims: Dict[str, Optional[int]] = {},
        params: Dict[str, Optional[Array]] = {},
        grads: Dict[str, Optional[Array]] = {},
        layers: Sequence["Model"] = [],
        shims: List[Shim] = [],
        attrs: Dict[str, Any] = {},
        refs: Dict[str, Optional["Model"]] = {},
        ops: Optional[Union[NumpyOps, CupyOps]] = None,
    ):
        """Initialize a new model."""
        self.name = name
        # Assign to callable attrs: https://github.com/python/mypy/issues/2427
        setattr(self, "_func", forward)
        setattr(self, "_init", init)
        self.ops = ops if ops is not None else get_current_ops()
        self._mem = Memory(self.ops)
        self._dims = dict(dims)
        self._attrs = dict(attrs)
        self._refs = dict(refs)
        self._layers = list(layers)
        self._shims = list(shims)
        self.__class__.global_id += 1
        self.id = self.__class__.global_id
        self._params = {}
        self._grads = {}
        for name, value in params.items():
            self._params[name] = None
            if value is not None:
                self.set_param(name, value)
        for name, value in grads.items():
            self._grads[name] = None
            if value is not None:
                self.set_grad(name, value)

    @property
    def layers(self) -> List["Model"]:
        """A list of child layers of the model. You can append to it to add
        layers but not reassign it.
        """
        return self._layers

    @property
    def shims(self) -> List[Shim]:
        return self._shims

    @property
    def param_names(self) -> Tuple[str, ...]:
        """Get the names of registered parameter (including unset)."""
        return tuple(self._params.keys())

    @property
    def grad_names(self) -> Tuple[str, ...]:
        """Get the names of parameters with registered gradients (including unset)."""
        return tuple([name for name in self.param_names if self.has_grad(name)])

    @property
    def dim_names(self) -> Tuple[str, ...]:
        """Get the names of registered dimensions (including unset)."""
        return tuple(self._dims.keys())

    @property
    def attr_names(self) -> Tuple[str, ...]:
        """Get the names of attributes."""
        return tuple(self._attrs.keys())

    @property
    def ref_names(self) -> Tuple[str, ...]:
        """Get the names of registered node references (including unset)."""
        return tuple(self._refs.keys())

    @classmethod
    @contextlib.contextmanager
    def define_operators(cls, operators: Dict[str, Callable]):
        """Bind arbitrary binary functions to Python operators, for use in any
        `Model` instance. Can (and should) be used as a contextmanager.

        EXAMPLE:
            with Model.define_operators({">>": chain}):
                model = ReLu(512) >> ReLu(512) >> Softmax()
        """
        curr_operators = dict(cls._thread_local.operators)
        cls._thread_local.operators = dict(operators)
        yield
        cls._thread_local.operators = dict(curr_operators)

    def has_dim(self, name: str) -> Optional[bool]:
        """Check whether the model has a dimension of a given name. If the
        dimension is registered but the value is unset, returns None.
        """
        if name not in self._dims:
            return False
        elif self._dims[name] is not None:
            return True
        else:
            return None

    def get_dim(self, name: str) -> int:
        """Retrieve the value of a dimension of the given name."""
        if name not in self._dims:
            raise KeyError(f"Cannot get dimension '{name}' for model '{self.name}'")
        value = self._dims[name]
        if value is None:
            err = f"Cannot get dimension '{name}' for model '{self.name}': value unset"
            raise ValueError(err)
        else:
            return value

    def set_dim(self, name: str, value: int) -> None:
        """Set a value for a dimension."""
        if name not in self._dims:
            raise KeyError(f"Cannot set dimension '{name}' for model '{self.name}'.")
        old_value = self._dims[name]
        if old_value is not None and old_value != value:
            err = f"Attempt to change dimension '{name}' for model '{self.name}' from {old_value} to {value}"
            raise ValueError(err)
        self._dims[name] = value

    def has_param(self, name: str) -> Optional[bool]:
        """Check whether the model has a weights parameter of the given name.

        Returns None if the parameter is registered but currently unset.
        """
        if name not in self._params:
            return False
        elif self._params[name] is not None:
            return True
        else:
            return None

    def get_param(self, name: str) -> Array:
        """Retrieve a weights parameter by name."""
        if name not in self._params:
            raise KeyError(f"Unknown param: '{name}' for model '{self.name}'.")
        key = (self.id, name)
        if key not in self._mem:
            raise KeyError(
                f"Parameter '{name}' for model '{self.name}' has not been allocated yet."
            )
        return self._mem[key]

    def set_param(self, name: str, value: Optional[Array]) -> None:
        """Set a weights parameter's value."""
        if value is None:
            self._params[name] = None
        else:
            key = (self.id, name)
            if key not in self._mem:
                self._mem.add(key, value.shape)
            data = self._mem[(self.id, name)]
            try:
                copy_array(dst=data, src=value)
            except ValueError as e:
                raise ValueError(
                    f"Cannot set param '{name}' for model '{self.name}': {e}."
                )
            self._params[name] = True

    def inc_grad(self, name: str, value: Array) -> None:
        """Check whether the model has a gradient of the given name."""
        grad_name = f"d_{name}"
        key = (self.id, grad_name)
        param_key = (self.id, name)
        if key in self._mem:
            grad = self._mem[key]
        else:
            grad = self._mem.add_gradient(key, param_key)

        if grad.shape != value.shape:
            raise ValueError(
                f"Cannot add a value to the gradient of parameter '{name}' for model '{self.name}' "
                f"as the shapes should match, but found {grad.shape} for the original gradient, and "
                f"{value.shape} for the value that should be added."
            )

        grad += value
        self._grads[grad_name] = True

    def has_grad(self, name: str) -> Optional[bool]:
        """Check whether the model has a non-zero gradient for a parameter.
        Returns None if the gradient is allocated but currently 0.
        """
        grad_name = f"d_{name}"
        key = (self.id, grad_name)
        if key not in self._mem:
            return False
        elif not self._mem[key].any():
            return None
        else:
            return True

    def get_grad(self, name: str) -> Array:
        """Get a gradient from the model."""
        grad_name = f"d_{name}"
        key = (self.id, grad_name)
        if key not in self._mem:
            err = f"Gradient '{grad_name}' has not been allocated yet for model '{self.name}'"
            raise KeyError(err)
        return self._mem[key]

    def set_grad(self, name: str, value: Array) -> None:
        """Set a gradient value for the model."""
        grad_name = f"d_{name}"
        key = (self.id, grad_name)
        if key not in self._mem:
            self.inc_grad(name, value)
        else:
            data = self._mem[key]
            try:
                copy_array(dst=data, src=value)
            except ValueError as e:
                err = f"Cannot set grad '{grad_name}' for model '{self.name}': {e}"
                raise ValueError(err)

    def has_attr(self, name: str) -> bool:
        """Check whether the model has the given attribute."""
        return name in self._attrs

    def get_attr(self, name: str) -> Any:
        """Get the attribute. Raises KeyError if not present."""
        if name not in self._attrs:
            raise KeyError(f"Cannot get attribute '{name}' for model '{self.name}'.")
        return self._attrs[name]

    def set_attr(self, name: str, value: Any) -> None:
        """Set the attribute to the given value."""
        self._attrs[name] = value

    def has_ref(self, name: str) -> Optional[bool]:
        """Check whether the model has a reference of a given name. If the
        reference is registered but the value is unset, returns None.
        """
        if name not in self._refs:
            return False
        elif self._refs[name] is not None:
            return True
        else:
            return None

    def get_ref(self, name: str) -> "Model":
        """Retrieve the value of a reference of the given name."""
        if name not in self._refs:
            raise KeyError(f"Cannot get reference '{name} for model '{self.name}'.")
        value = self._refs[name]
        if value is None:
            err = f"Cannot get reference '{name}' for model '{self.name}': value unset."
            raise ValueError(err)
        else:
            return value

    def set_ref(self, name: str, value: Optional["Model"]) -> None:
        """Set a value for a reference."""
        if value is None:
            self._refs[name] = value
        elif value in self.walk():
            self._refs[name] = value
        else:
            raise ValueError("Cannot add reference to node not in tree.")

    def __call__(self, X: InT, is_train: bool) -> Tuple[OutT, Callable]:
        """Call the model's `forward` function, returning the output and a
        callback to compute the gradients via backpropagation."""
        return self._func(self, X, is_train=is_train)

    def initialize(self, X: Optional[InT] = None, Y: Optional[OutT] = None) -> "Model":
        """Finish initialization of the model, optionally providing a batch of
        example input and output data to perform shape inference."""
        if self._init is not None:
            self._init(self, X=X, Y=Y)
        return self

    def begin_update(self, X: InT) -> Tuple[OutT, Callable[[InT], OutT]]:
        """Run the model over a batch of data, returning the output and a
        callback to complete the backward pass. A tuple (Y, finish_update),
        where Y is a batch of output data, and finish_update is a callback that
        takes the gradient with respect to the output and an optimizer function,
        and returns the gradient with respect to the input.
        """
        return self._func(self, X, is_train=True)

    def predict(self, X: InT) -> OutT:
        """Call the model's `forward` function with `is_train=False`, and return
        only the output, instead of the `(output, callback)` tuple.
        """
        return self._func(self, X, is_train=False)[0]

    def finish_update(self, optimizer: Optimizer) -> None:
        """Update parameters with current gradients. The optimizer is called
        with each parameter and gradient of the model.
{'fraction_non_alphanumeric': 0.11325310155263353, 'fraction_numerical': 0.0014114850308298047, 'mean_word_length': 2.1042315231177215, 'pattern_counts': {'":': 8, '<': 3, '<?xml version=': 0, '>': 75, 'https://': 2, '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': '2432065', 'n_tokens_mistral': 7750, 'n_tokens_neox': 7617, 'n_words': 2666}