cost_fn fixes and clarity (#22)

* cost_fn fixes and clarity

Author: rbs333

.github/workflows/test.yml

@@ -0,0 +1,57 @@
+name: Tests
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    services:
+      redis:
+        image: redis/redis-stack:latest
+        ports:
+          - 6379:6379
+        options: >-
+          --health-cmd "redis-cli ping"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install Poetry
+      uses: snok/install-poetry@v1
+      with:
+        version: latest
+        virtualenvs-create: true
+        virtualenvs-in-project: true
+
+    - name: Load cached venv
+      id: cached-poetry-dependencies
+      uses: actions/cache@v3
+      with:
+        path: .venv
+        key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock') }}
+
+    - name: Install dependencies
+      if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
+      run: poetry install --all-extras
+
+    - name: Run tests
+      run: poetry run test
+      env:
+        REDIS_URL: redis://localhost:6379/0

README.md

@@ -136,7 +136,7 @@ metrics = run_grid_study(
 ```
 
 #### Example output
-| search_method | model                                      | avg_query_time | recall@k | precision | ndcg@k   |
+| search_method | model                                      | avg_query_time | recall | precision | ndcg  |
 |----------------|---------------------------------------------|----------------|-----------|-----------|----------|
 | weighted_rrf   | sentence-transformers/all-MiniLM-L6-v2     | 0.006608       | 0.156129  | 0.261056  | 0.204241 |
 | rerank         | sentence-transformers/all-MiniLM-L6-v2     | 0.127574       | 0.156039  | 0.260437  | 0.190298 |
@@ -169,7 +169,7 @@ index_settings:
 optimization_settings:
   # defines weight of each metric in optimization function
   metric_weights:
-    f1_at_k: 1
+    f1: 1
     total_indexing_time: 1
   algorithms: ["hnsw"] # indexing algorithm to be included in the study
   vector_data_types: ["float16", "float32"] # data types to be included in the study
@@ -214,7 +214,7 @@ metrics = run_bayes_study(
 ```
 
 #### Example output
-| search_method | algorithm | vector_data_type | ef_construction | ef_runtime | m  | avg_query_time | total_indexing_time | f1@k    |
+| search_method | algorithm | vector_data_type | ef_construction | ef_runtime | m  | avg_query_time | total_indexing_time | f1    |
 |---------------|-----------|------------------|------------------|------------|----|----------------|----------------------|---------|
 | hybrid        | hnsw      | float16          | 200              | 50         | 8  | 0.004628       | 3.559                | 0.130712|
 | hybrid        | hnsw      | float16          | 200              | 50         | 64 | 0.004498       | 4.804                | 0.130712|
@@ -296,7 +296,7 @@ cache = SemanticCache(
 
 # Add some data to the cache
 paris_key = cache.store(
-    prompt="what is the capital of france?", 
+    prompt="what is the capital of france?",
     response="paris"
 )
 
@@ -307,7 +307,7 @@ test_data = [
         "query_match": paris_key  # Expected cache hit
     },
     {
-        "query": "What's the capital of Britain?", 
+        "query": "What's the capital of Britain?",
         "query_match": ""  # Expected cache miss
     }
 ]
@@ -337,7 +337,7 @@ routes = [
         distance_threshold=0.5,
     ),
     Route(
-        name="farewell", 
+        name="farewell",
         references=["bye", "goodbye"],
         metadata={"type": "farewell"},
         distance_threshold=0.5,
@@ -509,7 +509,7 @@ metrics = run_grid_study(
 
 ### Example output
 
-| search_method     | model                                      | avg_query_time | recall@k | precision | ndcg@k   |
+| search_method     | model                                      | avg_query_time | recall | precision | ndcg   |
 |-------------------|---------------------------------------------|----------------|-----------|-----------|----------|
 | pre_filter_vector | sentence-transformers/all-MiniLM-L6-v2     | 0.001177       | 1.0       | 0.25      | 0.914903 |
 | basic_vector      | sentence-transformers/all-MiniLM-L6-v2     | 0.002605       | 0.9       | 0.23      | 0.717676 |

docs/examples/bayesian_optimization/00_bayes_study.ipynb

@@ -17,8 +17,12 @@ index_settings:
 optimization_settings:
   # defines weight of each metric in optimization function
   metric_weights:
-    f1_at_k: 1
-    total_indexing_time: 1
+    f1: 2
+    total_indexing_time: 2
+    avg_query_time: 2
+    recall: 2
+    ndcg: 2
+    precision: 2
   algorithms: ["hnsw"] # indexing algorithm to be included in the study
   vector_data_types: ["float16", "float32"] # data types to be included in the study
   distance_metrics: ["cosine"] # distance metrics to be included in the study

docs/examples/bayesian_optimization/bayes_study_config.yaml

@@ -3641,7 +3641,7 @@
     }
    ],
    "source": [
-    "metrics[[\"search_method\", \"model\", \"model_dim\", 'total_indexing_time', \"avg_query_time\", \"recall@k\", \"precision\", \"ndcg@k\"]].sort_values(by=\"ndcg@k\", ascending=False)"
+    "metrics[[\"search_method\", \"model\", \"model_dim\", 'total_indexing_time', \"avg_query_time\", \"recall\", \"precision\", \"ndcg\"]].sort_values(by=\"ndcg\", ascending=False)"
    ]
   },
   {

docs/examples/comparison/00_comparison.ipynb

@@ -1501,7 +1501,7 @@
     }
    ],
    "source": [
-    "metrics[[\"search_method\", \"model\", \"avg_query_time\", \"recall@k\", \"precision\", \"ndcg@k\"]].sort_values(by=\"ndcg@k\", ascending=False)"
+    "metrics[[\"search_method\", \"model\", \"avg_query_time\", \"recall\", \"precision\", \"ndcg\"]].sort_values(by=\"ndcg\", ascending=False)"
    ]
   }
  ],

