Autofaiss Index selection demo slider

[1]:
# Install autofaiss
!pip install autofaiss &> /dev/null
[4]:
## Import automatic index selection function
from autofaiss.external.optimize import get_optimal_index_keys_v2
from autofaiss.external.metadata import IndexMetadata
[7]:
#@title Demo
import ipywidgets as widgets
from IPython.display import display
from ipywidgets import HBox, Label

continuous_update = True

output = widgets.Output()
nb_vectors = widgets.FloatLogSlider(value = 10000, min=1, max=9, description="", readout_format=".0f", continuous_update=continuous_update)
dim_vectors = widgets.FloatLogSlider(base=2, value = 128, min=0, max=12, description="", readout_format=".0f", continuous_update=continuous_update)
max_index_memory_usage = widgets.FloatLogSlider(value = 100, min=0, max=5, description="", readout_format=".0f", continuous_update=continuous_update)

def get_text():
    _nb_vectors = round(nb_vectors.value)
    _dim_vectors = round(dim_vectors.value)
    m = round(max_index_memory_usage.value)
    _max_index_memory_usage = f"{m}MB"# if m < 1000 else f"{m/1000}GB"

    s = ""
    #s += f"nb_vectors={_nb_vectors}, dim_vectors={_dim_vectors}, max_index_memory_usage={_max_index_memory_usage}\n"
    opti_keys = get_optimal_index_keys_v2(_nb_vectors, _dim_vectors, _max_index_memory_usage)

    if not opti_keys:
        s += 'Impossible to find index parameters for such a large dataset and low memory'
    else:
        index_metadada = IndexMetadata(opti_keys[0], _nb_vectors, _dim_vectors)

        s += f"The optimal index would be: {opti_keys[0]}"
        s+= "\n\n"
        s+= index_metadada.get_index_description(tunable_parameters_infos=True)

    return s

text = widgets.Textarea(
    value=str(get_text()),
    placeholder='...',
    description='Result:',
    disabled=False,
    rows = 50,
    continuous_update=False,
    layout={'width': '100%'}
)

def update(b):
  # Display the message within the output widget.
  with output:
    text.value = get_text()

nb_vectors.observe(update)
dim_vectors.observe(update)
max_index_memory_usage.observe(update)

display(HBox([Label('Number of vectors in the dataset'), nb_vectors]))
display(HBox([Label('Dimension of vectors in the dataset'), dim_vectors]))
display(HBox([Label('Max size of index in MB'), max_index_memory_usage]))
display(output)
display(text)