Semantic-Vector-Search/config.py at main · itsyourdecide/Semantic-Vector-Search · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
"""
RustySearch — Configuration
============================
All tuneable parameters in one place.
Change these values to adapt the project to your data and hardware.
"""

import os

# ─── Paths ────────────────────────────────────────────────────────────────────
BASE_DIR  = os.path.dirname(os.path.abspath(__file__))
DATA_DIR  = os.path.join(BASE_DIR, "data")

INDEX_FILE = os.path.join(DATA_DIR, "news_index.bin")   # Rust IVF index
TEXTS_FILE = os.path.join(DATA_DIR, "news_texts.json")  # raw document store

# ─── Embedding model ──────────────────────────────────────────────────────────
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
DIMENSION       = 384   # output vector size of the model above

# ─── Dataset (used by build_real_db.py) ───────────────────────────────────────
DATASET_NAME  = "dbpedia_14"   # HuggingFace dataset name
DATASET_SPLIT = "train"
NUM_RECORDS   = 100_000     # how many articles to embed and index
                             # 10_000 → ~2 min   |   100_000 → ~15 min

# ─── Index building ───────────────────────────────────────────────────────────
KMEANS_MAX_ITER   = 15      # K-means iterations (more = better quality, slower)
ENCODE_BATCH_SIZE = 256     # sentences per GPU/CPU batch during encoding
LOAD_CHUNK_SIZE   = 10_000  # vectors pushed to Rust per progress-bar step

# ─── Search ───────────────────────────────────────────────────────────────────
N_PROBE      = 10    # IVF clusters probed per query (higher = more recall, slower)
DEFAULT_TOP_K = 5   # default number of results returned

# ─── API server ───────────────────────────────────────────────────────────────
API_HOST    = "0.0.0.0"
API_PORT    = 8000
API_TITLE   = "RustySearch AI API"
API_VERSION = "1.0.0"