docs/examples/grid_study/00_grid_study.ipynb

@@ -562,7 +562,7 @@
     }
    ],
    "source": [
-    "metrics[[\"search_method\", \"model\", \"avg_query_time\", \"recall@k\", \"precision\", \"ndcg@k\"]].sort_values(by=\"ndcg@k\", ascending=False)"
+    "metrics[[\"search_method\", \"model\", \"avg_query_time\", \"recall\", \"precision\", \"ndcg\"]].sort_values(by=\"ndcg\", ascending=False)"
    ]
   }
  ],

docs/examples/grid_study/01_custom_grid_study.ipynb

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "redis-retrieval-optimizer"
-version = "0.2.0"
+version = "0.2.1"
 description = "A tool to help optimize information retrieval with the Redis Query Engine."
 authors = [ "Robert Shelton <[email protected]>" ]
 license = "MIT"

pyproject.toml

@@ -28,16 +28,17 @@
     "model": [],
     "model_dim": [],
     "ret_k": [],
-    "recall@k": [],
-    "ndcg@k": [],
-    "f1@k": [],
+    "recall": [],
+    "ndcg": [],
+    "f1": [],
     "precision": [],
     "algorithm": [],
     "ef_construction": [],
     "ef_runtime": [],
     "m": [],
     "distance_metric": [],
     "vector_data_type": [],
+    "objective_value": [],
 }
 
 
@@ -52,12 +53,13 @@ def update_metric_row(trial_settings: TrialSettings, trial_metrics: dict):
     METRICS["vector_data_type"].append(trial_settings.index_settings.vector_data_type)
     METRICS["model"].append(trial_settings.embedding.model)
     METRICS["model_dim"].append(trial_settings.embedding.dim)
-    METRICS["recall@k"].append(trial_metrics["recall"])
-    METRICS["ndcg@k"].append(trial_metrics["ndcg"])
+    METRICS["recall"].append(trial_metrics["recall"])
+    METRICS["ndcg"].append(trial_metrics["ndcg"])
     METRICS["precision"].append(trial_metrics["precision"])
