orcalib.torch#

OrcaLookupCacheBuilder #

OrcaLookupCacheBuilder(
    db,
    index_name,
    num_memories,
    embedding_col_name,
    memory_column_aliases,
    drop_exact_match=False,
    exact_match_threshold=EXACT_MATCH_THRESHOLD,
    batch_size=1000,
)

This class allows you to precache your lookup results into a HuggingFace Dataset for faster training. It also allows you to inject the pre-cached lookup results into your model during training and inference.

In the future, we may extend this class to support additional dataset types and frameworks.

Examples:

First, configure the lookup cache builder with the necessary information about your lookups.

lookup_cache_builder = OrcaLookupCacheBuilder(
    db=OrcaDatabase(DB_NAME),
    index_name=INDEX_NAME,
    num_memories=10,
    embedding_col_name="embeddings",
    memory_column_aliases={"$score": "memory_scores", "label": "memory_labels"},
)

train_data = load_dataset(DATASET_NAME)

Next, add the lookup results to the HuggingFace Dataset. lookup_cache_builder.add_lookups_to_hf_dataset(train_data, "text")

Finally, inject the lookup results into your model during training and inference.

class MyModel(OrcaLookupModule):
    def __init__(self, cache_builder: OrcaLookupCacheBuilder):
        ...
        # Orca modules that use lookup layers don't need to know about the lookup cache builder.
        self.my_orca_layer = OrcaLookupLayer(...)
        self.cache_builder = cache_builder

    def forward(self, x, memory_scores, memory_labels):
        # Finally, inject the lookup results into the model. Downstream lookup layers will use these results.
        self.cache_builder.inject_lookup_results(self, memory_scores=memory_scores, memory_labels=memory_labels)
        ...

Parameters:

db (OrcaDatabase) –

The OrcaDatabase instance that contains the memory table we’ll use for lookups.
index_name (str) –

The name of the memory index we’ll use for lookups.
num_memories (int) –

The number of memories to fetch for each element in the data set.
embedding_col_name (str | None) –

The name of the feature that will be created to store the embedding of the query column. If None, the embedding will not be stored (for when embeddings are alredy processed).
memory_column_aliases (Dict[str, str]) –

Maps the lookup column names to the feature name that will be added to the Dataset. For example, {"$score": "memory_scores", "label": "memory_labels"} will add a memory_scores column and a memory_labels column to the Dataset that contains the scores and labels of the memories, respectively. It’s a good idea to align the aliases to match the inputs to your model’s forward() method, e.g., forward(x, memory_scores, memory_labels).
drop_exact_match (bool, default: False ) –

Whether to drop the highest match that’s above the exact_match_threshold.
exact_match_threshold (float, default: EXACT_MATCH_THRESHOLD ) –

The similarity threshold for considering a memory an exact match.
batch_size (int, default: 1000 ) –

The batch size to use for the vector scan. Defaults to 1000.

memory_column_aliases `instance-attribute` #

memory_column_aliases = memory_column_aliases

A mapping of the lookup column names to the feature name that will be added to the dataset

from_model `classmethod` #

from_model(
    model,
    embedding_col_name,
    memory_column_aliases=None,
    memory_feature_prefix="memory_",
)

This is a convenient way to create an OrcaLookupCacheBuilder that is configured with the same lookup settings as a model or module. It will create memory column aliases for each lookup column name in the model’s settings by prepending the memory_feature_prefix to the column name. For example, if the model has a lookup column named label, the memory column alias will be memory_label, because memory_feature_prefix defaults to "memory_".

Parameters:

model (OrcaLookupModule) –

The model whose settings will be used to create the OrcaLookupCacheBuilder.
embedding_col_name (str) –

The name of the feature that will be created to store the embedding of the query column.
memory_column_aliases (Dict[str, str], default: None ) –

A mapping of the lookup column names to the feature name that will be added to the Dataset. This can help when there are conflicts between the lookup column names and special columns. For example, if you’re looking up both “score” and “$score", you can provide an alias for "$score”, which would both map to “memory_score”. This will help avoid conflicts. Defaults to None.
memory_feature_prefix (str, default: 'memory_' ) –

The prefix to prepend to create the names of the features that will be added to the Dataset to hold the lookup results. Defaults to "memory_". For example, if the model has a lookup column named label, the memory column alias will be memory_label. For special columns, e.g., $score, the $ will be removed: memory_score.

Returns:

OrcaLookupCacheBuilder –

An OrcaLookupCacheBuilder instance that is configured with the same settings as the model.

inject_lookup_results #

inject_lookup_results(model, **features)

Sets (or clears) the lookup result override for an OrcaLookupModule, e.g., OrcaLookupLayer, OrcaModel. All downstream lookup layers will use these results instead of performing a lookup by contacting the database. When the feature values are None, the lookup result override will be cleared instead.

Parameters:

model (OrcaLookupModule) –

The OrcaLookupModule to inject the lookup results into.
features (Dict[str, list[list[Any]] | Tensor], default: {} ) –

A mapping of the memory column aliases to their values. These values will be converted to the correct type to be used as the lookup results. Important: The feature values should all be None or all be non-None. If the feature values are None, the lookup result override will be cleared.

Example

class MyModel(OrcaLookupModule):
    ...

    def forward(self, x, memory_scores, memory_labels):
        self.cache_builder.inject_lookup_results(self, memory_scores=memory_scores, memory_labels=memory_labels)
        ...

get_lookup_result #

get_lookup_result(**features)

Returns a BatchedScanResult that contains lookup results that contain the provided features.

Parameters:

features (Dict[str, list[list[Any]] | Tensor], default: {} ) –

A mapping of the memory column aliases to their values. These values will be converted based on the column/artifact type to be used as the lookup results.

Returns:

BatchedScanResult –

The lookup results that contain the provided features for each memory.

add_lookups_to_hf_dataset #

add_lookups_to_hf_dataset(
    ds, query_column_name, map_cache_file_name=None
)

Adds the lookup results as columns (i.e., features) to a HuggingFace Dataset. This function will perform a vector scan on the memory index to fetch the lookup results for each example in the dataset. The feature names for the memories will be the same as the memory_column_aliases provided in the constructor. The embedding of the query column will be stored in the embedding_col_name provided

Parameters:

ds (Dataset) –

The HuggingFace dataset to add the lookup results to.
query_column_name (str) –

The name of the column that contains the query text to lookup in the memory index.
map_cache_file_name (str | None, default: None ) –

The name of the cache file to use for the mapping of the dataset file. Defaults to None.

Returns:

Dataset –

The HuggingFace dataset with the lookup results added as features.

Examples:

First, configure the lookup cache builder with the necessary information about your lookups.

lookup_cache_builder = OrcaLookupCacheBuilder(
    db=OrcaDatabase(DB_NAME),
    index_name=INDEX_NAME,
    num_memories=10,
    embedding_col_name="embeddings",
    memory_column_aliases={"$score": "memory_scores", "label": "memory_labels"},
)

Now, load the HuggingFace dataset and add the lookup results to it.

train_data = load_dataset(DATASET_NAME) # Load the HuggingFace dataset
lookup_cache_builder.add_lookups_to_hf_dataset(train_data, "text")

OrcaCrossAttentionClassificationHead #

OrcaCrossAttentionClassificationHead(
    model_dim,
    num_classes,
    num_memories,
    num_layers=1,
    num_heads=8,
    classification_mode=ClassificationMode.DIRECT,
    activation=F.relu,
    dropout=0.1,
    deep_residuals=True,
    split_retrieval_path=False,
    memory_guide_weight=0.0,
    single_lookup=True,
    database=None,
    curate_enabled=False,
    memory_index_name=None,
    drop_exact_match=None,
    exact_match_threshold=None,
    shuffle_memories=False,
    label_column_name=None,
)

Bases: OrcaLookupModule, LabelColumnNameMixin

A transformer decoder layer block that does cross attention with memory lookup

Examples:

import torch
from orcalib.orca_torch import OrcaModule, OrcaClassificationHead

class MyModule(OrcaModule):
    def __init__(self):
        super().__init__()
        self.trunk = torch.nn.Linear(10, 10)
        self.classifier = OrcaClassificationHead(model_dim=10, num_classes=5, "my_index", "my_label", num_memories=10)

    def forward(self, x):
        x = self.trunk(x) # N x 10
        x = self.classifier(x)
        return x # N x 5, e.g., where each row may become logits for a softmax

Parameters:

model_dim (int) –

The dimension of the input vector and hidden layers.
num_classes (int) –

The size of the output vector.
num_memories (int) –

The number of memory vectors to be returned from the lookup.
num_layers (int, default: 1 ) –

The number of attention blocks to be used, copies of OrcaClassificationCrossAttentionLayer.
num_heads (int, default: 8 ) –

The number of heads to be used in the multi-head attention layer.
classification_mode (ClassificationMode, default: DIRECT ) –

The mode of classification to be used.
activation (Callable[[Tensor], Tensor], default: relu ) –

The activation function.
dropout (float, default: 0.1 ) –

The dropout rate.
deep_residuals (bool, default: True ) –

Whether to use deep residuals.
split_retrieval_path (bool, default: False ) –

