Skip to content

orcalib.orca_torch#

OrcaModule #

1
2
3
4
5
6
7
8
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 #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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.

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 #

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 #

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 #

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 #

lookup_result_override

The lookup result to use instead of performing a lookup.

lookup_database property writable #

lookup_database

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

memory_index_name property writable #

memory_index_name

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

lookup_column_names property writable #

lookup_column_names

The names of the columns to retrieve for each memory.

num_memories property writable #

num_memories

The number of memories to look up.

drop_exact_match property writable #

drop_exact_match

Whether to drop exact matches from the results.

exact_match_threshold property writable #

exact_match_threshold

The similarity threshold for exact matches.

shuffle_memories property writable #

shuffle_memories

Whether to shuffle the looked up memories.

curate_database property writable #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_effective_lookup_settings #

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 #

get_lookup_database_instance()

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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 #

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 #

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 #

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 #

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 #

lookup_result_override

The lookup result to use instead of performing a lookup.

lookup_database property writable #

lookup_database

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

memory_index_name property writable #

memory_index_name

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

lookup_column_names property writable #

lookup_column_names

The names of the columns to retrieve for each memory.

num_memories property writable #

num_memories

The number of memories to look up.

drop_exact_match property writable #

drop_exact_match

Whether to drop exact matches from the results.

exact_match_threshold property writable #

exact_match_threshold

The similarity threshold for exact matches.

shuffle_memories property writable #

shuffle_memories

Whether to shuffle the looked up memories.

curate_database property writable #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_effective_lookup_settings #

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 #

get_lookup_database_instance()

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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 #

get_lookup_setting_summary()

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

enable_memory #

enable_memory()

Enables memory access for the model and all its children.

disable_memory #

disable_memory()

Disables memory access for the model and all its children.

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 #

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 #

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 #

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 #

lookup_result_override

The lookup result to use instead of performing a lookup.

lookup_database property writable #

lookup_database

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

memory_index_name property writable #

memory_index_name

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

lookup_column_names property writable #

lookup_column_names

The names of the columns to retrieve for each memory.

num_memories property writable #

num_memories

The number of memories to look up.

drop_exact_match property writable #

drop_exact_match

Whether to drop exact matches from the results.

exact_match_threshold property writable #

exact_match_threshold

The similarity threshold for exact matches.

shuffle_memories property writable #

shuffle_memories

Whether to shuffle the looked up memories.

curate_database property writable #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_effective_lookup_settings #

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 #

get_lookup_database_instance()

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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 #

get_lookup_setting_summary()

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

forward #

1
2
3
4
5
6
7
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:

OrcaLabelLookupLayer #

OrcaLabelLookupLayer(
    database=None,
    curate_enabled=False,
    memory_index_name=None,
    label_column_name=None,
    num_memories=None,
    drop_exact_match=None,
    exact_match_threshold=None,
    shuffle_memories=False,
    freeze_num_memories=False,
)

Bases: OrcaLookupLayer, LabelColumnNameMixin

A layer to lookup embeddings and label from a database index.

This is a convenience layer around OrcaLookupLayer that does not just perform the lookups but also converts the results to tensors.

Note

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

Examples:

import torch
from orcalib import OrcaModule, OrcaLabelLookupLayer

class MyModule(OrcaModule):
    def __init__(self):
        super().__init__()
        self.lookup = OrcaLabelLookupLayer(
            index_name="my_index",
            label_column_name="my_label",
            num_memories=10
        )

    def forward(self, x):
        embeddings, labels = self.lookup(x)
        # embeddings is a float tensor of shape (batch_size, num_memories, embedding_dim)
        # labels is an int tensor of shape (batch_size, num_memories)

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.

  • label_column_name (ColumnName | None, default: None ) –

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

  • 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. When set to True, an error will be raised if an attempt is made to change the number of memories.

lookup_result_transforms property writable #

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 #

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 #

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 #

lookup_result_override

The lookup result to use instead of performing a lookup.