-    METRICS["f1@k"].append(trial_metrics["f1"])
+    METRICS["f1"].append(trial_metrics["f1"])
     METRICS["total_indexing_time"].append(trial_metrics["total_indexing_time"])
     METRICS["avg_query_time"].append(trial_metrics["avg_query_time"])
+    METRICS["objective_value"].append(trial_metrics["objective_value"])
 
 
 def persist_metrics(
@@ -70,17 +72,30 @@ def persist_metrics(
     client.json().set(f"study:{study_id}", Path.root_path(), METRICS)
 
 
+def norm_metric(value: float):
+    """Normalize a metric value using 1/(1+value) formula.
+
+    Handles edge cases:
+    - When value is -1, returns a large positive number (infinity equivalent)
+    - When value is very negative, returns a large positive number
+    - When value is very positive, returns a small positive number
+    """
+    if value == -1:
+        # Return a large positive number to represent "infinity" for optimization
+        return 1000.0
+    return 1 / (1 + value)
+
+
 def cost_fn(metrics: dict, weights: dict):
     objective = 0
     for key in metrics:
-        objective += weights.get(key, 0) * metrics[key]
+        if key == "avg_query_time" or key == "total_indexing_time":
+            objective += weights.get(key, 0) * -norm_metric(metrics[key])
+        else:
+            objective += weights.get(key, 0) * metrics[key]
     return objective
 
 
-def norm_metric(value: float):
-    return 1 / (1 + value)
-
-
 def objective(trial, study_config, redis_url, corpus_processor, search_method_map):
 
     # optimizer will select hyperparameters from available option in study_config
@@ -152,19 +167,19 @@ def objective(trial, study_config, redis_url, corpus_processor, search_method_ma
     search_method_output = search_fn(search_input)
 
     trial_metrics = utils.eval_trial_metrics(qrels, search_method_output.run)
-    trial_metrics["total_indexing_time"] = -(total_indexing_time)
-    trial_metrics["avg_query_time"] = -(
-        utils.get_query_time_stats(search_method_output.query_metrics.query_times)[
-            "avg_query_time"
-        ]
+    trial_metrics["total_indexing_time"] = total_indexing_time
+    trial_metrics["avg_query_time"] = utils.get_query_time_stats(
+        search_method_output.query_metrics.query_times
+    )["avg_query_time"]
+
+    trial_metrics["objective_value"] = cost_fn(
+        trial_metrics, study_config.optimization_settings.metric_weights.model_dump()
     )
 
     # save results as we go in case of failure
     persist_metrics(redis_url, trial_settings, trial_metrics, study_config.study_id)
 
-    return cost_fn(
-        trial_metrics, study_config.optimization_settings.metric_weights.model_dump()
-    )
+    return trial_metrics["objective_value"]
 
 
 def run_bayes_study(

redis_retrieval_optimizer/bayes_study.py

@@ -29,10 +29,10 @@ def update_metric_row(
     )
     metrics["model"].append(embedding_settings.model)
     metrics["model_dim"].append(embedding_settings.dim)
-    metrics["recall@k"].append(trial_metrics["recall"])
-    metrics["ndcg@k"].append(trial_metrics["ndcg"])
+    metrics["recall"].append(trial_metrics["recall"])
+    metrics["ndcg"].append(trial_metrics["ndcg"])
     metrics["precision"].append(trial_metrics["precision"])
-    metrics["f1@k"].append(trial_metrics["f1"])
+    metrics["f1"].append(trial_metrics["f1"])
     metrics["total_indexing_time"].append(trial_metrics["total_indexing_time"])
     metrics["avg_query_time"].append(trial_metrics["query_stats"]["avg_query_time"])
     return metrics
@@ -125,9 +125,9 @@ def run_grid_study(
         "search_method": [],
         "total_indexing_time": [],
         "avg_query_time": [],
-        "recall@k": [],
-        "ndcg@k": [],
-        "f1@k": [],
+        "recall": [],
+        "ndcg": [],
+        "f1": [],
         "precision": [],
         "ret_k": [],
         "algorithm": [],