Whether to split the retrieval path.
memory_guide_weight (float, default: 0.0 ) –

The weight of the memory guide.
single_lookup (bool, default: True ) –

Whether to use a single lookup.
database (OrcaDatabase | str | None, default: None ) –

The OrcaDatabase instance to use for lookups and curate tracking.
curate_enabled (bool, default: False ) –

Whether Curate tracking is enabled.
memory_index_name (str | None, default: None ) –

The name of the index to use for lookups.
drop_exact_match (DropExactMatchOption | None, default: None ) –

Choose to drop the exact match (if found) always, never, or only during training or inference.
exact_match_threshold (float | None, default: None ) –

Minimum similarity score for something to be considered the exact match
shuffle_memories (bool, default: False ) –

Whether to shuffle the memories before returning them.
label_column_name (ColumnName | None, default: None ) –

The name of the label column to return from the index.

lookup_result_transforms `property` `writable` #

1	`lookup_result_transforms`

A list of transforms to apply to the lookup result. NOTE: This will be applied even when lookup_result_override is set.

extra_lookup_column_names `property` `writable` #

1	`extra_lookup_column_names`

While set, all lookups will include these additional columns. They may inclue columns on the indexed table as well as index-specific columns, e.g., $score, $embedding.

lookup_query_override `property` `writable` #

1	`lookup_query_override`

The query to use instead of performing a lookup. NOTE: This will be ignored if lookup_result_override is also set.

lookup_result_override `property` `writable` #

1	`lookup_result_override`

The lookup result to use instead of performing a lookup.

lookup_database `property` `writable` #

1	`lookup_database`

The name of the database to use for looking up memories.

memory_index_name `property` `writable` #

1	`memory_index_name`

The name of the index to use for looking up memories.

lookup_column_names `property` `writable` #

1	`lookup_column_names`

The names of the columns to retrieve for each memory.

num_memories `property` `writable` #

1	`num_memories`

The number of memories to look up.

drop_exact_match `property` `writable` #

1	`drop_exact_match`

Whether to drop exact matches from the results.

exact_match_threshold `property` `writable` #

1	`exact_match_threshold`

The similarity threshold for exact matches.

shuffle_memories `property` `writable` #

1	`shuffle_memories`

Whether to shuffle the looked up memories.

curate_database `property` `writable` #

1	`curate_database`

The name of the database to use for saving curate tracking data.

curate_next_run_settings `property` `writable` #

1	`curate_next_run_settings`

The settings for the next curate model run.

curate_model_id `property` `writable` #

1	`curate_model_id`

The model id to associate with curated model runs.

curate_model_version `property` `writable` #

1	`curate_model_version`

The model version to associate with curated model runs.

curate_metadata `property` `writable` #

1	`curate_metadata`

The metadata to attach to curated model runs.

curate_tags `property` `writable` #

1	`curate_tags`

The tags to attach to the curated model runs.

curate_seq_id `property` `writable` #

1	`curate_seq_id`

The sequence id to associate with curated model runs.

curate_batch_size `property` `writable` #

1	`curate_batch_size`

The batch size of the model run to track curate data for, usually inferred automatically.

last_curate_run_ids `property` `writable` #

1	`last_curate_run_ids`

The run ids of the last model run for which curate tracking data was collected.

last_curate_run_settings `property` `writable` #

1	`last_curate_run_settings`

The settings of the last model run for which curate tracking data was collected.

get_effective_lookup_settings #

1	`get_effective_lookup_settings()`

Returns the effective lookup settings for this module, with any inherited settings applied.

Returns:

LookupSettings –

The effective lookup settings for this module. Practically, this be the lookup settings
LookupSettings –

set on this module. For any settings that are not set on this module, the inherited settings
LookupSettings –

will be used instead.

get_lookup_database_instance #

1	`get_lookup_database_instance()`

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

get_orca_modules_recursively(
    max_depth=None, include_self=True, filter_type=None
)

Recursively yields all children of this module that are instances of the specified filter type.

All parent nodes will be processed before their children
This will search through all children —even those that are not a subclass of the filter_type— but it only returns children that are a subclass of filter_type.

Parameters:

max_depth (int | None, default: None ) –
The maximum depth to search.
- Setting this to 0 will only include this module.
- Setting this to 1 will include only this module and its children.
- Setting it to None (the default) will search through all modules.
- Modules that are not of filter_type or OrcaModule do not increment the depth.
include_self (bool, default: True ) –

Whether to include the current OrcaModule in the results.
filter_type (type[ORCA_MODULE_TYPE] | None, default: None ) –

The subtype of OrcaModule to filter for. If None, any subtypes of OrcaModule will be returned.

Yields:

ORCA_MODULE_TYPE –

modules of type filter_type that are used in the children of this module.

enable_curate #

enable_curate(recursive=True)

Enable Curate tracking for the model and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to enable Curate tracking recursively.

disable_curate #

disable_curate(recursive=True)

Disable Curate tracking for this module and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to disable Curate tracking recursively.

update_curate_settings #

update_curate_settings(
    model_id=None,
    model_version=None,
    tags=None,
    extra_tags=None,
    metadata=None,
    extra_metadata=None,
    batch_size=None,
    seq_id=None,
    enabled=None,
    enable_recursive=True,
)

Update curate tracking settings for the module and all its children.

Parameters:

model_id (str | None, default: None ) –

The ID of the model.
model_version (str | None, default: None ) –

The version of the model.
tags (Iterable[str] | None, default: None ) –

The new tags to be added to the model.
extra_tags (Iterable[str] | None, default: None ) –

The extra tags to be added to the model.
metadata (OrcaMetadataDict | None, default: None ) –

The new metadata to be added to the model.
extra_metadata (OrcaMetadataDict | None, default: None ) –

The extra metadata to be added to the model.
batch_size (int | None, default: None ) –

The batch size to be used for the model.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the model.

record_next_model_memory_lookups #

record_next_model_memory_lookups(
    tags=None, metadata=None, batch_size=None, seq_id=None
)

Sets up curate tracking for the memory lookups during the next forward pass only.

Parameters:

tags (Iterable[str] | None, default: None ) –

Additional tags to be recorded on the next model run.
metadata (OrcaMetadataDict | None, default: None ) –

Additional metadata to be recorded on the next model run.
batch_size (int | None, default: None ) –

The batch size to be used for the next model run.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the next model run.

record_model_feedback #

record_model_feedback(
    val, name="default", kind=FeedbackKind.CONTINUOUS
)

Records feedback for the last model runs for which memory lookups were recorded by curate.

Parameters:

val (list[float] | float | int | list[int]) –

The feedback to be recorded.
name (str, default: 'default' ) –

The name of the feedback.
kind (FeedbackKind, default: CONTINUOUS ) –

The kind of feedback.

record_model_input_output #

record_model_input_output(inputs, outputs)

Records the inputs and outputs of the last model runs for which memory lookups were recorded by curate.

Parameters:

inputs (list[Any] | Any) –

The inputs to be recorded.
outputs (list[Any] | Any) –

The outputs to be recorded.

get_lookup_setting_summary #

1	`get_lookup_setting_summary()`

Returns a summary of the lookup settings for each OrcaLookupLayer in this module and its descendants.

forward #

forward(x, ctx=None, labels=None, memory_key=None)

Generate logits based on the input vector and memory context.

Parameters:

x (Tensor) –

The input tensor of shape (batch_size, embedding_dim)
ctx (Tensor | None, default: None ) –

The memory context tensor of shape (batch_size, num_memories, embedding_dim). If None, the memory context will be looked up based on the memory_key or input tensor.
labels (Tensor | None, default: None ) –

The memory label tensor of shape (batch_size, num_memories) containing integer labels. If None, the labels will be looked up from the index based on the memory_key or input tensor.
memory_key (Tensor | None, default: None ) –

The memory key tensor of shape (batch_size, embedding_dim) to use for lookup. If None, the input tensor will be used.

Returns:

Tensor –

The logits tensor of shape (batch_size, num_classes)

OrcaKnnClassifier #

OrcaKnnClassifier(
    num_classes,
    num_memories,
    label_column_name,
    weigh_memories=True,
    database=None,
    curate_enabled=False,
    memory_index_name=None,
    drop_exact_match=DropExactMatchOption.TRAINING_ONLY,
    exact_match_threshold=None,
)

Bases: OrcaLookupModule

A simple KNN layer that returns the average label of the K nearest memories to the input vector.

Examples:

import torch
from orcalib import OrcaModule, OrcaKnnClassifier

class MyModule(OrcaModule):
    def __init__(self):
        super().__init__()
        self.knn_head = OrcaKnnClassifier(
            memory_index_name="my_index",
            label_column_name="my_label",
            num_memories=10,
            num_classes=5,
        )

    def forward(self, x):
        logits = self.knn_head(x)
        return logits

Parameters:

num_classes (int) –

The size of the output vector.
num_memories (int) –

The number of memory vectors to be returned from the lookup.
weigh_memories (bool, default: True ) –

Whether to weigh the memories by their scores.
database (OrcaDatabase | str | None, default: None ) –

