vectordb-bench 0.0.11__py3-none-any.whl → 0.0.13__py3-none-any.whl

vectordb_bench/frontend/components/custom/displayCustomCase.py

@@ -0,0 +1,31 @@
+
+from vectordb_bench.frontend.components.custom.getCustomConfig import CustomCaseConfig
+
+
+def displayCustomCase(customCase: CustomCaseConfig, st, key):
+
+    columns = st.columns([1, 2])
+    customCase.dataset_config.name = columns[0].text_input(
+        "Name", key=f"{key}_name", value=customCase.dataset_config.name)
+    customCase.name = f"{customCase.dataset_config.name} (Performace Case)"
+    customCase.dataset_config.dir = columns[1].text_input(
+        "Folder Path", key=f"{key}_dir", value=customCase.dataset_config.dir)
+
+    columns = st.columns(4)
+    customCase.dataset_config.dim = columns[0].number_input(
+        "dim", key=f"{key}_dim", value=customCase.dataset_config.dim)
+    customCase.dataset_config.size = columns[1].number_input(
+        "size", key=f"{key}_size", value=customCase.dataset_config.size)
+    customCase.dataset_config.metric_type = columns[2].selectbox(
+        "metric type", key=f"{key}_metric_type", options=["L2", "Cosine", "IP"])
+    customCase.dataset_config.file_count = columns[3].number_input(
+        "train file count", key=f"{key}_file_count", value=customCase.dataset_config.file_count)
+
+    columns = st.columns(4)
+    customCase.dataset_config.use_shuffled = columns[0].checkbox(
+        "use shuffled data",  key=f"{key}_use_shuffled", value=customCase.dataset_config.use_shuffled)
+    customCase.dataset_config.with_gt = columns[1].checkbox(
+        "with groundtruth",  key=f"{key}_with_gt", value=customCase.dataset_config.with_gt)
+
+    customCase.description = st.text_area(
+        "description", key=f"{key}_description", value=customCase.description)

vectordb_bench/frontend/components/custom/displaypPrams.py

@@ -0,0 +1,11 @@
+def displayParams(st):
+    st.markdown("""
+- `Folder Path` - The path to the folder containing all the files. Please ensure that all files in the folder are in the `Parquet` format.
+  - Vectors data files: The file must be named `train.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`.
+  - Query test vectors: The file must be named `test.parquet` and should have two columns: `id` as an incrementing `int` and `emb` as an array of `float32`.
+  - Ground truth file: The file must be named `neighbors.parquet` and should have two columns: `id` corresponding to query vectors and `neighbors_id` as an array of `int`.
+
+- `Train File Count` - If the vector file is too large, you can consider splitting it into multiple files. The naming format for the split files should be `train-[index]-of-[file_count].parquet`. For example, `train-01-of-10.parquet` represents the second file (0-indexed) among 10 split files.
+
+- `Use Shuffled Data` - If you check this option, the vector data files need to be modified. VectorDBBench will load the data labeled with `shuffle`. For example, use `shuffle_train.parquet` instead of `train.parquet` and `shuffle_train-04-of-10.parquet` instead of `train-04-of-10.parquet`. The `id` column in the shuffled data can be in any order.
+""")

vectordb_bench/frontend/components/custom/getCustomConfig.py

@@ -0,0 +1,40 @@
+import json
+
+from pydantic import BaseModel
+
+from vectordb_bench import config
+
+
+class CustomDatasetConfig(BaseModel):
+    name: str = "custom_dataset"
+    dir: str = ""
+    size: int = 0
+    dim: int = 0
+    metric_type: str = "L2"
+    file_count: int = 1
+    use_shuffled: bool = False
+    with_gt: bool = True
+
+
+class CustomCaseConfig(BaseModel):
+    name: str = "custom_dataset (Performace Case)"
+    description: str = ""
+    load_timeout: int = 36000
+    optimize_timeout: int = 36000
+    dataset_config: CustomDatasetConfig = CustomDatasetConfig()
+
+
+def get_custom_configs():
+    with open(config.CUSTOM_CONFIG_DIR, "r") as f:
+        custom_configs = json.load(f)
+        return [CustomCaseConfig(**custom_config) for custom_config in custom_configs]
+
+
+def save_custom_configs(custom_configs: list[CustomDatasetConfig]):
+    with open(config.CUSTOM_CONFIG_DIR, "w") as f:
+        json.dump([custom_config.dict()
+                   for custom_config in custom_configs], f, indent=4)
+
+
+def generate_custom_case():
+    return CustomCaseConfig()