lookup_database property writable #

lookup_database

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

memory_index_name property writable #

memory_index_name

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

lookup_column_names property writable #

lookup_column_names

The names of the columns to retrieve for each memory.

num_memories property writable #

num_memories

The number of memories to look up.

drop_exact_match property writable #

drop_exact_match

Whether to drop exact matches from the results.

exact_match_threshold property writable #

exact_match_threshold

The similarity threshold for exact matches.

shuffle_memories property writable #

shuffle_memories

Whether to shuffle the looked up memories.

curate_database property writable #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_effective_lookup_settings #

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 #

get_lookup_database_instance()

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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 #

get_lookup_setting_summary()

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

forward #

1
2
3
4
5
6
7
forward(
    x,
    orca_db_instance=None,
    index_name=None,
    label_column_name=None,
    num_memories=None,
)

Look up embeddings and labels

Parameters:

  • x (Tensor) –

    The input tensor of shape (batch_size, embedding_dim)

  • orca_db_instance (OrcaDatabase | None, default: None ) –

    Optional override for the OrcaDatabase instance to use.

  • index_name (str | None, default: None ) –

    Optional override for the name of the index to use.

  • label_column_name (ColumnName | None, default: None ) –

    Optional override for the name of the label column to return from the index.

  • num_memories (int | None, default: None ) –

    Optional override for the number of memories to return from the index.

Returns:

  • embeddings ( Tensor ) –

    A tensor with the same type as the input x of shape (batch_size, num_memories, embedding_dim)

  • labels ( Tensor ) –

    An int64 tensor of shape (batch_size, num_memories)

OrcaClassificationMemoryGuideLayer #