The OrcaDatabase instance to use for lookups and curate tracking.
curate_enabled (bool, default: False ) –

Whether Curate tracking is enabled.
memory_index_name (str | None, default: None ) –

The name of the index to use for lookups.
drop_exact_match (DropExactMatchOption, default: TRAINING_ONLY ) –

Choose to drop the exact match (if found) always, never, or only during training or inference.
exact_match_threshold (float | None, default: None ) –

Minimum similarity score for something to be considered the exact match
label_column_name (ColumnName) –

The name of the label column to return from the index.

lookup_result_transforms `property` `writable` #

1	`lookup_result_transforms`

A list of transforms to apply to the lookup result. NOTE: This will be applied even when lookup_result_override is set.

extra_lookup_column_names `property` `writable` #

1	`extra_lookup_column_names`

While set, all lookups will include these additional columns. They may inclue columns on the indexed table as well as index-specific columns, e.g., $score, $embedding.

lookup_query_override `property` `writable` #

1	`lookup_query_override`

The query to use instead of performing a lookup. NOTE: This will be ignored if lookup_result_override is also set.

lookup_result_override `property` `writable` #

1	`lookup_result_override`

The lookup result to use instead of performing a lookup.

lookup_database `property` `writable` #

1	`lookup_database`

The name of the database to use for looking up memories.

memory_index_name `property` `writable` #

1	`memory_index_name`

The name of the index to use for looking up memories.

lookup_column_names `property` `writable` #

1	`lookup_column_names`

The names of the columns to retrieve for each memory.

num_memories `property` `writable` #

1	`num_memories`

The number of memories to look up.

drop_exact_match `property` `writable` #

1	`drop_exact_match`

Whether to drop exact matches from the results.

exact_match_threshold `property` `writable` #

1	`exact_match_threshold`

The similarity threshold for exact matches.

shuffle_memories `property` `writable` #

1	`shuffle_memories`

Whether to shuffle the looked up memories.

curate_database `property` `writable` #

1	`curate_database`

The name of the database to use for saving curate tracking data.

curate_next_run_settings `property` `writable` #

1	`curate_next_run_settings`

The settings for the next curate model run.

curate_model_id `property` `writable` #

1	`curate_model_id`

The model id to associate with curated model runs.

curate_model_version `property` `writable` #

1	`curate_model_version`

The model version to associate with curated model runs.

curate_metadata `property` `writable` #

1	`curate_metadata`

The metadata to attach to curated model runs.

curate_tags `property` `writable` #

1	`curate_tags`

The tags to attach to the curated model runs.

curate_seq_id `property` `writable` #

1	`curate_seq_id`

The sequence id to associate with curated model runs.

curate_batch_size `property` `writable` #

1	`curate_batch_size`

The batch size of the model run to track curate data for, usually inferred automatically.

last_curate_run_ids `property` `writable` #

1	`last_curate_run_ids`

The run ids of the last model run for which curate tracking data was collected.

last_curate_run_settings `property` `writable` #

1	`last_curate_run_settings`

The settings of the last model run for which curate tracking data was collected.

get_effective_lookup_settings #

1	`get_effective_lookup_settings()`

Returns the effective lookup settings for this module, with any inherited settings applied.

Returns:

LookupSettings –

The effective lookup settings for this module. Practically, this be the lookup settings
LookupSettings –

set on this module. For any settings that are not set on this module, the inherited settings
LookupSettings –

will be used instead.

get_lookup_database_instance #

1	`get_lookup_database_instance()`

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

get_orca_modules_recursively(
    max_depth=None, include_self=True, filter_type=None
)

Recursively yields all children of this module that are instances of the specified filter type.

All parent nodes will be processed before their children
This will search through all children —even those that are not a subclass of the filter_type— but it only returns children that are a subclass of filter_type.

Parameters:

max_depth (int | None, default: None ) –
The maximum depth to search.
- Setting this to 0 will only include this module.
- Setting this to 1 will include only this module and its children.
- Setting it to None (the default) will search through all modules.
- Modules that are not of filter_type or OrcaModule do not increment the depth.
include_self (bool, default: True ) –

Whether to include the current OrcaModule in the results.
filter_type (type[ORCA_MODULE_TYPE] | None, default: None ) –

The subtype of OrcaModule to filter for. If None, any subtypes of OrcaModule will be returned.

Yields:

ORCA_MODULE_TYPE –

modules of type filter_type that are used in the children of this module.

enable_curate #

enable_curate(recursive=True)

Enable Curate tracking for the model and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to enable Curate tracking recursively.

disable_curate #

disable_curate(recursive=True)

Disable Curate tracking for this module and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to disable Curate tracking recursively.

update_curate_settings #

update_curate_settings(
    model_id=None,
    model_version=None,
    tags=None,
    extra_tags=None,
    metadata=None,
    extra_metadata=None,
    batch_size=None,
    seq_id=None,
    enabled=None,
    enable_recursive=True,
)

Update curate tracking settings for the module and all its children.

Parameters:

model_id (str | None, default: None ) –

The ID of the model.
model_version (str | None, default: None ) –

The version of the model.
tags (Iterable[str] | None, default: None ) –

The new tags to be added to the model.
extra_tags (Iterable[str] | None, default: None ) –

The extra tags to be added to the model.
metadata (OrcaMetadataDict | None, default: None ) –

The new metadata to be added to the model.
extra_metadata (OrcaMetadataDict | None, default: None ) –

The extra metadata to be added to the model.
batch_size (int | None, default: None ) –

The batch size to be used for the model.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the model.

record_next_model_memory_lookups #

record_next_model_memory_lookups(
    tags=None, metadata=None, batch_size=None, seq_id=None
)

Sets up curate tracking for the memory lookups during the next forward pass only.

Parameters:

tags (Iterable[str] | None, default: None ) –

Additional tags to be recorded on the next model run.
metadata (OrcaMetadataDict | None, default: None ) –

Additional metadata to be recorded on the next model run.
batch_size (int | None, default: None ) –

The batch size to be used for the next model run.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the next model run.

record_model_feedback #

record_model_feedback(
    val, name="default", kind=FeedbackKind.CONTINUOUS
)

Records feedback for the last model runs for which memory lookups were recorded by curate.

Parameters:

val (list[float] | float | int | list[int]) –

The feedback to be recorded.
name (str, default: 'default' ) –

The name of the feedback.
kind (FeedbackKind, default: CONTINUOUS ) –

The kind of feedback.

record_model_input_output #

record_model_input_output(inputs, outputs)

Records the inputs and outputs of the last model runs for which memory lookups were recorded by curate.

Parameters:

inputs (list[Any] | Any) –

The inputs to be recorded.
outputs (list[Any] | Any) –

The outputs to be recorded.

get_lookup_setting_summary #

1	`get_lookup_setting_summary()`

Returns a summary of the lookup settings for each OrcaLookupLayer in this module and its descendants.

forward #

forward(x=None, ctx_labels=None, ctx_scores=None)

Generate logits based on the nearest neighbors of the input vector.

Parameters:

x (Tensor | None, default: None ) –

The input tensor of shape (batch_size, embedding_dim), can be omitted if labels and scores are provided directly.
ctx_labels (Tensor | None, default: None ) –

The memory label tensor of shape (batch_size, num_memories) contains integer labels. If this is None, the labels will be looked up from the index based on the input tensor.
ctx_scores (Tensor | None, default: None ) –

The memory score tensor of shape (batch_size, num_memories) contains float scores. If this is None, the scores will be looked up from the index based on the input tensor.

Returns:

Tensor –

The output tensor of shape (batch_size, num_classes), if neither x nor scores are provided the dtype will be float32, otherwise it will be the same as the scores or input tensor.

OrcaMoeClassificationHead #

OrcaMoeClassificationHead(
    model_dim,
    num_classes,
    num_memories,
    label_column_name,
    gate_layers=1,
    hidden_layers=0,
    activation=F.relu,
    dropout=0.1,
    database=None,
    curate_enabled=False,
    memory_index_name=None,
    drop_exact_match=DropExactMatchOption.TRAINING_ONLY,
    exact_match_threshold=None,
)

Bases: OrcaLookupModule

A mixture of experts classification head that combines a KNN classifier with a linear classifier.

Examples:

import torch
from orcalib import OrcaModel, OrcaMoeClassificationHead

class MyModule(OrcaModel):
    def __init__(self):
        super().__init__()
        self.trunk = torch.nn.Linear(10, 10)
        self.classifier = OrcaMoeClassificationHead(model_dim=10, num_classes=5, num_memories=10)

    def forward(self, x):
        x = self.trunk(x)
        x = self.classifier(x)
        return x

Parameters:

model_dim (int) –

The dimension of the input vector and hidden layers.
num_classes (int) –

The size of the output vector.
num_memories (int) –

The number of memory vectors to be returned from the lookup.
label_column_name (ColumnName) –

The name of the label column to return from the index.
gate_layers (int, default: 1 ) –

The number of layers to use in the gating network.
hidden_layers (int, default: 0 ) –

The number of hidden layers to use in the linear classifier.
activation (Callable[[Tensor], Tensor], default: relu ) –