vectordb_bench/frontend/components/custom/initStyle.py

@@ -0,0 +1,15 @@
+def initStyle(st):
+    st.markdown(
+        """<style>
+            /* expander - header */
+            .main div[data-testid='stExpander'] summary p {font-size: 20px; font-weight: 600;}
+            /* 
+            button {
+                height: auto;
+                padding-left: 8px !important;
+                padding-right: 6px !important;
+            }
+            */
+        </style>""",
+        unsafe_allow_html=True,
+    )

vectordb_bench/frontend/components/run_test/autoRefresh.py

@@ -1,5 +1,5 @@
 from streamlit_autorefresh import st_autorefresh
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.frontend.config.styles import *
 
 
 def autoRefresh():

vectordb_bench/frontend/components/run_test/caseSelector.py

@@ -1,9 +1,13 @@
-from vectordb_bench.frontend.const.styles import *
+
+from vectordb_bench.frontend.config.styles import *
 from vectordb_bench.backend.cases import CaseType
-from vectordb_bench.frontend.const.dbCaseConfigs import *
+from vectordb_bench.frontend.config.dbCaseConfigs import *
+from collections import defaultdict
+
+from vectordb_bench.frontend.utils import addHorizontalLine
 
 
-def caseSelector(st, activedDbList):
+def caseSelector(st, activedDbList: list[DB]):
     st.markdown(
         "<div style='height: 24px;'></div>",
         unsafe_allow_html=True,
@@ -14,41 +18,49 @@ def caseSelector(st, activedDbList):
         unsafe_allow_html=True,
     )
 
-    caseIsActived = {case: False for case in CASE_LIST}
-    allCaseConfigs = {db: {case: {} for case in CASE_LIST} for db in DB_LIST}
-    for caseOrDivider in CASE_LIST_WITH_DIVIDER:
-        if caseOrDivider == DIVIDER:
-            caseItemContainer.markdown(
-                "<div style='border: 1px solid #cccccc60; margin-bottom: 24px;'></div>",
-                unsafe_allow_html=True,
-            )
+    activedCaseList: list[CaseConfig] = []
+    dbToCaseClusterConfigs = defaultdict(lambda: defaultdict(dict))
+    dbToCaseConfigs = defaultdict(lambda: defaultdict(dict))
+    caseClusters = UI_CASE_CLUSTERS + [get_custom_case_cluter()]
+    for caseCluster in caseClusters:
+        activedCaseList += caseClusterExpander(
+            st, caseCluster, dbToCaseClusterConfigs, activedDbList)
+    for db in dbToCaseClusterConfigs:
+        for uiCaseItem in dbToCaseClusterConfigs[db]:
+            for case in uiCaseItem.cases:
+                dbToCaseConfigs[db][case] = dbToCaseClusterConfigs[db][uiCaseItem]
+
+    return activedCaseList, dbToCaseConfigs
+
+
+def caseClusterExpander(st, caseCluster: UICaseItemCluster, dbToCaseClusterConfigs, activedDbList: list[DB]):
+    expander = st.expander(caseCluster.label, False)
+    activedCases: list[CaseConfig] = []
+    for uiCaseItem in caseCluster.uiCaseItems:
+        if uiCaseItem.isLine:
+            addHorizontalLine(expander)
         else:
-            case = caseOrDivider
-            caseItemContainer = st.container()
-            caseIsActived[case] = caseItem(
-                caseItemContainer, allCaseConfigs, case, activedDbList
-            )
-    activedCaseList = [case for case in CASE_LIST if caseIsActived[case]]
-    return activedCaseList, allCaseConfigs
+            activedCases += caseItemCheckbox(expander,
+                                             dbToCaseClusterConfigs, uiCaseItem, activedDbList)
+    return activedCases
 
 
-def caseItem(st, allCaseConfigs, case: CaseType, activedDbList):
-    selected = st.checkbox(case.case_name)
+def caseItemCheckbox(st, dbToCaseClusterConfigs, uiCaseItem: UICaseItem, activedDbList: list[DB]):
+    selected = st.checkbox(uiCaseItem.label)
     st.markdown(
-        f"<div style='color: #1D2939; margin: -8px 0 20px {CHECKBOX_INDENT}px; font-size: 14px;'>{case.case_description}</div>",
+        f"<div style='color: #1D2939; margin: -8px 0 20px {CHECKBOX_INDENT}px; font-size: 14px;'>{uiCaseItem.description}</div>",
         unsafe_allow_html=True,
     )
 
     if selected:
