Embeddings search engine

Installing dependencies

Install dependencies via pip or uv

pip install -r requirements.txt

Layout

src/
├── engine.py
├── api/
│   ├── __init__.py
│   ├── __pycache__
│   ├── endpoints.py
│   └── requests.py
├── evaluator/
│   ├── __init__.py
│   ├── evaluator.py
│   └── metrics.py
├── finetuning/
│   ├── __init__.py
│   └── tuning.py
└── models/
    ├── __init__.py
    └── document.py

---

Part 1

Part 1 of the project is a simple embeddings-based search engine with three endpoints:

• POST /api/documents: ingest a collection of documents
• DELETE /api/documents: deletes a collection of documents
• POST /api/search: searches the indexed collection by a text query

The implementation uses:

• ChromaDB as the in-process vector store (with PersistentClient for on-disk persistence)
• Sentence-Transformers embeddings via Chroma’s SentenceTransformerEmbeddingFunction using Qwen Embeddings
• FastAPI as the web server

To run the app just run:

fastapi run app.py

API

To see the API documentation, check openapi.json

---

Part 2

Part 2 of the project is an evaluator, to run it, simply run

python run_evaluator.py --split [train, test, val] --top-k [1, 2, ...]

The results yield:

• Recall@10: 0.9360
• MRR@10: 0.8297
• nDCG@10: 0.8565

---

Part 3

To finetune the initial model, since the model is too large, I made a container with a conda environment env.yml in order to run the code on a HPC cluster. Since texts tend to be long and the model contains around 600 million parameters, I also implemented gradient accumulation in order to be able to use a larger batch size and keep variance between batches low. Otherwise, even the best GPUs run out of VRAM when fine-tuning on a single unit.

We use MultipleNegativesRankingLoss as the loss function since it is widely recognised as one of the best loss functions for semantic search and retrieval since it allows positive pairs to be pulled together in embedding space.

The job to finetune the model on the cluster is defined in this file and the script is here.

The plot to visualize loss can be seen here