The activation function.
dropout (float, default: 0.1 ) –

The dropout rate.
database (OrcaDatabase | str | None, default: None ) –

The OrcaDatabase instance to use for lookups and curate tracking.
curate_enabled (bool, default: False ) –

Whether Curate tracking is enabled.
memory_index_name (str | None, default: None ) –

The name of the index to use for lookups.
drop_exact_match (DropExactMatchOption, default: TRAINING_ONLY ) –

Choose to drop the exact match (if found) always, never, or only during training or inference.
exact_match_threshold (float | None, default: None ) –

Minimum similarity score for something to be considered the exact match

lookup_result_transforms `property` `writable` #

1	`lookup_result_transforms`

A list of transforms to apply to the lookup result. NOTE: This will be applied even when lookup_result_override is set.

extra_lookup_column_names `property` `writable` #

1	`extra_lookup_column_names`

While set, all lookups will include these additional columns. They may inclue columns on the indexed table as well as index-specific columns, e.g., $score, $embedding.

lookup_query_override `property` `writable` #

1	`lookup_query_override`

The query to use instead of performing a lookup. NOTE: This will be ignored if lookup_result_override is also set.

lookup_result_override `property` `writable` #

1	`lookup_result_override`

The lookup result to use instead of performing a lookup.

lookup_database `property` `writable` #

1	`lookup_database`

The name of the database to use for looking up memories.

memory_index_name `property` `writable` #

1	`memory_index_name`

The name of the index to use for looking up memories.

lookup_column_names `property` `writable` #

1	`lookup_column_names`

The names of the columns to retrieve for each memory.

num_memories `property` `writable` #

1	`num_memories`

The number of memories to look up.

drop_exact_match `property` `writable` #

1	`drop_exact_match`

Whether to drop exact matches from the results.

exact_match_threshold `property` `writable` #

1	`exact_match_threshold`

The similarity threshold for exact matches.

shuffle_memories `property` `writable` #

1	`shuffle_memories`

Whether to shuffle the looked up memories.

curate_database `property` `writable` #

1	`curate_database`

The name of the database to use for saving curate tracking data.

curate_next_run_settings `property` `writable` #

1	`curate_next_run_settings`

The settings for the next curate model run.

curate_model_id `property` `writable` #

1	`curate_model_id`

The model id to associate with curated model runs.

curate_model_version `property` `writable` #

1	`curate_model_version`

The model version to associate with curated model runs.

curate_metadata `property` `writable` #

1	`curate_metadata`

The metadata to attach to curated model runs.

curate_tags `property` `writable` #

1	`curate_tags`

The tags to attach to the curated model runs.

curate_seq_id `property` `writable` #

1	`curate_seq_id`

The sequence id to associate with curated model runs.

curate_batch_size `property` `writable` #

1	`curate_batch_size`

The batch size of the model run to track curate data for, usually inferred automatically.

last_curate_run_ids `property` `writable` #

1	`last_curate_run_ids`

The run ids of the last model run for which curate tracking data was collected.

last_curate_run_settings `property` `writable` #

1	`last_curate_run_settings`

The settings of the last model run for which curate tracking data was collected.

get_effective_lookup_settings #

1	`get_effective_lookup_settings()`

Returns the effective lookup settings for this module, with any inherited settings applied.

Returns:

LookupSettings –

The effective lookup settings for this module. Practically, this be the lookup settings
LookupSettings –

set on this module. For any settings that are not set on this module, the inherited settings
LookupSettings –

will be used instead.

get_lookup_database_instance #

1	`get_lookup_database_instance()`

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

get_orca_modules_recursively(
    max_depth=None, include_self=True, filter_type=None
)

Recursively yields all children of this module that are instances of the specified filter type.

All parent nodes will be processed before their children
This will search through all children —even those that are not a subclass of the filter_type— but it only returns children that are a subclass of filter_type.

Parameters:

max_depth (int | None, default: None ) –
The maximum depth to search.
- Setting this to 0 will only include this module.
- Setting this to 1 will include only this module and its children.
- Setting it to None (the default) will search through all modules.
- Modules that are not of filter_type or OrcaModule do not increment the depth.
include_self (bool, default: True ) –

Whether to include the current OrcaModule in the results.
filter_type (type[ORCA_MODULE_TYPE] | None, default: None ) –

The subtype of OrcaModule to filter for. If None, any subtypes of OrcaModule will be returned.

Yields:

ORCA_MODULE_TYPE –

modules of type filter_type that are used in the children of this module.

enable_curate #

enable_curate(recursive=True)

Enable Curate tracking for the model and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to enable Curate tracking recursively.

disable_curate #

disable_curate(recursive=True)

Disable Curate tracking for this module and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to disable Curate tracking recursively.

update_curate_settings #

update_curate_settings(
    model_id=None,
    model_version=None,
    tags=None,
    extra_tags=None,
    metadata=None,
    extra_metadata=None,
    batch_size=None,
    seq_id=None,
    enabled=None,
    enable_recursive=True,
)

Update curate tracking settings for the module and all its children.

Parameters:

model_id (str | None, default: None ) –

The ID of the model.
model_version (str | None, default: None ) –

The version of the model.
tags (Iterable[str] | None, default: None ) –

The new tags to be added to the model.
extra_tags (Iterable[str] | None, default: None ) –

The extra tags to be added to the model.
metadata (OrcaMetadataDict | None, default: None ) –

The new metadata to be added to the model.
extra_metadata (OrcaMetadataDict | None, default: None ) –

The extra metadata to be added to the model.
batch_size (int | None, default: None ) –

The batch size to be used for the model.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the model.

record_next_model_memory_lookups #

record_next_model_memory_lookups(
    tags=None, metadata=None, batch_size=None, seq_id=None
)

Sets up curate tracking for the memory lookups during the next forward pass only.

Parameters:

tags (Iterable[str] | None, default: None ) –

Additional tags to be recorded on the next model run.
metadata (OrcaMetadataDict | None, default: None ) –

Additional metadata to be recorded on the next model run.
batch_size (int | None, default: None ) –

The batch size to be used for the next model run.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the next model run.

record_model_feedback #

record_model_feedback(
    val, name="default", kind=FeedbackKind.CONTINUOUS
)

Records feedback for the last model runs for which memory lookups were recorded by curate.

Parameters:

val (list[float] | float | int | list[int]) –

The feedback to be recorded.
name (str, default: 'default' ) –

The name of the feedback.
kind (FeedbackKind, default: CONTINUOUS ) –

The kind of feedback.

record_model_input_output #

record_model_input_output(inputs, outputs)

Records the inputs and outputs of the last model runs for which memory lookups were recorded by curate.

Parameters:

inputs (list[Any] | Any) –

The inputs to be recorded.
outputs (list[Any] | Any) –

The outputs to be recorded.

get_lookup_setting_summary #

1	`get_lookup_setting_summary()`

Returns a summary of the lookup settings for each OrcaLookupLayer in this module and its descendants.

forward #

forward(x, ctx_scores=None, ctx_labels=None)

Generate logits based on the input vector and memory context.

Parameters:

x (Tensor) –

The input tensor of shape (batch_size, embedding_dim).
ctx_scores (Tensor | None, default: None ) –

The memory scores tensor of shape (batch_size, num_memories). If this is None, the scores will be looked up from the index based on the input tensor.
ctx_labels (Tensor | None, default: None ) –

The memory labels tensor of shape (batch_size, num_memories). If this is None, the labels will be looked up from the index based on the input tensor

Returns:

Tensor –

The logits tensor of shape (batch_size, num_classes).

DropExactMatchOption #

Bases: str, Enum

Determines when to drop exact matches from the results.

Attributes:

ALWAYS –

Always drop exact matches from the results.
NEVER –

Never drop exact matches from the results.
TRAINING_ONLY –

Drop exact matches from the results only during training.
INFERENCE_ONLY –

Drop exact matches from the results only during inference.

OrcaLookupLayer #

OrcaLookupLayer(
    database=None,
    curate_enabled=False,
    memory_index_name=None,
    lookup_column_names=None,
    num_memories=None,
    drop_exact_match=None,
    exact_match_threshold=None,
    shuffle_memories=False,
    freeze_num_memories=False,
    cache_ttl=None,
    layer_name=None,
)

Bases: OrcaLookupModule

A layer to perform memory lookups from a database index.

Note

This requires a database to be attached to the model, with the index already created.

Examples:

import torch
from orcalib import OrcaModule, OrcaLookupLayer