OrcaClassificationMemoryGuideLayer(
    num_classes,
    num_memories,
    enable_in_training=False,
    guide_weight=0.1,
    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 PyTorch module that implements a memory-guided classification layer.

This layer biases the output of a classification model towards a set of memories. The bias is controlled by a weight parameter, which determines how strongly the model should be biased towards the memories.

Parameters:

  • num_classes (int) –

    The number of classes in the classification task.

  • num_memories (int) –

    The number of memories the layer should use.

  • enable_in_training (bool, default: False ) –

    Whether to enable the layer during training.

  • guide_weight (float, default: 0.1 ) –

    The weight of the memory guide.

  • 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 #

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 #

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 #

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 #

lookup_result_override

The lookup result to use instead of performing a lookup.

lookup_database property writable #

lookup_database

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

memory_index_name property writable #

memory_index_name

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

lookup_column_names property writable #

lookup_column_names

The names of the columns to retrieve for each memory.

num_memories property writable #

num_memories

The number of memories to look up.

drop_exact_match property writable #

drop_exact_match

Whether to drop exact matches from the results.

exact_match_threshold property writable #

exact_match_threshold

The similarity threshold for exact matches.

shuffle_memories property writable #

shuffle_memories

Whether to shuffle the looked up memories.

curate_database property writable #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_effective_lookup_settings #

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 #

get_lookup_database_instance()

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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 #

get_lookup_setting_summary()

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

forward #

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

Ground the logits based on the memory context.

Parameters:

  • logits (Tensor) –

    Input tensor of shape (batch_size, num_classes)

  • memory_key (Tensor) –

    Memory key tensor of shape (batch_size, embedding_dim)

  • ctx (Tensor | None, default: None ) –

    Memory embeddings tensor of shape (batch_size, num_memories, embedding_dim). If None, the memory context is looked up based on the memory key.

  • labels (Tensor | None, default: None ) –

    Memory label tensor of shape (batch_size, num_memories). If None, the labels are looked up along with the memory context.

Returns:

  • Tensor

    Output tensor of shape (batch_size, num_classes).

OrcaClassificationCrossAttentionLayer #

OrcaClassificationCrossAttentionLayer(
    model_dim,
    num_heads,
    num_classes,
    num_memories,
    dropout=0.1,
    activation=F.relu,
    projection_mode=ProjectionMode.LABEL,
    split_retrieval_path=False,
    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

Note that this is Classification-specific, and the labels returned by the lookup layer are used as the value-weights for the cross attention.

The block contains the typical transformer components: multi-head attention, feed forward, and layer norm. The block also contains a lookup layer that looks up a vector in an OrcaDatabase index and returns the top k results. These results are used as the memory context for the cross attention.

Parameters:

  • model_dim (int) –

    The dimension of the input vector and hidden layers.

  • num_heads (int) –

    The number of heads to be used in the multi-head attention layer.

  • num_classes (int) –

    The number of classes for the output classification and weights for cross attention.

  • num_memories (int) –

    The number of memory vectors to be returned from the lookup.

  • dropout (float, default: 0.1 ) –

    The dropout rate.

  • activation (Callable[[Tensor], Tensor], default: relu ) –

    The activation function.

  • projection_mode (ProjectionMode, default: LABEL ) –

    The projection mode to use for the memory labels.

  • split_retrieval_path (bool, default: False ) –

    Whether to split the retrieval path.

  • 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 #

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 #

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 #

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 #

lookup_result_override

The lookup result to use instead of performing a lookup.

lookup_database property writable #

lookup_database

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

memory_index_name property writable #

memory_index_name

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

lookup_column_names property writable #

lookup_column_names

The names of the columns to retrieve for each memory.

num_memories property writable #

num_memories

The number of memories to look up.

drop_exact_match property writable #

drop_exact_match

Whether to drop exact matches from the results.

exact_match_threshold property writable #

exact_match_threshold

The similarity threshold for exact matches.

shuffle_memories property writable #

shuffle_memories

Whether to shuffle the looked up memories.

curate_database property writable #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_effective_lookup_settings #

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 #

get_lookup_database_instance()

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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 #

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)

Performs cross attention on the input tensor and memory context.

Cross Attention Mechanism:

x, ctx, labels act as Q, K, V for the cross attention layer.

When ctx is None:

  • If split_retrieval_path is False, x is used as both Q and K.
  • If split_retrieval_path is True, memory_key is used as K (instead of x)

The labels are used as V and projected into the embedding space.

Parameters:

  • x (Tensor) –

    The input tensor of shape (batch_size, embedding_dim)

  • ctx (Tensor | None, default: None ) –

    The memory embeddings tensor of shape (batch_size, num_memories, embedding_dim).

  • labels (Tensor | None, default: None ) –

    The memory label tensor of shape (batch_size, num_memories).

  • memory_key (Tensor | None, default: None ) –

    The memory lookup tensor of shape (batch_size, embedding_dim).

Returns:

  • Tensor

    The output tensor of shape (batch_size, embedding_dim).

OrcaMemoryBindingLayer #

OrcaMemoryBindingLayer(
    num_memories,
    num_classes,
    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

Memory binding layer that transforms positional logits (which describe the memories that the model predicted to be relevant) into regular logits, which describe the class the model predicts.

Parameters:

  • num_memories (int) –

    The number of memory vectors to be returned from the lookup.

  • num_classes (int) –

    The number of classes for the output classification and weights for cross attention.

  • 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 #

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 #

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 #

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 #

lookup_result_override

The lookup result to use instead of performing a lookup.

lookup_database property writable #

lookup_database

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

memory_index_name property writable #

memory_index_name

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

lookup_column_names property writable #

lookup_column_names

The names of the columns to retrieve for each memory.

num_memories property writable #

num_memories

The number of memories to look up.

drop_exact_match property writable #

drop_exact_match

Whether to drop exact matches from the results.

exact_match_threshold property writable #

exact_match_threshold

The similarity threshold for exact matches.

shuffle_memories property writable #

shuffle_memories

Whether to shuffle the looked up memories.

curate_database property writable #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_effective_lookup_settings #

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 #

get_lookup_database_instance()

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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 #

get_lookup_setting_summary()

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

forward #

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

Parameters:

  • logits (Tensor) –

    Input tensor with positional logits of shape (batch_size, num_memories).

  • memory_key (Tensor | None, default: None ) –

    Memory key tensor of shape (batch_size, embedding_dim).

  • ctx (Tensor | None, default: None ) –

    Memory context tensor of shape (batch_size, num_memories, embedding_dim). If None, the memory context is looked up based on the memory key.

  • labels (Tensor | None, default: None ) –

    Memory Label tensor of shape (batch_size, num_memories). If None, the labels are looked up along with the memory context.

Returns:

  • Tensor

    Output tensor of shape (batch_size, num_classes).

OrcaLLMMemoryGuideLayer #

OrcaLLMMemoryGuideLayer(
    num_memories,
    alpha,
    beta,
    vocab_size,
    tokenizer,
    S_min=3,
    S_max=10,
    enable_in_training=False,
    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 PyTorch module that implements a memory-guided generation layer for Language Models.

This layer biases the output distribution of the model towards a set of memories.

Parameters:

  • num_memories (int) –

    The number of memories.

  • alpha (float) –

    The alpha parameter for the memory guide.

  • beta (float) –

    The beta parameter for the memory guide.

  • vocab_size (int) –

    The size of the vocabulary.

  • tokenizer (Callable[[str | list[str]], list[int] | list[list[int]]]) –

    The tokenizer function.

  • S_min (int, default: 3 ) –

    The minimum length of the suffixes to search for.

  • S_max (int, default: 10 ) –

    The maximum length of the suffixes to search for.

  • enable_in_training (bool, default: False ) –

    Whether to enable the memory guide layer during training.

  • 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 (str | None, default: None ) –

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

lookup_result_transforms property writable #

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 #

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 #

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 #

lookup_result_override

The lookup result to use instead of performing a lookup.

lookup_database property writable #

lookup_database

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

memory_index_name property writable #

memory_index_name

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

lookup_column_names property writable #

lookup_column_names

The names of the columns to retrieve for each memory.

num_memories property writable #

num_memories

The number of memories to look up.

drop_exact_match property writable #

drop_exact_match

Whether to drop exact matches from the results.

exact_match_threshold property writable #

exact_match_threshold

The similarity threshold for exact matches.

shuffle_memories property writable #

shuffle_memories

Whether to shuffle the looked up memories.

curate_database property writable #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_effective_lookup_settings #

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 #

get_lookup_database_instance()

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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 #

get_lookup_setting_summary()

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

forward #

forward(memory_key, logits, inpt_tokens)

Applies memory guidance to the model logits.

Parameters:

  • memory_key (Tensor) –

    The memory key to use for memory lookup.

  • logits (Tensor) –

    The original model logits.

  • inpt_tokens (list[int]) –

    The input tokens.

Returns:

  • Tensor

    The updated logits.

OrcaRankingCrossAttentionLayer #

OrcaRankingCrossAttentionLayer(
    model_dim,
    num_heads,
    num_memories,
    num_ranks,
    activation=F.relu,
    dropout=0.1,
    split_retrieval_path=False,
    projection_mode=ProjectionMode.LABEL,
    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 for rankings.

Note that this is Ranking-specific, and the rankings returned by the lookup layer are used as the value-weights for the cross attention.

The module contains the typical transformer components: multi-head attention, feed forward, and layer norm. The module also contains a lookup layer that looks up a vector in an Orca index and returns the top k results. These results are used as the memory context for the cross attention.

Parameters:

  • model_dim (int) –

    The dimension of the input vector and hidden layers.

  • num_heads (int) –

    The number of heads to be used in the multi-head attention layer.

  • num_memories (int) –

    The number of memory vectors to be returned from the lookup.

  • num_ranks (int) –

    The number of ranks to be used for the memory context.

  • activation (Callable[[Tensor], Tensor], default: relu ) –

    The activation function.

  • dropout (float, default: 0.1 ) –

    The dropout rate.

  • split_retrieval_path (bool, default: False ) –

    Whether to split the retrieval path. This is used when the memory key is different from the input vector.

  • projection_mode (ProjectionMode, default: LABEL ) –

    The mode of projection to be used.

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

    The database 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 (str | None, default: None ) –

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

lookup_result_transforms property writable #

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 #

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 #

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 #

lookup_result_override

The lookup result to use instead of performing a lookup.

lookup_database property writable #

lookup_database

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

memory_index_name property writable #

memory_index_name

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

lookup_column_names property writable #

lookup_column_names

The names of the columns to retrieve for each memory.

num_memories property writable #

num_memories

The number of memories to look up.

drop_exact_match property writable #

drop_exact_match

Whether to drop exact matches from the results.

exact_match_threshold property writable #

exact_match_threshold

The similarity threshold for exact matches.

shuffle_memories property writable #

shuffle_memories

Whether to shuffle the looked up memories.

curate_database property writable #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_effective_lookup_settings #

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 #

get_lookup_database_instance()

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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 #

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, ranks=None, memory_key=None)

Performs cross attention on the input tensor and memory context.

x, ctx, ranks act as Q, K, V for the cross attention layer.

When ctx is None:

  • If split_retrieval_path is False, x is used as both Q and K.
  • If split_retrieval_path is True, memory_key is used as K (instead of x)

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).

  • ranks (Tensor | None, default: None ) –

    The memory rank tensor of shape (batch_size, num_memories).

  • memory_key (Tensor | None, default: None ) –

    The memory key tensor of shape (batch_size, embedding_dim).

Returns:

  • Tensor

    The output tensor of shape (batch_size, embedding_dim).

OrcaRankingHead #

OrcaRankingHead(
    model_dim,
    num_memories,
    num_ranks,
    num_layers=1,
    num_heads=8,
    activation=F.relu,
    dropout=0.1,
    split_retrieval_path=False,
    projection_mode=ProjectionMode.LABEL,
    memory_guide_weight=0.0,
    single_lookup=True,
    deep_residuals=False,
    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 for ranking problems

Parameters:

  • model_dim (int) –

    The dimension of the input vector and hidden layers.

  • num_memories (int) –

    The number of memory vectors to be returned from the lookup.

  • num_ranks (int) –

    The number of ranks to be used for the memory context.

  • 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.

  • activation (Callable[[Tensor], Tensor], default: relu ) –

    The activation function.

  • dropout (float, default: 0.1 ) –

    The dropout rate.

  • split_retrieval_path (bool, default: False ) –

    Whether to split the retrieval path.

  • projection_mode (ProjectionMode, default: LABEL ) –

    The mode of projection to be used.

  • memory_guide_weight (float, default: 0.0 ) –

    The weight of the memory guide.

  • single_lookup (bool, default: True ) –

    Whether to use a single lookup.

  • deep_residuals (bool, default: False ) –

    Whether to use deep residuals.

  • 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 (str | None, default: None ) –

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

lookup_result_transforms property writable #

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 #

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 #

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 #

lookup_result_override

The lookup result to use instead of performing a lookup.

lookup_database property writable #

lookup_database

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

memory_index_name property writable #

memory_index_name

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

lookup_column_names property writable #

lookup_column_names

The names of the columns to retrieve for each memory.

num_memories property writable #

num_memories

The number of memories to look up.

drop_exact_match property writable #

drop_exact_match

Whether to drop exact matches from the results.

exact_match_threshold property writable #

exact_match_threshold

The similarity threshold for exact matches.

shuffle_memories property writable #

shuffle_memories

Whether to shuffle the looked up memories.

curate_database property writable #

curate_database

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

curate_next_run_settings property writable #

curate_next_run_settings

The settings for the next curate model run.

curate_model_id property writable #

curate_model_id

The model id to associate with curated model runs.

curate_model_version property writable #

curate_model_version

The model version to associate with curated model runs.

curate_metadata property writable #

curate_metadata

The metadata to attach to curated model runs.

curate_tags property writable #

curate_tags

The tags to attach to the curated model runs.

curate_seq_id property writable #

curate_seq_id

The sequence id to associate with curated model runs.

curate_batch_size property writable #

curate_batch_size

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

last_curate_run_ids property writable #

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 #

last_curate_run_settings

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

get_effective_lookup_settings #

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 #

get_lookup_database_instance()

Returns the OrcaDatabase instance to use for looking up memories.

get_orca_modules_recursively #

1
2
3
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 #

1
2
3
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 #

1
2
3
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 #

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, ctx_ranks=None, memory_key=None)

Performs cross attention on the input tensor 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 is looked up based on the memory_key or input tensor.

  • ctx_ranks (Tensor | None, default: None ) –

    The memory rank tensor of shape (batch_size, num_memories). If None, the ranks are looked up along with the memory context.

  • memory_key (Tensor | None, default: None ) –

    The memory key tensor of shape (batch_size, embedding_dim). If None, the memory key is the input tensor.

Returns:

  • Tensor

    The output tensor of shape (batch_size, 1). The output is the rank of the input vector.

OrcaMemoryDataset #

OrcaMemoryDataset(
    db,
    index,
    columns,
    memory_index,
    mem_columns,
    num_memories,
    *,
    page_size=1000,
    verbose=False,
    drop_exact_match=False,
    exact_match_threshold=EXACT_MATCH_THRESHOLD,
    shuffle_memories=False,
    cache_file_path=None
)

Bases: Dataset

A PyTorch Dataset that allows loading all data needed in a model run from the Orca database upfront which is more efficient. The dataset is generated by joining an index that contains the embeddings for the training data with an index that contains the embeddings for the memories (those might be the same). The join operation happens in large pages which speeds up the calculation over doing individual lookups for each batch.

The dataset returns a tuple that contains:

  • The embedding for the training data: a Tensor of shape (embedding_dim, 1)
  • Any other columns that are specified to be retrieved from the training data table
  • The embeddings for the memories: a Tensor of shape (num_memories, embedding_dim)
  • Any other columns that are specified to be retrieved from the memory data table
  • The scores for the memories: a Tensor of shape (num_memories, 1)

Examples:

import torch
from torch.utils.data import DataLoader
from orcalib.orca_torch import OrcaMemoryDataset

dataset = OrcaMemoryDataset(
    db,
    index="train_index",
    columns="label",
    memory_index="memories_index",
    mem_columns="label",
    num_memories=10,
)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# training loop assuming we have a `model`
optimizer = torch.optim.Adam(model.parameters())
for x, x_labels, ctx, ctx_labels, ctx_scores in dataloader:
    optimizer.zero_grad()
    logits = model(x, ctx, torch.stack(ctx_labels).T, ctx_scores)
    loss = torch.nn.functional.cross_entropy(logits, x_labels)
    loss.backward()
    optimizer.step()

Parameters:

  • db (OrcaDatabase) –

    The OrcaDatabase to fetch the index data from.

  • index (str) –

    The name of the index to fetch the data from.

  • columns (list[str | ColumnHandle] | str | ColumnHandle) –

    The columns to fetch from the index. Can be a single column name, or a list of column names.

  • memory_index (str) –

    The name of the memory index to fetch the data from. (Generally the same as the index)

  • mem_columns (list[str | ColumnHandle] | str | ColumnHandle) –

    The columns to fetch from the memory index. Can be a single column name, or a list of column names.

  • num_memories (int) –

    The number of memory vectors to fetch for each item vector.

  • page_size (int, default: 1000 ) –

    The page size to use when fetching the data from the database.

  • verbose (bool, default: False ) –

    Whether to print verbose logging information.

  • drop_exact_match (bool, default: False ) –

    Whether to drop the exact match (if found) always, never, or only during training or inference.

  • exact_match_threshold (float, default: EXACT_MATCH_THRESHOLD ) –

    The minimum similarity score for something to be considered the exact match.

  • shuffle_memories (bool, default: False ) –

    Whether to shuffle the memories before returning them.

  • cache_file_path (str | None, default: None ) –

    The path for the pickle file to use for caching the dataset. If None, caching is disabled.

__getitem__ #

__getitem__(i)

Get the item at the given location i in the dataset.

If the item is on a page that has not been loaded yet, the page is fetched from the database.

Parameters:

  • i (int) –

    The index of the item to get.

Returns:

  • x ( Tensor ) –

    The embedding for the item of shape (embedding_dim, 1).

  • x_columns ( list[Any] | Any ) –

    The value for the column or columns for the item that were fetched.

  • ctx ( Tensor ) –

    The memory embeddings of shape (num_memories, embedding_dim).

  • ctx_columns ( list[list[Any] | Any] ) –

    The values for the column or columns for the memory embeddings that were fetched.

  • ctx_scores ( Tensor ) –

    The scores between the memory embeddings and the item embedding of shape (num_memories).

__len__ #

__len__()

Returns the length of the dataset

get_dict #

get_dict(i)

Returns the dictionary for the items at the given location

Parameters:

  • i (int) –

    The location of the item to get.

Returns:

  • dict

    A dictionary containing the item and memory data.

set_cache_file_path #

set_cache_file_path(path)

Set the cache file path for the dataset.

Parameters:

  • path (str) –

    The path where the cache file will be stored.

save #

save()

Save the dataset to a file.

load #

load(path)

Load the dataset from a file.

Parameters:

  • path (str) –

    The path to the pickle file on disk to load the dataset from.

OrcaTextClassificationTrainer #

OrcaTextClassificationTrainer(
    model,
    tokenizer,
    trainloader,
    testloader,
    use_memory=True,
    memory_dropout=0.0,
    device_override=None,
    param_groups=None,
    verbosity=0,
)

A simple trainer class for Text Classification Problems with Orca. Intended for quick prototyping, not to outperform a custom training loop.

Deprecated:

This class is not the recommended way to train models anymore.

Parameters:

  • model (OrcaModel) –

    The model to train.

  • tokenizer (PreTrainedTokenizerBase) –

    The tokenizer to use for encoding the input data.

  • trainloader (DataLoader) –

    The DataLoader for the training data.

  • testloader (DataLoader) –

    The DataLoader for the test data.

  • use_memory (bool, default: True ) –

    Whether to use memory for the model.

  • memory_dropout (float, default: 0.0 ) –

    The dropout rate to use for the memory.

  • device_override (str | None, default: None ) –

    The device to use for training. If None, the device will be automatically selected based on the availability of CUDA.

  • param_groups (None | list[dict[str, Any]], default: None ) –

    The parameter groups to use for training. If None, the model parameters will be used.

  • verbosity (int, default: 0 ) –

    The verbosity level to use for training.

get_test_accuracy #

get_test_accuracy(testloader_override=None)

Computes and returns the average accuracy of the model either on the main testset (from the constructor) or on the provided testloader_override.

Parameters:

  • testloader_override (DataLoader | None, default: None ) –

    The DataLoader to use for the testset. If None, the main testset will be used.

Returns:

  • float

    The average accuracy of the model on the testset.

train #

train(num_epochs=10)

Trains the model for the given number of epochs.

Parameters:

  • num_epochs (int, default: 10 ) –

    The number of epochs to train for.

embed #

embed(strings, model=EmbeddingModel.SENTENCE_TRANSFORMER)

Generate embeddings for a list of strings.

Parameters:

  • strings (list[str]) –

    The list of strings to generate embeddings for.

  • model (EmbeddingModel, default: SENTENCE_TRANSFORMER ) –

    The embedding model to use.

Returns:

  • Tensor

    the embeddings of shape (len(strings), embedding_dim)

Examples:

>>> embed(["hello", "world"]).shape
torch.Size([2, 768])