Cognitive-rag / 03_embedding
README.md

03_embedding

This step generates vector embeddings for chunked text and user queries.

Current setup uses NVIDIA NIM embeddings via langchain-nvidia-ai-endpoints with nvidia/llama-3.2-nv-embedqa-1b-v2. Embeddings are L2-normalized when --normalize is provided, which is recommended for cosine similarity.

Inputs

  • Chunked JSONL from chunking step: 02_chunking/output/chunked_mtsamples.jsonl

Outputs

  • Chunk embeddings JSONL: output/embeddings_mtsamples.jsonl
  • Run stats: output/embedding_stats.json
  • Query embedding JSON: output/query_embedding.json

Note: --output-dir is relative to your current working directory. If you run from the repo root, outputs go to ./output. If you run from 03_embedding, outputs go to 03_embedding/output.

Install

pip install -r 03_embedding/requirements.txt

Set your NVIDIA API key:

set NVIDIA_API_KEY=your_key_here

Embed chunks

python 03_embedding/embed_chunks.py --normalize

Optional overrides:

python 03_embedding/embed_chunks.py --input 02_chunking/output/chunked_mtsamples.jsonl --output-dir 03_embedding/output --normalize

Embed a query

python 03_embedding/embed_query.py --query "How long is parental leave?" --normalize

WSL Usage (Embeddings Only)

Run these from WSL after chunking is complete.

cd /mnt/c/Users/ajwaa/OneDrive/Desktop/CognitiveRAG
source .venv/bin/activate
export NVIDIA_API_KEY=your_key_here

# Text
python3 03_embedding/embed_chunks.py \
  --input 02_chunking/output/test_text/chunked_documents.jsonl \
  --output-dir 03_embedding/output/test_text \
  --normalize \
  --model nvidia/llama-3.2-nv-embedqa-1b-v2

# Audio
python3 03_embedding/embed_chunks.py \
  --input 02_chunking/output/test_audio/chunked_documents.jsonl \
  --output-dir 03_embedding/output/test_audio \
  --normalize \
  --model nvidia/llama-3.2-nv-embedqa-1b-v2

# Image
python3 03_embedding/embed_chunks.py \
  --input 02_chunking/output/test_image/chunked_documents.jsonl \
  --output-dir 03_embedding/output/test_image \
  --normalize \
  --model nvidia/llama-3.2-nv-embedqa-1b-v2

# Video
python3 03_embedding/embed_chunks.py \
  --input 02_chunking/output/test_video/chunked_documents.jsonl \
  --output-dir 03_embedding/output/test_video \
  --normalize \
  --model nvidia/llama-3.2-nv-embedqa-1b-v2

Notes

  • Model: nvidia/llama-3.2-nv-embedqa-1b-v2
  • Normalization: set --normalize for cosine similarity workflows