-        caseConfigSettingContainer = st.container()
         caseConfigSetting(
-            caseConfigSettingContainer, allCaseConfigs, case, activedDbList
+            st.container(), dbToCaseClusterConfigs, uiCaseItem, activedDbList
         )
 
-    return selected
+    return uiCaseItem.cases if selected else []
 
 
-def caseConfigSetting(st, allCaseConfigs, case, activedDbList):
+def caseConfigSetting(st, dbToCaseClusterConfigs, uiCaseItem: UICaseItem, activedDbList: list[DB]):
     for db in activedDbList:
         columns = st.columns(1 + CASE_CONFIG_SETTING_COLUMNS)
         # column 0 - title
@@ -57,12 +69,12 @@ def caseConfigSetting(st, allCaseConfigs, case, activedDbList):
             f"<div style='margin: 0 0 24px {CHECKBOX_INDENT}px; font-size: 18px; font-weight: 600;'>{db.name}</div>",
             unsafe_allow_html=True,
         )
-        caseConfig = allCaseConfigs[db][case]
         k = 0
-        for config in CASE_CONFIG_MAP.get(db, {}).get(case.case_cls().label, []):
+        caseConfig = dbToCaseClusterConfigs[db][uiCaseItem]
+        for config in CASE_CONFIG_MAP.get(db, {}).get(uiCaseItem.caseLabel, []):
             if config.isDisplayed(caseConfig):
                 column = columns[1 + k % CASE_CONFIG_SETTING_COLUMNS]
-                key = "%s-%s-%s" % (db, case, config.label.value)
+                key = "%s-%s-%s" % (db, uiCaseItem.label, config.label.value)
                 if config.inputType == InputType.Text:
                     caseConfig[config.label] = column.text_input(
                         config.displayLabel if config.displayLabel else config.label.value,
@@ -88,6 +100,16 @@ def caseConfigSetting(st, allCaseConfigs, case, activedDbList):
                         value=config.inputConfig["value"],
                         help=config.inputHelp,
                     )
+                elif config.inputType == InputType.Float:
+                    caseConfig[config.label] = column.number_input(
+                        config.displayLabel if config.displayLabel else config.label.value,
+                        step=config.inputConfig.get("step", 0.1),
+                        min_value=config.inputConfig["min"],
+                        max_value=config.inputConfig["max"],
+                        key=key,
+                        value=config.inputConfig["value"],
+                        help=config.inputHelp,
+                    )
                 k += 1
         if k == 0:
             columns[1].write("Auto")

vectordb_bench/frontend/components/run_test/dbConfigSetting.py

@@ -1,13 +1,10 @@
 from pydantic import ValidationError
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.backend.clients import DB
+from vectordb_bench.frontend.config.styles import DB_CONFIG_SETTING_COLUMNS
 from vectordb_bench.frontend.utils import inputIsPassword
 
 
-def dbConfigSettings(st, activedDbList):
-    st.markdown(
-        "<style> .streamlit-expanderHeader p {font-size: 20px; font-weight: 600;}</style>",
-        unsafe_allow_html=True,
-    )
+def dbConfigSettings(st, activedDbList: list[DB]):
     expander = st.expander("Configurations for the selected databases", True)
 
     dbConfigs = {}
@@ -31,7 +28,7 @@ def dbConfigSettings(st, activedDbList):
     return dbConfigs, isAllValid
 
 
