DAGE-107: Add a script to download bbc news dataset and support batch indexing in docker-services

rre-tools/README.md

@@ -12,12 +12,12 @@ documents from a search engine, generate synthetic queries, and score the releva
 
 ### [Embedding Model Evaluator](docs/embedding_model_evaluator/README.md)
 
-This tool provide a flexible tool to test a HuggingFace embedding model to ensure that works as expected with exact
-vector search.
+This tool extends MTEB benchmarking tool to test a HuggingFace embedding model performance on both Retrieval and Reranking
+tasks based on custom datasets.
 
 ### [Approximate Search Evaluator](docs/approximate_search_evaluator/README.md)
 
-This tool provide a flexible tool to deply RRE and extract metrics to test your search engine collection given a 
+This tool provides a flexible tool to deply RRE and extract metrics to test your search engine collection given a 
 [template](https://github.com/SeaseLtd/rated-ranking-evaluator/wiki/What%20We%20Need%20To%20Provide#query-templates).
 
 ## Quickstart: tools installation
@@ -73,7 +73,7 @@ uv run embedding_model_evaluator --config <path-to-config-yaml>
 By default, the CLI is pointing to the  
 [configs/](configs/embedding_model_evaluator/embedding_model_evaluator_config.yaml) inside the `configs/` directory.
 
-## Running Embedding Model Evaluator
+## Running Approximate Search Evaluator
 For a detailed description to fill in configuration file (e.g., 
 [Config](configs/approximate_search_evaluator/approximate_search_evaluator_config.yaml)) you can look at the 
 [README](docs/approximate_search_evaluator/README.md).

rre-tools/configs/dataset_generator/dataset_generator_config.yaml

@@ -39,7 +39,7 @@ doc_fields:
 
 # (Optional) File containing predefined queries to use instead of or alongside generated ones
 # If available, must be in a txt format.
-#queries: "queries.txt"
+#queries: "templates/queries.txt"
 
 # (Optional) Whether to generate queries from documents
 # Default: true; set to false to disable query generation from documents

rre-tools/configs/dataset_generator/llm_config.yaml

@@ -7,7 +7,7 @@ name: openai
 model: gpt-5-nano-2025-08-07
 
 # Maximum number of tokens the model may return
-max_tokens: 2000
+max_tokens: 4000
 
 # Environment variable where LLM API key is stored
 api_key_env: OPENAI_API_KEY

rre-tools/configs/embedding_model_evaluator/embedding_model_evaluator_config.yaml

@@ -1,5 +1,5 @@
 # Model ID (HuggingFace models).
-model_id: "sentence-transformers/all-MiniLM-L6-v2"
+model_id: "sentence-transformers/all-MiniLM-L12-v2"
 
 # Accepted values: 'retrieval', 'reranking'
 task_to_evaluate: "retrieval"

rre-tools/docker-services/scripts/README.md

@@ -0,0 +1,13 @@
+## Extract BBC News Dataset 
+
+[The script](extract_bbc_news_dataset.py) downloads and processes BBC News articles from the [RealTimeData/bbc_news_alltime](https://huggingface.co/datasets/RealTimeData/bbc_news_alltime)
+dataset using HuggingFace.
+
+It filters, deduplicates, truncates long content, adds UUIDs, and saves the results to a JSON file.
+
+```
+python extract_bbc_news_dataset.py --filename output.json
+```
+Arguments:
+
+`--filename`: Output JSON filename (default: dataset.json)

rre-tools/docker-services/scripts/extract_bbc_news_dataset.py

@@ -0,0 +1,99 @@
+import argparse
+import json
+import uuid
+
+from datasets import load_dataset
+from tqdm import tqdm
+
+CONTENT_MAX_LEN = 1000
+DATASET_SIZE = 100000
+
+
+def truncate_content(txt: str) -> str:
+    if len(txt) <= CONTENT_MAX_LEN:
+        return txt
+
+    truncated_txt: str = txt[:CONTENT_MAX_LEN]
+    last_dot_index = truncated_txt.rfind('.')
+
+    if last_dot_index != -1:
+        # text size <= CONTENT_MAX_LEN
+        return txt[:last_dot_index + 1]
+
+    # find the first dot after CONTENT_MAX_LEN
+    next_dot_index = txt.find('.', CONTENT_MAX_LEN)
+
+    if next_dot_index != -1:
+        return txt[:next_dot_index + 1]
+
+    return truncated_txt
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser(description="Extract BBC News Dataset")
+    parser.add_argument('--filename', type=str, default='dataset.json', help='Output filename')
+    args = parser.parse_args()
+
+    months = []
+    end_year = 2025
+    end_month = 6
+
+    start_year = 2017
+    start_month = 1
+
+    curr_year = end_year
+    curr_month = end_month
+    while curr_year >= start_year and curr_month >= start_month:
+        months.append(f"{curr_year}-{curr_month:02d}")
+        curr_month -= 1
+        if curr_month == 0:
+            curr_month = 12
+            curr_year -= 1
+
+    all_results = []
+    seen_links = set()
+
+    for month in tqdm(months):
+        ds = load_dataset("RealTimeData/bbc_news_alltime", month)
+
+        for elem in ds['train']:
+            # skip if section=empty/None
+            if not elem.get("section") or elem.get("section") is None:
+                continue
+
+            if not elem.get("title"):
+                continue
+
+            # skip duplicates based on the web link
+            link = elem.get("link")
+            if not link or link in seen_links:
+                continue
+            seen_links.add(link)
+
+            # truncate long content
+            content = elem.get("content", "")
+            if not content:
+                continue
+            text = truncate_content(content)
+            elem["content"] = text
+
+            elem["id"] = str(uuid.uuid4())
+
+            #  id first shows up in the json file
+            id_val = elem.pop("id")
+            new_elem = {"id": id_val, **elem}
+
+            all_results.append(new_elem)
+            if len(all_results) == DATASET_SIZE:
+                break
+        if len(all_results) == DATASET_SIZE:
+            break
+
+    # for solr + vespa
+    with open(args.filename, "w", encoding="utf-8") as f:
+        json.dump(all_results, f, ensure_ascii=False, indent=4)
+    # for opensearch + elasticsearch
+    # with open("dataset.jsonl", "w", encoding="utf-8") as f:
+    #   for obj in all_results:
+    #      f.write(json.dumps(obj, ensure_ascii=False) + "\n")

rre-tools/docker-services/solr-init/solr_init.py

@@ -21,8 +21,9 @@
 EMBEDDINGS_FILE = os.path.join(EMBEDDINGS_FOLDER, "documents_embeddings.jsonl")
 TMP_FILE = os.getenv("TMP_FILE", "/tmp/merged_dataset.json")
 
-DEFAULT_TIMEOUT = float(os.getenv("DEFAULT_TIMEOUT", "10.0"))
+DEFAULT_TIMEOUT = int(os.getenv("DEFAULT_TIMEOUT", "600"))
 FORCE_REINDEX = os.getenv("FORCE_REINDEX", "false").lower() == "true"
+INDEX_BATCH_SIZE = 1000
 
 logging.basicConfig(
     stream=sys.stdout,
@@ -32,7 +33,7 @@
 log = logging.getLogger("solr_init")
 
 
-def wait_for_solr_core(endpoint: str, timeout: int = 30, interval: float = 1.0) -> None:
+def wait_for_solr_core(endpoint: str, timeout: int, interval: float = 1.0) -> None:
     """Waits until Solr core /admin/ping endpoint returns 200 or timeouts."""
     ping_url = f"{endpoint.rstrip('/')}/admin/ping?wt=json"
     log.info("Waiting for Solr core at %s ...", ping_url)
@@ -175,23 +176,50 @@ def create_vector_field(endpoint: str, dimension: int) -> None:
 
 def index_documents(endpoint: str, docs: list[dict[str, Any]]) -> None:
     """
-    Sends documents to /update endpoint and commit.
+    Sends documents to /update endpoint in batches and commit at the end.
     """
-    update_url = f"{endpoint.rstrip('/')}/update?commit=true"
-    log.info("Indexing %d documents into Solr", len(docs))
-    try:
-        response = requests.post(update_url, json=docs, timeout=60.0)
-        response.raise_for_status()
-        log.info("Indexing successful (status=%s)", response.status_code)
-    except requests.RequestException as e:
-        log.error("Failed to index documents: %s", e)
-        raise
+    total_docs = len(docs)
+    if total_docs == 0:
+        log.info("No documents provided for indexing.")
+        return
+
+    log.info("Indexing %d documents into Solr", total_docs)
+    num_batches = (total_docs + INDEX_BATCH_SIZE - 1) // INDEX_BATCH_SIZE
+
+    update_url_no_commit = f"{endpoint.rstrip('/')}/update?commit=false"
+    update_url_commit = f"{endpoint.rstrip('/')}/update?commit=true"
+
+    for i in range(num_batches):
+        start_index = i * INDEX_BATCH_SIZE
+        end_index = min((i + 1) * INDEX_BATCH_SIZE, total_docs)
+        batch = docs[start_index:end_index]
+
+        if not batch:
+            continue
+
+        is_last_batch = (i == num_batches - 1)
+
+        current_update_url = update_url_commit if is_last_batch else update_url_no_commit
+        commit_status = "true" if is_last_batch else "false"
+
+        log.info(f"Sending Batch {i + 1}/{num_batches} ({len(batch)} docs, commit={commit_status})")
+
+        try:
+            response = requests.post(current_update_url, json=batch, timeout=DEFAULT_TIMEOUT)
+            response.raise_for_status()
+            log.debug(f"Batch {i + 1} indexing successful (status={response.status_code})")
+
+        except requests.RequestException as e:
+            log.error(f"Failed to index batch {i + 1}: {e}")
+            raise Exception(f"Failed during batch {i + 1} indexing.") from e
+
+    log.info("Successfully indexed %d documents in %d batches.", total_docs, num_batches)
 
 
 def main() -> int:
     log.info("Starting solr_init.py")
     try:
-        wait_for_solr_core(COLLECTION_ENDPOINT, timeout=60, interval=1.0)
+        wait_for_solr_core(COLLECTION_ENDPOINT, timeout=DEFAULT_TIMEOUT, interval=1.0)
     except Exception as e:
         log.error("Solr core not available: %s", e)
         sys.exit(1)

rre-tools/src/rre_tools/embedding_model_evaluator/main.py

@@ -12,6 +12,7 @@
 
 import argparse
 import logging
+import time
 from typing import Any
 
 import mteb
@@ -101,14 +102,17 @@ def main() -> None:
 
     # --- Evaluation (in-memory) ---
     log.info("Starting MTEB evaluation...")
+    start = time.time()
     evaluation = mteb.MTEB(tasks=[task])
     evaluation.run(
         model=model_with_cache,
         output_folder=config.output_dest,
         overwrite_results=True,
         config=config,
     )
+    end = time.time()
     log.info("Finished MTEB evaluation.")
+    log.info(f"Time took for MTEB evaluation: {end - start:.2f} seconds")
 
     # --- Optional: write embeddings (kept for parity with previous behavior) ---
     writer = EmbeddingWriter(

rre-tools/templates/queries.txt

@@ -0,0 +1,10 @@
+artificial intelligence
+President Trump
+vacation
+covid-19
+Ukraine war
+bitcoin crypto
+tech layoffs
+job market
+Gaza strip
+Cannes film festival

rre-tools/templates/template_solr.json

@@ -1,7 +1,7 @@
 {
   "q": "$query",
   "defType": "edismax",
-  "qf": "title^2 description content",
+  "qf": "title description content",
   "mm": 1,
   "rows": 10,
   "fl": "id,title,description,content,score"