class MyLookupModule(OrcaModule):
    def __init__(self):
        super().__init__()
        self.lookup = OrcaLookupLayer(
            memory_index_name="text_index",
            lookup_column_names=["label, "$embedding"],
            num_memories=10
        )

    def forward(self, x: Tensor) -> Tuple[Tensor, Tensor]:
        res = self.lookup(x)
        # ctx is a float tensor of shape (batch_size, num_memories, embedding_dim)
        ctx = res.to_tensor("$embedding", dtype=x.dtype, device=x.device)
        # ctx_labels is an integer tensor of shape (batch_size, num_memories)
        ctx_labels = res.to_tensor("label", dtype=torch.int64, device=x.device).squeeze(-1)
        return ctx, ctx_labels

Parameters:

database (OrcaDatabase | str | None, default: None ) –

The OrcaDatabase instance to use for lookups and curate tracking.
curate_enabled (bool, default: False ) –

Whether Curate tracking is enabled.
memory_index_name (str | None, default: None ) –

The name of the index to use for lookups.
lookup_column_names (list[str] | None, default: None ) –

The names of the columns to return from the index during a lookup.
num_memories (int | None, default: None ) –

The number of memories to return from the index during a lookup.
drop_exact_match (DropExactMatchOption | None, default: None ) –

Choose to drop the exact match (if found) always, never, or only during training or inference.
exact_match_threshold (float | None, default: None ) –

Minimum similarity score for something to be considered the exact match
shuffle_memories (bool, default: False ) –

Whether to shuffle the memories before returning them.
freeze_num_memories (bool, default: False ) –

Whether the number of memories should be frozen.
cache_ttl (int | None, default: None ) –

The time-to-live for the lookup cache.

lookup_result_transforms `property` `writable` #

1	`lookup_result_transforms`

A list of transforms to apply to the lookup result. NOTE: This will be applied even when lookup_result_override is set.

extra_lookup_column_names `property` `writable` #

1	`extra_lookup_column_names`

While set, all lookups will include these additional columns. They may inclue columns on the indexed table as well as index-specific columns, e.g., $score, $embedding.

lookup_query_override `property` `writable` #

1	`lookup_query_override`

The query to use instead of performing a lookup. NOTE: This will be ignored if lookup_result_override is also set.

lookup_result_override `property` `writable` #

1	`lookup_result_override`

The lookup result to use instead of performing a lookup.

lookup_database `property` `writable` #

1	`lookup_database`

The name of the database to use for looking up memories.

memory_index_name `property` `writable` #

1	`memory_index_name`

The name of the index to use for looking up memories.

lookup_column_names `property` `writable` #

1	`lookup_column_names`

The names of the columns to retrieve for each memory.

num_memories `property` `writable` #

1	`num_memories`

The number of memories to look up.

drop_exact_match `property` `writable` #

1	`drop_exact_match`

Whether to drop exact matches from the results.

exact_match_threshold `property` `writable` #

1	`exact_match_threshold`

The similarity threshold for exact matches.

shuffle_memories `property` `writable` #

1	`shuffle_memories`

Whether to shuffle the looked up memories.

curate_database `property` `writable` #

1	`curate_database`

The name of the database to use for saving curate tracking data.

curate_next_run_settings `property` `writable` #

1	`curate_next_run_settings`

The settings for the next curate model run.

curate_model_id `property` `writable` #

1	`curate_model_id`

The model id to associate with curated model runs.

curate_model_version `property` `writable` #

1	`curate_model_version`

The model version to associate with curated model runs.

curate_metadata `property` `writable` #

1	`curate_metadata`

The metadata to attach to curated model runs.

curate_tags `property` `writable` #

1	`curate_tags`

The tags to attach to the curated model runs.

curate_seq_id `property` `writable` #

1	`curate_seq_id`

The sequence id to associate with curated model runs.

curate_batch_size `property` `writable` #

1	`curate_batch_size`

The batch size of the model run to track curate data for, usually inferred automatically.

last_curate_run_ids `property` `writable` #

1	`last_curate_run_ids`

The run ids of the last model run for which curate tracking data was collected.

last_curate_run_settings `property` `writable` #

1	`last_curate_run_settings`

The settings of the last model run for which curate tracking data was collected.

get_effective_lookup_settings #

1	`get_effective_lookup_settings()`

Returns the effective lookup settings for this module, with any inherited settings applied.

Returns:

LookupSettings –

The effective lookup settings for this module. Practically, this be the lookup settings
LookupSettings –

set on this module. For any settings that are not set on this module, the inherited settings
LookupSettings –

will be used instead.

get_lookup_database_instance #

1	`get_lookup_database_instance()`

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

get_orca_modules_recursively(
    max_depth=None, include_self=True, filter_type=None
)

Recursively yields all children of this module that are instances of the specified filter type.

All parent nodes will be processed before their children
This will search through all children —even those that are not a subclass of the filter_type— but it only returns children that are a subclass of filter_type.

Parameters:

max_depth (int | None, default: None ) –
The maximum depth to search.
- Setting this to 0 will only include this module.
- Setting this to 1 will include only this module and its children.
- Setting it to None (the default) will search through all modules.
- Modules that are not of filter_type or OrcaModule do not increment the depth.
include_self (bool, default: True ) –

Whether to include the current OrcaModule in the results.
filter_type (type[ORCA_MODULE_TYPE] | None, default: None ) –

The subtype of OrcaModule to filter for. If None, any subtypes of OrcaModule will be returned.

Yields:

ORCA_MODULE_TYPE –

modules of type filter_type that are used in the children of this module.

enable_curate #

enable_curate(recursive=True)

Enable Curate tracking for the model and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to enable Curate tracking recursively.

disable_curate #

disable_curate(recursive=True)

Disable Curate tracking for this module and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to disable Curate tracking recursively.

update_curate_settings #

update_curate_settings(
    model_id=None,
    model_version=None,
    tags=None,
    extra_tags=None,
    metadata=None,
    extra_metadata=None,
    batch_size=None,
    seq_id=None,
    enabled=None,
    enable_recursive=True,
)

Update curate tracking settings for the module and all its children.

Parameters:

model_id (str | None, default: None ) –

The ID of the model.
model_version (str | None, default: None ) –

The version of the model.
tags (Iterable[str] | None, default: None ) –

The new tags to be added to the model.
extra_tags (Iterable[str] | None, default: None ) –

The extra tags to be added to the model.
metadata (OrcaMetadataDict | None, default: None ) –

The new metadata to be added to the model.
extra_metadata (OrcaMetadataDict | None, default: None ) –

The extra metadata to be added to the model.
batch_size (int | None, default: None ) –

The batch size to be used for the model.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the model.

record_next_model_memory_lookups #

record_next_model_memory_lookups(
    tags=None, metadata=None, batch_size=None, seq_id=None
)

Sets up curate tracking for the memory lookups during the next forward pass only.

Parameters:

tags (Iterable[str] | None, default: None ) –

Additional tags to be recorded on the next model run.
metadata (OrcaMetadataDict | None, default: None ) –

Additional metadata to be recorded on the next model run.
batch_size (int | None, default: None ) –

The batch size to be used for the next model run.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the next model run.

record_model_feedback #

record_model_feedback(
    val, name="default", kind=FeedbackKind.CONTINUOUS
)

Records feedback for the last model runs for which memory lookups were recorded by curate.

Parameters:

val (list[float] | float | int | list[int]) –

The feedback to be recorded.
name (str, default: 'default' ) –

The name of the feedback.
kind (FeedbackKind, default: CONTINUOUS ) –

The kind of feedback.

record_model_input_output #

record_model_input_output(inputs, outputs)

Records the inputs and outputs of the last model runs for which memory lookups were recorded by curate.

Parameters:

inputs (list[Any] | Any) –

The inputs to be recorded.
outputs (list[Any] | Any) –

The outputs to be recorded.

get_lookup_setting_summary #

1	`get_lookup_setting_summary()`

Returns a summary of the lookup settings for each OrcaLookupLayer in this module and its descendants.

forward #

forward(
    x,
    orca_db_instance=None,
    index_name=None,
    lookup_column_names=None,
    num_memories=None,
)

Perform a vector index scan and return the top num_memories results.

Parameters:

x (Tensor | str | list[str]) –

The input tensor of shape (batch_size, embedding_dim)
orca_db_instance (OrcaDatabase | None, default: None ) –

Override for the database to use.
index_name (str | None, default: None ) –

Override for the name of the index to use.
lookup_column_names (list[str] | None, default: None ) –

Override for the names of the columns to return.
num_memories (int | None, default: None ) –

Override for the number of memories to return.

Returns:

BatchedScanResult –

The batch of lookup results, use to_tensor to convert columns from the result to tensors.

OrcaLookupModule #

OrcaLookupModule(
    database=None,
    model_id=None,
    model_version=None,
    metadata=None,
    curate_enabled=False,
    tags=None,
    memory_index_name=None,
    lookup_column_names=None,
    num_memories=None,
    drop_exact_match=None,
    exact_match_threshold=None,
    shuffle_memories=False,
    freeze_num_memories=False,
    propagate_lookup_settings=True,
)

Bases: OrcaModule, LookupSettingsMixin, PostInitMixin

OrcaLookupModule is the base class for all Orca PyTorch layers that support memory lookups —–either directly or through their children— in addition to curate tracking.

Note

Lookup settings are propagated to all children of this module, by recursively setting the same lookup settings instance for all its children.

Parameters:

database (OrcaDatabase | str | None, default: None ) –

The OrcaDatabase instance to use for Curate and memory lookups.
model_id (str | None, default: None ) –

The ID of the model.
model_version (str | None, default: None ) –

The version of the model.
metadata (OrcaMetadataDict | None, default: None ) –

The metadata to be stored with Curate runs.
curate_enabled (bool, default: False ) –

Whether Curate tracking is enabled.
tags (Iterable[str] | None, default: None ) –

The tags to be included in Curate runs.
memory_index_name (str | None, default: None ) –

The name of the index to use for lookups.
lookup_column_names (list[str] | None, default: None ) –

The names of the columns to return from the index during a lookup.
num_memories (int | None, default: None ) –

The number of memories to return from the index during a lookup.
freeze_num_memories (bool, default: False ) –

Whether the number of memories should be frozen. If set to True, an error will be raised if an attempt is made to change the number of memories.
drop_exact_match (DropExactMatchOption | None, default: None ) –

Choose to drop the exact match (if found) always, never, or only during training or inference.
exact_match_threshold (float | None, default: None ) –

Minimum similarity score for something to be considered the exact match
shuffle_memories (bool, default: False ) –

Whether to shuffle the memories before returning them.
propagate_lookup_settings (bool, default: True ) –

Whether to propagate lookup settings to all children.

lookup_result_transforms `property` `writable` #

1	`lookup_result_transforms`

A list of transforms to apply to the lookup result. NOTE: This will be applied even when lookup_result_override is set.

extra_lookup_column_names `property` `writable` #

1	`extra_lookup_column_names`

While set, all lookups will include these additional columns. They may inclue columns on the indexed table as well as index-specific columns, e.g., $score, $embedding.

lookup_query_override `property` `writable` #

1	`lookup_query_override`

The query to use instead of performing a lookup. NOTE: This will be ignored if lookup_result_override is also set.

lookup_result_override `property` `writable` #

1	`lookup_result_override`

The lookup result to use instead of performing a lookup.

lookup_database `property` `writable` #

1	`lookup_database`

The name of the database to use for looking up memories.

memory_index_name `property` `writable` #

1	`memory_index_name`

The name of the index to use for looking up memories.

lookup_column_names `property` `writable` #

1	`lookup_column_names`

The names of the columns to retrieve for each memory.

num_memories `property` `writable` #

1	`num_memories`

The number of memories to look up.

drop_exact_match `property` `writable` #

1	`drop_exact_match`

Whether to drop exact matches from the results.

exact_match_threshold `property` `writable` #

1	`exact_match_threshold`

The similarity threshold for exact matches.

shuffle_memories `property` `writable` #

1	`shuffle_memories`

Whether to shuffle the looked up memories.

curate_database `property` `writable` #

1	`curate_database`

The name of the database to use for saving curate tracking data.

curate_next_run_settings `property` `writable` #

1	`curate_next_run_settings`

The settings for the next curate model run.

curate_model_id `property` `writable` #

1	`curate_model_id`

The model id to associate with curated model runs.

curate_model_version `property` `writable` #

1	`curate_model_version`

The model version to associate with curated model runs.

curate_metadata `property` `writable` #

1	`curate_metadata`

The metadata to attach to curated model runs.

curate_tags `property` `writable` #

1	`curate_tags`

The tags to attach to the curated model runs.

curate_seq_id `property` `writable` #

1	`curate_seq_id`

The sequence id to associate with curated model runs.

curate_batch_size `property` `writable` #

1	`curate_batch_size`

The batch size of the model run to track curate data for, usually inferred automatically.

last_curate_run_ids `property` `writable` #

1	`last_curate_run_ids`

The run ids of the last model run for which curate tracking data was collected.

last_curate_run_settings `property` `writable` #

1	`last_curate_run_settings`

The settings of the last model run for which curate tracking data was collected.

get_effective_lookup_settings #

1	`get_effective_lookup_settings()`

Returns the effective lookup settings for this module, with any inherited settings applied.

Returns:

LookupSettings –

The effective lookup settings for this module. Practically, this be the lookup settings
LookupSettings –

set on this module. For any settings that are not set on this module, the inherited settings
LookupSettings –

will be used instead.

get_lookup_database_instance #

1	`get_lookup_database_instance()`

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

get_orca_modules_recursively(
    max_depth=None, include_self=True, filter_type=None
)

Recursively yields all children of this module that are instances of the specified filter type.

All parent nodes will be processed before their children
This will search through all children —even those that are not a subclass of the filter_type— but it only returns children that are a subclass of filter_type.

Parameters:

max_depth (int | None, default: None ) –
The maximum depth to search.
- Setting this to 0 will only include this module.
- Setting this to 1 will include only this module and its children.
- Setting it to None (the default) will search through all modules.
- Modules that are not of filter_type or OrcaModule do not increment the depth.
include_self (bool, default: True ) –

Whether to include the current OrcaModule in the results.
filter_type (type[ORCA_MODULE_TYPE] | None, default: None ) –

The subtype of OrcaModule to filter for. If None, any subtypes of OrcaModule will be returned.

Yields:

ORCA_MODULE_TYPE –

modules of type filter_type that are used in the children of this module.

enable_curate #

enable_curate(recursive=True)

Enable Curate tracking for the model and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to enable Curate tracking recursively.

disable_curate #

disable_curate(recursive=True)

Disable Curate tracking for this module and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to disable Curate tracking recursively.

update_curate_settings #

update_curate_settings(
    model_id=None,
    model_version=None,
    tags=None,
    extra_tags=None,
    metadata=None,
    extra_metadata=None,
    batch_size=None,
    seq_id=None,
    enabled=None,
    enable_recursive=True,
)

Update curate tracking settings for the module and all its children.

Parameters:

model_id (str | None, default: None ) –

The ID of the model.
model_version (str | None, default: None ) –

The version of the model.
tags (Iterable[str] | None, default: None ) –

The new tags to be added to the model.
extra_tags (Iterable[str] | None, default: None ) –

The extra tags to be added to the model.
metadata (OrcaMetadataDict | None, default: None ) –

The new metadata to be added to the model.
extra_metadata (OrcaMetadataDict | None, default: None ) –

The extra metadata to be added to the model.
batch_size (int | None, default: None ) –

The batch size to be used for the model.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the model.

record_next_model_memory_lookups #

record_next_model_memory_lookups(
    tags=None, metadata=None, batch_size=None, seq_id=None
)

Sets up curate tracking for the memory lookups during the next forward pass only.

Parameters:

tags (Iterable[str] | None, default: None ) –

Additional tags to be recorded on the next model run.
metadata (OrcaMetadataDict | None, default: None ) –

Additional metadata to be recorded on the next model run.
batch_size (int | None, default: None ) –

The batch size to be used for the next model run.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the next model run.

record_model_feedback #

record_model_feedback(
    val, name="default", kind=FeedbackKind.CONTINUOUS
)

Records feedback for the last model runs for which memory lookups were recorded by curate.

Parameters:

val (list[float] | float | int | list[int]) –

The feedback to be recorded.
name (str, default: 'default' ) –

The name of the feedback.
kind (FeedbackKind, default: CONTINUOUS ) –

The kind of feedback.

record_model_input_output #

record_model_input_output(inputs, outputs)

Records the inputs and outputs of the last model runs for which memory lookups were recorded by curate.

Parameters:

inputs (list[Any] | Any) –

The inputs to be recorded.
outputs (list[Any] | Any) –

The outputs to be recorded.

get_lookup_setting_summary #

1	`get_lookup_setting_summary()`

Returns a summary of the lookup settings for each OrcaLookupLayer in this module and its descendants.

OrcaModel #

OrcaModel(
    database=None,
    model_id=None,
    model_version=None,
    metadata=None,
    curate_enabled=False,
    tags=None,
    memory_index_name=None,
    lookup_column_names=None,
    num_memories=None,
    drop_exact_match=None,
    exact_match_threshold=None,
    shuffle_memories=False,
    freeze_num_memories=False,
    propagate_lookup_settings=True,
)

Bases: OrcaLookupModule, PostInitMixin, PreForwardMixin

OrcaModel should be the base class for all PyTorch models that include Orca layers.

This class is responsible for:

Propagating Curate and memory lookup settings to all children of the model.
Getting curate run ids and preparing tracking before the forward pass.
Building layer names for all children of the model.

This class provides functions to let you:

Enable and disable Curate tracking for the model and all its children.
Record Curate scores for the last run.
Record model input and output for the last run.
Enable and disable memory access for the model and all its children.
Update Curate tracking settings for the model and all its children.

When using OrcaModel, you can set global settings for Curate tracking and memory lookups that will be propagated to all children of the model. This allows you to set these settings once for the entire model and have them automatically applied to all layers.

Parameters:

database (OrcaDatabase | str | None, default: None ) –

The OrcaDatabase instance to use for Curate and memory lookups.
model_id (str | None, default: None ) –

model_id will be included in all runs tracked with curate.
model_version (str | None, default: None ) –

model_version will be included in all runs tracked with curate.
metadata (OrcaMetadataDict | None, default: None ) –

metadata is a dictionary of additional information to be stored with Curate runs.
curate_enabled (bool, default: False ) –

Whether curate tracking is enabled.
tags (Iterable[str] | None, default: None ) –

tags is a set of strings to be included in all runs tracked with curate.
memory_index_name (str | None, default: None ) –

The name of the index to use for lookups.
lookup_column_names (list[str] | None, default: None ) –

The names of the columns to return from the index during a lookup.
num_memories (int | None, default: None ) –

The number of memories to return from the index during a lookup.
drop_exact_match (DropExactMatchOption | None, default: None ) –

Choose to drop the exact match (if found) always, never, or only during training
exact_match_threshold (float | None, default: None ) –

Minimum similarity score for something to be considered the exact match
shuffle_memories (bool, default: False ) –

Whether to shuffle the memories before returning them. (default: False)
freeze_num_memories (bool, default: False ) –

Whether the number of memories should be frozen after initialization. (default: False) When
propagate_lookup_settings (bool, default: True ) –

Whether to propagate lookup settings to all children. (default: False)

or inference. (default: NEVER) set to True, an error will be raised if an attempt is made to change the number of memories.

lookup_result_transforms `property` `writable` #

1	`lookup_result_transforms`

A list of transforms to apply to the lookup result. NOTE: This will be applied even when lookup_result_override is set.

extra_lookup_column_names `property` `writable` #

1	`extra_lookup_column_names`

While set, all lookups will include these additional columns. They may inclue columns on the indexed table as well as index-specific columns, e.g., $score, $embedding.

lookup_query_override `property` `writable` #

1	`lookup_query_override`

The query to use instead of performing a lookup. NOTE: This will be ignored if lookup_result_override is also set.

lookup_result_override `property` `writable` #

1	`lookup_result_override`

The lookup result to use instead of performing a lookup.

lookup_database `property` `writable` #

1	`lookup_database`

The name of the database to use for looking up memories.

memory_index_name `property` `writable` #

1	`memory_index_name`

The name of the index to use for looking up memories.

lookup_column_names `property` `writable` #

1	`lookup_column_names`

The names of the columns to retrieve for each memory.

num_memories `property` `writable` #

1	`num_memories`

The number of memories to look up.

drop_exact_match `property` `writable` #

1	`drop_exact_match`

Whether to drop exact matches from the results.

exact_match_threshold `property` `writable` #

1	`exact_match_threshold`

The similarity threshold for exact matches.

shuffle_memories `property` `writable` #

1	`shuffle_memories`

Whether to shuffle the looked up memories.

curate_database `property` `writable` #

1	`curate_database`

The name of the database to use for saving curate tracking data.

curate_next_run_settings `property` `writable` #

1	`curate_next_run_settings`

The settings for the next curate model run.

curate_model_id `property` `writable` #

1	`curate_model_id`

The model id to associate with curated model runs.

curate_model_version `property` `writable` #

1	`curate_model_version`

The model version to associate with curated model runs.

curate_metadata `property` `writable` #

1	`curate_metadata`

The metadata to attach to curated model runs.

curate_tags `property` `writable` #

1	`curate_tags`

The tags to attach to the curated model runs.

curate_seq_id `property` `writable` #

1	`curate_seq_id`

The sequence id to associate with curated model runs.

curate_batch_size `property` `writable` #

1	`curate_batch_size`

The batch size of the model run to track curate data for, usually inferred automatically.

last_curate_run_ids `property` `writable` #

1	`last_curate_run_ids`

The run ids of the last model run for which curate tracking data was collected.

last_curate_run_settings `property` `writable` #

1	`last_curate_run_settings`

The settings of the last model run for which curate tracking data was collected.

get_effective_lookup_settings #

1	`get_effective_lookup_settings()`

Returns the effective lookup settings for this module, with any inherited settings applied.

Returns:

LookupSettings –

The effective lookup settings for this module. Practically, this be the lookup settings
LookupSettings –

set on this module. For any settings that are not set on this module, the inherited settings
LookupSettings –

will be used instead.

get_lookup_database_instance #

1	`get_lookup_database_instance()`

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

get_orca_modules_recursively(
    max_depth=None, include_self=True, filter_type=None
)

Recursively yields all children of this module that are instances of the specified filter type.

All parent nodes will be processed before their children
This will search through all children —even those that are not a subclass of the filter_type— but it only returns children that are a subclass of filter_type.

Parameters:

max_depth (int | None, default: None ) –
The maximum depth to search.
- Setting this to 0 will only include this module.
- Setting this to 1 will include only this module and its children.
- Setting it to None (the default) will search through all modules.
- Modules that are not of filter_type or OrcaModule do not increment the depth.
include_self (bool, default: True ) –

Whether to include the current OrcaModule in the results.
filter_type (type[ORCA_MODULE_TYPE] | None, default: None ) –

The subtype of OrcaModule to filter for. If None, any subtypes of OrcaModule will be returned.

Yields:

ORCA_MODULE_TYPE –

modules of type filter_type that are used in the children of this module.

enable_curate #

enable_curate(recursive=True)

Enable Curate tracking for the model and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to enable Curate tracking recursively.

disable_curate #

disable_curate(recursive=True)

Disable Curate tracking for this module and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to disable Curate tracking recursively.

update_curate_settings #

update_curate_settings(
    model_id=None,
    model_version=None,
    tags=None,
    extra_tags=None,
    metadata=None,
    extra_metadata=None,
    batch_size=None,
    seq_id=None,
    enabled=None,
    enable_recursive=True,
)

Update curate tracking settings for the module and all its children.

Parameters:

model_id (str | None, default: None ) –

The ID of the model.
model_version (str | None, default: None ) –

The version of the model.
tags (Iterable[str] | None, default: None ) –

The new tags to be added to the model.
extra_tags (Iterable[str] | None, default: None ) –

The extra tags to be added to the model.
metadata (OrcaMetadataDict | None, default: None ) –

The new metadata to be added to the model.
extra_metadata (OrcaMetadataDict | None, default: None ) –

The extra metadata to be added to the model.
batch_size (int | None, default: None ) –

The batch size to be used for the model.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the model.

record_next_model_memory_lookups #

record_next_model_memory_lookups(
    tags=None, metadata=None, batch_size=None, seq_id=None
)

Sets up curate tracking for the memory lookups during the next forward pass only.

Parameters:

tags (Iterable[str] | None, default: None ) –

Additional tags to be recorded on the next model run.
metadata (OrcaMetadataDict | None, default: None ) –

Additional metadata to be recorded on the next model run.
batch_size (int | None, default: None ) –

The batch size to be used for the next model run.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the next model run.

record_model_feedback #

record_model_feedback(
    val, name="default", kind=FeedbackKind.CONTINUOUS
)

Records feedback for the last model runs for which memory lookups were recorded by curate.

Parameters:

val (list[float] | float | int | list[int]) –

The feedback to be recorded.
name (str, default: 'default' ) –

The name of the feedback.
kind (FeedbackKind, default: CONTINUOUS ) –

The kind of feedback.

record_model_input_output #

record_model_input_output(inputs, outputs)

Records the inputs and outputs of the last model runs for which memory lookups were recorded by curate.

Parameters:

inputs (list[Any] | Any) –

The inputs to be recorded.
outputs (list[Any] | Any) –

The outputs to be recorded.

get_lookup_setting_summary #

1	`get_lookup_setting_summary()`

Returns a summary of the lookup settings for each OrcaLookupLayer in this module and its descendants.

enable_memory #

1	`enable_memory()`

Enables memory access for the model and all its children.

disable_memory #

1	`disable_memory()`

Disables memory access for the model and all its children.

OrcaModule #

OrcaModule(
    curate_database=None,
    model_id=None,
    model_version=None,
    metadata=None,
    curate_enabled=False,
    tags=None,
)

Bases: Module, CurateSettingsMixin

OrcaModule is a special PyTorch Module that support tracking curate data during forward passes. It is the base class for all Orca PyTorch layers.

Note

Curate setting propagation is handled by the OrcaModel (model, not module) class, which recursively sets the curate settings for all children of the model to use the same instance of the curate settings object.

Parameters:

curate_database (OrcaDatabase | str | None, default: None ) –

The OrcaDatabase instance to use for Curate tracking.
model_id (str | None, default: None ) –

The ID of the model.
model_version (str | None, default: None ) –

The version of the model.
metadata (OrcaMetadataDict | None, default: None ) –

The metadata to be stored with Curate runs.
curate_enabled (bool, default: False ) –

Whether Curate tracking is enabled.
tags (Iterable[str] | None, default: None ) –

The tags to be included in Curate runs.

curate_database `property` `writable` #

1	`curate_database`

The name of the database to use for saving curate tracking data.

curate_next_run_settings `property` `writable` #

1	`curate_next_run_settings`

The settings for the next curate model run.

curate_model_id `property` `writable` #

1	`curate_model_id`

The model id to associate with curated model runs.

curate_model_version `property` `writable` #

1	`curate_model_version`

The model version to associate with curated model runs.

curate_metadata `property` `writable` #

1	`curate_metadata`

The metadata to attach to curated model runs.

curate_tags `property` `writable` #

1	`curate_tags`

The tags to attach to the curated model runs.

curate_seq_id `property` `writable` #

1	`curate_seq_id`

The sequence id to associate with curated model runs.

curate_batch_size `property` `writable` #

1	`curate_batch_size`

The batch size of the model run to track curate data for, usually inferred automatically.

last_curate_run_ids `property` `writable` #

1	`last_curate_run_ids`

The run ids of the last model run for which curate tracking data was collected.

last_curate_run_settings `property` `writable` #

1	`last_curate_run_settings`

The settings of the last model run for which curate tracking data was collected.

get_orca_modules_recursively #

get_orca_modules_recursively(
    max_depth=None, include_self=True, filter_type=None
)

Recursively yields all children of this module that are instances of the specified filter type.

All parent nodes will be processed before their children
This will search through all children —even those that are not a subclass of the filter_type— but it only returns children that are a subclass of filter_type.

Parameters:

max_depth (int | None, default: None ) –
The maximum depth to search.
- Setting this to 0 will only include this module.
- Setting this to 1 will include only this module and its children.
- Setting it to None (the default) will search through all modules.
- Modules that are not of filter_type or OrcaModule do not increment the depth.
include_self (bool, default: True ) –

Whether to include the current OrcaModule in the results.
filter_type (type[ORCA_MODULE_TYPE] | None, default: None ) –

The subtype of OrcaModule to filter for. If None, any subtypes of OrcaModule will be returned.

Yields:

ORCA_MODULE_TYPE –

modules of type filter_type that are used in the children of this module.

enable_curate #

enable_curate(recursive=True)

Enable Curate tracking for the model and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to enable Curate tracking recursively.

disable_curate #

disable_curate(recursive=True)

Disable Curate tracking for this module and (if recursive is True) for all its descendants.

Parameters:

recursive (bool, default: True ) –

Whether to disable Curate tracking recursively.

update_curate_settings #

update_curate_settings(
    model_id=None,
    model_version=None,
    tags=None,
    extra_tags=None,
    metadata=None,
    extra_metadata=None,
    batch_size=None,
    seq_id=None,
    enabled=None,
    enable_recursive=True,
)

Update curate tracking settings for the module and all its children.

Parameters:

model_id (str | None, default: None ) –

The ID of the model.
model_version (str | None, default: None ) –

The version of the model.
tags (Iterable[str] | None, default: None ) –

The new tags to be added to the model.
extra_tags (Iterable[str] | None, default: None ) –

The extra tags to be added to the model.
metadata (OrcaMetadataDict | None, default: None ) –

The new metadata to be added to the model.
extra_metadata (OrcaMetadataDict | None, default: None ) –

The extra metadata to be added to the model.
batch_size (int | None, default: None ) –

The batch size to be used for the model.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the model.

record_next_model_memory_lookups #

record_next_model_memory_lookups(
    tags=None, metadata=None, batch_size=None, seq_id=None
)

Sets up curate tracking for the memory lookups during the next forward pass only.

Parameters:

tags (Iterable[str] | None, default: None ) –

Additional tags to be recorded on the next model run.
metadata (OrcaMetadataDict | None, default: None ) –

Additional metadata to be recorded on the next model run.
batch_size (int | None, default: None ) –

The batch size to be used for the next model run.
seq_id (UUID | None, default: None ) –

The sequence ID to be used for the next model run.

record_model_feedback #

record_model_feedback(
    val, name="default", kind=FeedbackKind.CONTINUOUS
)

Records feedback for the last model runs for which memory lookups were recorded by curate.

Parameters:

val (list[float] | float | int | list[int]) –

The feedback to be recorded.
name (str, default: 'default' ) –

The name of the feedback.
kind (FeedbackKind, default: CONTINUOUS ) –

The kind of feedback.

record_model_input_output #

record_model_input_output(inputs, outputs)

Records the inputs and outputs of the last model runs for which memory lookups were recorded by curate.

Parameters:

inputs (list[Any] | Any) –

The inputs to be recorded.
outputs (list[Any] | Any) –

The outputs to be recorded.

orcalib.torch#

OrcaLookupCacheBuilder #

memory_column_aliases instance-attribute #

from_model classmethod #

inject_lookup_results #

get_lookup_result #

add_lookups_to_hf_dataset #

OrcaCrossAttentionClassificationHead #

lookup_result_transforms property writable #

extra_lookup_column_names property writable #

lookup_query_override property writable #

lookup_result_override property writable #

lookup_database property writable #

memory_index_name property writable #

lookup_column_names property writable #

num_memories property writable #

drop_exact_match property writable #

exact_match_threshold property writable #

shuffle_memories property writable #

curate_database property writable #

curate_next_run_settings property writable #

curate_model_id property writable #

curate_model_version property writable #

curate_metadata property writable #

curate_tags property writable #

curate_seq_id property writable #

curate_batch_size property writable #

last_curate_run_ids property writable #

last_curate_run_settings property writable #

get_effective_lookup_settings #

get_lookup_database_instance #

get_orca_modules_recursively #

enable_curate #

disable_curate #

update_curate_settings #

record_next_model_memory_lookups #

record_model_feedback #

record_model_input_output #

get_lookup_setting_summary #

forward #

OrcaKnnClassifier #

lookup_result_transforms property writable #

extra_lookup_column_names property writable #

lookup_query_override property writable #

lookup_result_override property writable #

lookup_database property writable #

memory_index_name property writable #

lookup_column_names property writable #

num_memories property writable #

drop_exact_match property writable #

exact_match_threshold property writable #

shuffle_memories property writable #

curate_database property writable #

curate_next_run_settings property writable #

curate_model_id property writable #

curate_model_version property writable #

curate_metadata property writable #

curate_tags property writable #

curate_seq_id property writable #

curate_batch_size property writable #

last_curate_run_ids property writable #

last_curate_run_settings property writable #

get_effective_lookup_settings #

get_lookup_database_instance #

get_orca_modules_recursively #

enable_curate #

disable_curate #

update_curate_settings #

record_next_model_memory_lookups #

record_model_feedback #

record_model_input_output #

get_lookup_setting_summary #

forward #

OrcaMoeClassificationHead #

lookup_result_transforms property writable #

extra_lookup_column_names property writable #

lookup_query_override property writable #

lookup_result_override property writable #

lookup_database property writable #

memory_index_name property writable #

memory_column_aliases `instance-attribute` #

from_model `classmethod` #

lookup_result_transforms `property` `writable` #

extra_lookup_column_names `property` `writable` #

lookup_query_override `property` `writable` #

lookup_result_override `property` `writable` #

lookup_database `property` `writable` #

memory_index_name `property` `writable` #

lookup_column_names `property` `writable` #

num_memories `property` `writable` #

drop_exact_match `property` `writable` #

exact_match_threshold `property` `writable` #

shuffle_memories `property` `writable` #

curate_database `property` `writable` #

curate_next_run_settings `property` `writable` #

curate_model_id `property` `writable` #

curate_model_version `property` `writable` #

curate_metadata `property` `writable` #

curate_tags `property` `writable` #

curate_seq_id `property` `writable` #

curate_batch_size `property` `writable` #

last_curate_run_ids `property` `writable` #

last_curate_run_settings `property` `writable` #

lookup_result_transforms `property` `writable` #

extra_lookup_column_names `property` `writable` #

lookup_query_override `property` `writable` #

lookup_result_override `property` `writable` #

lookup_database `property` `writable` #

memory_index_name `property` `writable` #

lookup_column_names `property` `writable` #

num_memories `property` `writable` #

drop_exact_match `property` `writable` #

exact_match_threshold `property` `writable` #

shuffle_memories `property` `writable` #

curate_database `property` `writable` #

curate_next_run_settings `property` `writable` #

curate_model_id `property` `writable` #

curate_model_version `property` `writable` #

curate_metadata `property` `writable` #

curate_tags `property` `writable` #

curate_seq_id `property` `writable` #

curate_batch_size `property` `writable` #

last_curate_run_ids `property` `writable` #

last_curate_run_settings `property` `writable` #

lookup_result_transforms `property` `writable` #

extra_lookup_column_names `property` `writable` #

lookup_query_override `property` `writable` #

lookup_result_override `property` `writable` #

lookup_database `property` `writable` #

memory_index_name `property` `writable` #

lookup_column_names `property` `writable` #

num_memories `property` `writable` #

drop_exact_match `property` `writable` #

exact_match_threshold `property` `writable` #

shuffle_memories `property` `writable` #

curate_database `property` `writable` #

curate_next_run_settings `property` `writable` #

curate_model_id `property` `writable` #

curate_model_version `property` `writable` #

curate_metadata `property` `writable` #

curate_tags `property` `writable` #

curate_seq_id `property` `writable` #

curate_batch_size `property` `writable` #

last_curate_run_ids `property` `writable` #

last_curate_run_settings `property` `writable` #

lookup_result_transforms `property` `writable` #

extra_lookup_column_names `property` `writable` #

lookup_query_override `property` `writable` #

lookup_result_override `property` `writable` #

lookup_database `property` `writable` #

memory_index_name `property` `writable` #

lookup_column_names `property` `writable` #

num_memories `property` `writable` #

drop_exact_match `property` `writable` #

exact_match_threshold `property` `writable` #

shuffle_memories `property` `writable` #

curate_database `property` `writable` #

curate_next_run_settings `property` `writable` #

curate_model_id `property` `writable` #

curate_model_version `property` `writable` #