-def dbConfigSettingItem(st, activeDb):
+def dbConfigSettingItem(st, activeDb: DB):
     st.markdown(
         f"<div style='font-weight: 600; font-size: 20px; margin-top: 16px;'>{activeDb.value}</div>",
         unsafe_allow_html=True,
@@ -40,20 +37,41 @@ def dbConfigSettingItem(st, activeDb):
 
     dbConfigClass = activeDb.config_cls
     properties = dbConfigClass.schema().get("properties")
-    propertiesItems = list(properties.items())
-    moveDBLabelToLast(propertiesItems)
     dbConfig = {}
-    for j, property in enumerate(propertiesItems):
-        column = columns[j % DB_CONFIG_SETTING_COLUMNS]
-        key, value = property
+    idx = 0
+
+    # db config (unique)
+    for key, property in properties.items():
+        if (
+            key not in dbConfigClass.common_short_configs()
+            and key not in dbConfigClass.common_long_configs()
+        ):
+            column = columns[idx % DB_CONFIG_SETTING_COLUMNS]
+            idx += 1
+            dbConfig[key] = column.text_input(
+                key,
+                key="%s-%s" % (activeDb.name, key),
+                value=property.get("default", ""),
+                type="password" if inputIsPassword(key) else "default",
+            )
+    # db config (common short labels)
+    for key in dbConfigClass.common_short_configs():
+        column = columns[idx % DB_CONFIG_SETTING_COLUMNS]
+        idx += 1
         dbConfig[key] = column.text_input(
             key,
-            key="%s-%s" % (activeDb, key),
-            value=value.get("default", ""),
-            type="password" if inputIsPassword(key) else "default",
+            key="%s-%s" % (activeDb.name, key),
+            value="",
+            type="default",
+            placeholder="optional, for labeling results",
         )
-    return dbConfig
 
-
-def moveDBLabelToLast(propertiesItems):
-    propertiesItems.sort(key=lambda x: 1 if x[0] == "db_label" else 0)
+    # db config (common long text_input)
+    for key in dbConfigClass.common_long_configs():
+        dbConfig[key] = st.text_area(
+            key,
+            key="%s-%s" % (activeDb.name, key),
+            value="",
+            placeholder="optional",
+        )
+    return dbConfig

vectordb_bench/frontend/components/run_test/dbSelector.py

@@ -1,7 +1,6 @@
 from streamlit.runtime.media_file_storage import MediaFileStorageError
-
-from vectordb_bench.frontend.const.styles import *
-from vectordb_bench.frontend.const.dbCaseConfigs import DB_LIST
+from vectordb_bench.frontend.config.styles import DB_SELECTOR_COLUMNS, DB_TO_ICON
+from vectordb_bench.frontend.config.dbCaseConfigs import DB_LIST
 
 
 def dbSelector(st):
@@ -18,17 +17,6 @@ def dbSelector(st):
     dbContainerColumns = st.columns(DB_SELECTOR_COLUMNS, gap="small")
     dbIsActived = {db: False for db in DB_LIST}
 
-    # style - image; column gap; checkbox font;
-    st.markdown(
-        """
-        <style>
-            div[data-testid='stImage'] {margin: auto;}
-            div[data-testid='stHorizontalBlock'] {gap: 8px;}
-            .stCheckbox p { color: #000; font-size: 18px; font-weight: 600; }
-        </style>
-        """,
-        unsafe_allow_html=True,
-    )
     for i, db in enumerate(DB_LIST):
         column = dbContainerColumns[i % DB_SELECTOR_COLUMNS]
         dbIsActived[db] = column.checkbox(db.name)

vectordb_bench/frontend/components/run_test/generateTasks.py

@@ -1,17 +1,15 @@
+from vectordb_bench.backend.clients import DB
 from vectordb_bench.models import CaseConfig, CaseConfigParamType, TaskConfig
 
 
-def generate_tasks(activedDbList, dbConfigs, activedCaseList, allCaseConfigs):
+def generate_tasks(activedDbList: list[DB], dbConfigs, activedCaseList: list[CaseConfig], allCaseConfigs):
     tasks = []
     for db in activedDbList:
         for case in activedCaseList:
             task = TaskConfig(
                     db=db.value,
                     db_config=dbConfigs[db],
-                    case_config=CaseConfig(
-                        case_id=case.value,
-                        custom_case={},
-                    ),
+                    case_config=case,
                     db_case_config=db.case_config_cls(
                         allCaseConfigs[db][case].get(CaseConfigParamType.IndexType, None)
                     )(**{key.value: value for key, value in allCaseConfigs[db][case].items()}),

vectordb_bench/frontend/components/run_test/initStyle.py

@@ -0,0 +1,16 @@
+def initStyle(st):
+    st.markdown(
+        """<style>
+            /* expander - header */
+            .main div[data-testid='stExpander'] p {font-size: 18px; font-weight: 600;}
+            /* db icon */
+            div[data-testid='stImage'] {margin: auto;}
+            /* db column gap */
+            div[data-testid='stHorizontalBlock'] {gap: 8px;}
+            /* check box */
+            .stCheckbox p { color: #000; font-size: 18px; font-weight: 600; }
+            /* db selector - db_name should not wrap */
+            div[data-testid="stVerticalBlockBorderWrapper"] div[data-testid="stCheckbox"] div[data-testid="stWidgetLabel"] p { white-space: nowrap; }
+        </style>""",
+        unsafe_allow_html=True,
+    )

vectordb_bench/frontend/components/run_test/submitTask.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from vectordb_bench.frontend.const.styles import *
+from vectordb_bench.frontend.config.styles import *
 from vectordb_bench.interface import benchMarkRunner