Jyosa Labs Blogs

FP8, KV Cache Quantization, and the "!!!!" Bug

2026-03-17T00:00:00+00:00

so at the end of my attention post, i teased about kv-cache and inference optimizations. well, here’s that follow up - except it turned into something way bigger than i planned. what started as “let me explain kv-cache quantization” became a full investigation into a production bug where a multimodal model just… starts outputting "!!!!!!!!!!!!!!!!". yeah, literally exclamation marks. nothing else.

to actually understand why that happens, we need to go deep - like, all the way down to individual bits in a floating point number. i know that sounds excessive, but i promise every piece connects. we’ll build up from binary fractions to FP8 format to quantization scales, and then watch the whole thing collapse like dominoes when one design decision goes wrong.

fair warning - this is a long one and it’s dense. the first half covers floating point fundamentals (how numbers actually live in 8 bits), and the second half traces the vLLM bug from root cause to garbage output. if you need to split it up, the transition between those two halves is a natural stopping point.

part 1: floating point fundamentals

okay so you know how decimal works, right? digits after the decimal point represent powers of 10:

0 . 3   5   7
    ↓   ↓   ↓
    1/10 1/100 1/1000
    10⁻¹ 10⁻² 10⁻³

0.357 = 3×(1/10) + 5×(1/100) + 7×(1/1000)
      = 0.3 + 0.05 + 0.007

each position is the previous position ÷ 10. simple enough.

well binary works the exact same way, just with base 2 instead of base 10:

0 . 1   0   1
    ↓   ↓   ↓
    1/2  1/4  1/8
    2⁻¹  2⁻²  2⁻³

0.101 (binary) = 1×(1/2) + 0×(1/4) + 1×(1/8)
               = 0.5 + 0 + 0.125
               = 0.625 (decimal)

each position is the previous position ÷ 2. that’s literally why it’s 1/2, 1/4, 1/8 - they’re just 2⁻¹, 2⁻², 2⁻³. nothing magical about it. just like decimal uses powers of 10, binary uses powers of 2.

the mantissa (where your actual digits live)

the mantissa (some people call it “significand”) holds the precision bits - the actual digits of your number. think of it as the “meat” of the number, while the exponent just tells you where to put the decimal point.

in FP8 E4M3 (which we’ll be using throughout this post), we have 3 mantissa bits. each bit represents a fractional position:

1st bit → 1/2  (0.5)
2nd bit → 1/4  (0.25)
3rd bit → 1/8  (0.125)

so if your mantissa bits are 101, you get: 1×(1/2) + 0×(1/4) + 1×(1/8) = 0.5 + 0 + 0.125 = 0.625

here’s all 8 possible mantissa values:

bits	calculation	fraction
000	0/2 + 0/4 + 0/8	0.000
001	0/2 + 0/4 + 1/8	0.125
010	0/2 + 1/4 + 0/8	0.250
011	0/2 + 1/4 + 1/8	0.375
100	1/2 + 0/4 + 0/8	0.500
101	1/2 + 0/4 + 1/8	0.625
110	1/2 + 1/4 + 0/8	0.750
111	1/2 + 1/4 + 1/8	0.875

with 3 bits, you get 8 levels of precision. that’s… not a lot. but hey, we’re cramming a number into 8 bits total, so sacrifices must be made.

and here’s the thing people don’t always get right away - these aren’t like integers where every value has its own exact slot. these 8 values act more like buckets. if your actual number is, say, 0.3, well tough luck - there’s no bucket for 0.3. your closest options are 0.250 and 0.375, and 0.3 gets rounded to whichever is nearest (0.250 in this case). the number 0.3 just doesn’t exist in this format, it gets snapped to the nearest bucket.

think of it like a ruler that only has marks every 0.125 apart. you can’t measure 0.3 exactly - you’d say “it’s about at the 0.25 mark.” the more mantissa bits you have, the finer the markings on your ruler. FP16 has 10 mantissa bits (1024 buckets per exponent range), FP32 has 23 (over 8 million buckets). we have 3 bits. 8 buckets. that’s the precision tax you pay for squeezing a number into a single byte.

and this bucket thing gets worse as numbers get bigger. near zero, your buckets might be 0.00195 apart (subnormal range). but near the max value of 448, the gap between adjacent representable numbers is 32 (the difference between 416 and 448). so a value like 430 gets rounded to either 416 or 448 - you lose up to 14 in either direction. same 8 buckets per exponent range, but the buckets themselves get wider as the numbers get larger. keep this in mind - it matters when we get to the quantization part.

the implicit leading 1 (free bit hack)

okay so here’s a clever trick that gives us a free extra bit of precision. to understand it, let’s think about scientific notation first.

in decimal, when you write a number in scientific notation, you always normalize it so there’s exactly one non-zero digit before the decimal point:

 → 1.234 × 10³
0056 → 5.6 × 10⁻³
7   → 4.27 × 10¹

you’d never write 0.001234 × 10⁶ or 1234.0 × 10⁰ - you always shift the decimal point until you have one non-zero digit up front.

binary does the same thing. you shift the binary point until you have a 1 before it:

11 (decimal) = 1011 (binary)
             = 1.011 × 2³    (shifted right 3 times)

0.125 (decimal) = 0.001 (binary)
                = 1.0 × 2⁻³  (shifted left 3 times)

but here’s what makes binary special - in decimal, that leading digit could be 1, 2, 3… up to 9. you have to actually store it because it could be any of those. but in binary? the only digits are 0 and 1. and since we shift until the leading digit isn’t zero, it can only be 1. there’s literally no other option.

so why store something that’s always 1? we don’t! we just agree that there’s a 1 there and don’t waste a bit on it:

what we STORE in the mantissa:  0 1 1   (3 bits)
what we actually MEAN:        1.0 1 1   (implicitly 1.011)
                              ↑
                              this 1 is free! never stored.

this means our 3 mantissa bits are actually giving us 4 digits of precision - the implicit 1 plus the 3 stored bits. for free. when you only have 8 bits total, getting one extra bit of precision without spending any storage is a pretty big deal.

so remember those 8 mantissa buckets from earlier (0.000 to 0.875)? with the implicit leading 1, they become 1.000 to 1.875. the mantissa value is always at least 1.0 and at most 1.875. the exponent then scales this up or down to wherever you need it.

the full formula for a normal number:

value = (-1)^sign × 2^(exponent - bias) × 1.mantissa
                                          ↑
                                    the free 1!

example: sign=0, exponent=1001 (9), mantissa=101
  = 1 × 2^(9-7) × 1.625
  = 1 × 4 × 1.625
  = 6.5

now quick thing - why is the bias 7? well, our exponent field can only store 0 to 15 (4 bits, all unsigned). but we want to represent both big numbers (like 256) and tiny numbers (like 0.015). the bias is just a trick to get both: instead of using the stored value directly, you subtract 7 from it.

stored 1  → 1 - 7 = -6  → 2⁻⁶ = 0.015  (tiny positive number)
stored 7  → 7 - 7 =  0  → 2⁰  = 1.0    (one, the middle)
stored 14 → 14 - 7 = 7  → 2⁷  = 128    (big number)
stored 15 → 15 - 7 = 8  → 2⁸  = 256    (biggest)

so stored value 7 maps to 2⁰ = 1.0, which sits right in the middle. below 7 you get exponents that shrink numbers (2⁻¹ = 0.5, 2⁻⁶ = 0.015 — still positive, just tiny!), above 7 you get exponents that grow them (2¹ = 2, 2⁸ = 256). it’s not about tiny numbers - it’s about whether the exponent makes your number bigger or smaller than 1.

why specifically 7 and not 8 or 6? it comes from the formula bias = 2^(bits-1) - 1. for 4 exponent bits: 2³ - 1 = 7. same logic everywhere - FP16 with 5 exponent bits uses bias 15, FP32 with 8 exponent bits uses bias 127. it just gives you a nice balanced split between the “bigger than 1” and “smaller than 1” side of things.

and just to really drive this home - what if we didn’t use a bias at all? then our stored exponent values 0 to 15 would just directly mean 2⁰ to 2¹⁵. the smallest exponent would be 2⁰ = 1, and everything we could represent would be ≥ 1.0. we’d have a much bigger range on the top end (up to 2¹⁵ = 32,768!) but we’d completely lose the ability to represent anything between 0 and 1. no 0.5, no 0.1, no 0.015 - none of it. for something like kv-cache values where activations can be small fractions, that would be useless. the bias is what gives us that whole sub-1.0 world by letting some exponents become negative (like -6), which makes the result a tiny positive number (2⁻⁶ = 0.015), not a negative one.

subnormal numbers (filling the gap near zero)

okay so there’s a problem. with normal numbers and the implicit leading 1, the smallest positive number we can represent is:

smallest exponent = 0001 (stored) → actual = 1-7 = -6
smallest mantissa = 000 → 1.000 (with implicit 1)

smallest normal = 2^-6 × 1.0 = 0.015625

that leaves a huge gap between 0 and 0.015625. you literally can’t represent 0.001 or 0.005 or 0.01!

without subnormals:
0 ──────────────────────────── 0.015625 ────→
↑                               ↑
zero                        smallest normal

can't represent anything in between!

the fix is subnormal mode. when the exponent bits are all zeros (0000), special rules kick in:

the exponent is fixed at 2⁻⁶ (doesn’t change)
the implicit leading 1 becomes a 0 instead (so it’s 0.mantissa not 1.mantissa)

now why 2⁻⁶ specifically? it’s not a random choice - it’s the same exponent as the smallest normal number. remember, the smallest normal has stored exponent 0001, which gives actual exponent 1 - 7 = -6, so it equals 2⁻⁶ × 1.0 = 0.015625. subnormals use that exact same 2⁻⁶, but since they drop the implicit leading 1 (using 0.mantissa instead of 1.mantissa), they produce values below 0.015625. this is the key insight - by sharing the same power of 2, subnormals extend smoothly downward from exactly where normals start.

this gives us 7 tiny values between 0 and the smallest normal number (and mantissa 000 with exponent 0000? that’s just plain old zero - 0.0 × 2⁻⁶ = 0):

mantissa	0.mantissa	value
000	0.0	0.0 × 2⁻⁶ = 0 (zero itself!)
001	0.125	0.125 × 2⁻⁶ = 0.00195
010	0.25	0.25 × 2⁻⁶ = 0.00391
011	0.375	0.375 × 2⁻⁶ = 0.00586
100	0.5	0.5 × 2⁻⁶ = 0.00781
101	0.625	0.625 × 2⁻⁶ = 0.00977
110	0.75	0.75 × 2⁻⁶ = 0.01172
111	0.875	0.875 × 2⁻⁶ = 0.01367

look at how nicely this stitches together. the largest subnormal is 0.875 × 2⁻⁶ = 0.01367, and the smallest normal is 1.0 × 2⁻⁶ = 0.01563. the gap between them is only 0.00195 - which is the same step size as between any two adjacent subnormals! there’s no sudden jump, no weird discontinuity. it’s a perfectly smooth transition from subnormal territory into normal territory. if they’d picked 2⁻⁷ instead, the largest subnormal would be 0.00684 and the smallest normal would still be 0.01563 - that’s a gap of 0.00879, way bigger than the subnormal step size. the numbers would “jump” at the boundary and the whole point of having subnormals would be undermined.

why does this matter? without subnormals, you get “sudden underflow”:

a = 0.012
b = 0.010
a - b = 0.002 → UNDERFLOWS TO 0! (wrong!)

with subnormals:
a - b = 0.002 → stored as subnormal (correct!)

subnormals give us extra buckets near zero. pretty important when you’re doing math with tiny numbers, which… yeah, we’ll see why that matters soon.

FP8 E4M3: the format that started the whole mess

alright, now let’s put it all together and understand FP8 E4M3 properly. this is the specific format vLLM uses for KV cache quantization, and understanding it is key to understanding the bug.

FP8 E4M3 format:
┌───┬─────────┬─────────┐
│ S │ E E E E │ M M M   │
└───┴─────────┴─────────┘
  1      4         3     = 8 bits total

S = sign bit (1 bit)
E = exponent bits (4 bits), bias = 7
M = mantissa bits (3 bits)

8 bits. that’s it. one byte per value compared to FP16’s 2 bytes - 50% memory reduction. sounds great for kv-cache, right? well… yes, but there’s a catch. actually there are several catches.

why is the max value exactly 448?

this is where it gets fun. let me walk you through it step by step, because this number is crucial to understanding the bug later.

step 1: biggest exponent we can use

4 exponent bits can store 0000 to 1111 (0 to 15). as we discussed earlier, the bias of 7 lets us split this into tiny positive numbers (exponents below 7 shrink things below 1.0) and big numbers (exponents above 7 grow things). you subtract 7 from the stored value to get the actual exponent:

stored 0  → 0 - 7 = -7 → 2^-7 = 0.008  (tiny positive number)
stored 7  → 7 - 7 =  0 → 2^0  = 1.0    (middle ground)
stored 15 → 15 - 7 = 8 → 2^8  = 256    (big number)

so max exponent = 8, meaning 2⁸ = 256.

step 2: biggest mantissa we can use

3 mantissa bits give us 8 possible values (000 to 111). with the implicit leading 1, mantissa 110 gives us 1.75. but wait - what about 111 which would give 1.875?

here’s the thing - 111 is reserved for NaN (Not a Number). we need a way to say “this calculation went wrong” (like 0/0 or sqrt(-1)). so the last mantissa pattern is sacrificed for that purpose.

step 3: multiply them together

max value = 2^8 × 1.75
          = 256 × 1.75
          = 448 ✓

if 111 wasn’t reserved for NaN, max would be 256 × 1.875 = 480. we lose 32 numbers of range just to have NaN. worth it? absolutely - you’ll see why when NaN starts appearing in our bug trace.

FP8 E4M3 has NO infinity!

this one surprised me when i first learned about it. in standard IEEE formats (FP16/FP32), there’s a special bit pattern for infinity:

FP16: S.11111.0000000000 = ±∞
FP32: S.11111111.00000000000000000000000 = ±∞

but FP8 E4M3 designers made a different choice. when the exponent is 1111:

standard IEEE:
  exponent = 1111, mantissa = 000  →  ±Infinity
  exponent = 1111, mantissa ≠ 000  →  NaN

FP8 E4M3 (different!):
  exponent = 1111, mantissa = 000  →  256 (normal number!)
  exponent = 1111, mantissa = 001  →  288 (normal number!)
  exponent = 1111, mantissa = 010  →  320 (normal number!)
  ...
  exponent = 1111, mantissa = 110  →  448 (normal number! max!)
  exponent = 1111, mantissa = 111  →  NaN (only special value!)

trade-off: no way to represent ∞, but we get 7 extra usable numbers (256-448). when you only have 256 total bit patterns, every one counts.

E4M3 vs E5M2: two flavors of FP8

quick aside - there’s another FP8 format called E5M2 with 5 exponent bits and 2 mantissa bits:

format	exponent	mantissa	max value	use case
E4M3	4 bits	3 bits	448	inference (better precision)
E5M2	5 bits	2 bits	57,344	training (larger range)

E4M3 gives you more precision (8 mantissa levels vs 4), E5M2 gives you more range (max 57,344 vs 448). for kv-cache quantization during inference, E4M3 is the standard choice because precision matters more than range… usually. when it doesn’t, well, that’s our bug.

part 2: the vLLM KV cache bug

alright, now we have all the building blocks. let’s trace this bug from the beginning.

the problem

when running Qwen3-VL (a multimodal model that handles both text and images) with FP8 KV cache quantization in vLLM:

initial requests may work fine
at some point, all subsequent requests start failing - both text AND image
output becomes: "!!!!!!!!!!!!!!!!!!!!!!"

the configuration that triggered it:

args: [
  '--kv-cache-dtype', 'fp8',
  '--calculate-kv-scales',
  # ... other args
]

two flags. that’s all it took to create a cascading failure. let me show you exactly why.

vLLM’s scale calculation: the root cause

so here’s the deal with FP8 quantization - you can’t just shove FP16 values into FP8 directly. you need a scale factor to map the range of your actual values into the representable range of FP8 (which maxes out at 448, remember?).

vLLM computes these scales dynamically during the first forward pass. every time forward() is called on an attention layer, it first checks whether scales need computing:

# from vllm/model_executor/layers/attention/attention.py

class Attention(nn.Module, AttentionLayerBase):
    def forward(self, query, key, value, ...):
        if self.calculate_kv_scales:
            torch.ops.vllm.maybe_calc_kv_scales(
                query, key, value, self.layer_name
            )
        # ... rest of attention computation

that maybe_calc_kv_scales is a gate - it only runs once:

# from vllm/model_executor/layers/attention/attention.py

def maybe_calc_kv_scales(query, key, value, layer_name):
    forward_context = get_forward_context()
    self = forward_context.no_compile_layers[layer_name]

    # Only calculate if the layer's flag is True
    # This flag gets set to False after the first forward pass
    if not self.calculate_kv_scales:
        return

    self.calc_kv_scales(query, key, value)

and here’s the actual scale computation - the part that does the math:

# from vllm/model_executor/layers/attention/attention.py

def calc_kv_scales(self, query, key, value):
    self._q_scale.copy_(torch.abs(query).max() / self.q_range)
    self._k_scale.copy_(torch.abs(key).max() / self.k_range)
    self._v_scale.copy_(torch.abs(value).max() / self.v_range)
    self._q_scale_float = self._q_scale.item()
    self._k_scale_float = self._k_scale.item()
    self._v_scale_float = self._v_scale.item()
    # We only calculate the scales once
    self.calculate_kv_scales = False  # ← LOCKED FOREVER

see that last line? self.calculate_kv_scales = False. once. done. never again. the q_range, k_range, and v_range divisors come from environment constants:

# from vllm/envs.py

# Divisor for dynamic query scale factor calculation for FP8 KV Cache
Q_SCALE_CONSTANT: int = 200
# Divisor for dynamic key scale factor calculation for FP8 KV Cache
K_SCALE_CONSTANT: int = 200
# Divisor for dynamic value scale factor calculation for FP8 KV Cache
V_SCALE_CONSTANT: int = 100

and those constants get turned into tensors during initialization:

# from vllm/model_executor/layers/attention/attention.py

def set_default_quant_scales(layer, ...):
    # Initialize q/k/v range constants used by calc_kv_scales
    layer.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
    layer.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
    layer.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)

so the scale formula is: k_scale = max(abs(key_values)) / 200. if the max absolute key value in the first request is 5.0, then k_scale = 5.0 / 200 = 0.025.

the quantization and dequantization using this scale is straightforward:

quantization (storing to cache):
  fp8_value = original_value / k_scale

dequantization (reading from cache):
  original_value = fp8_value × k_scale

and here’s the actual dequantization in vLLM’s Triton attention kernel:

# from vllm/v1/attention/ops/chunked_prefill_paged_decode.py

# K : (HEAD_SIZE, BLOCK_SIZE)
K_load = tl.load(key_cache_ptr + k_offset, ...)

if K_load.dtype.is_fp8():
    K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype)
else:
    K = K_load

# V : (BLOCK_SIZE, HEAD_SIZE)
V_load = tl.load(value_cache_ptr + v_offset, ...)

if V_load.dtype.is_fp8():
    V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype)
else:
    V = V_load

see? load the FP8 value, cast to float32, multiply by the scale. simple. but the scale is LOCKED from the first request.

the scale computed on the first request is used for ALL subsequent requests, regardless of their value distributions. let that sink in for a second.

why one-time scale calibration is dangerous

imagine this scenario:

request 1: K values range [-5.0, +5.0]
  k_scale = 5.0 / 200 = 0.025
  max storable = 448 × 0.025 = 11.2
  ✓ works fine, all values fit

request 2: K values range [-20.0, +20.0]
  k_scale STILL = 0.025 (locked!)
  value 20.0 → 20.0 / 0.025 = 800
  but FP8 max = 448
  CLIPPED to 448!
  retrieved: 448 × 0.025 = 11.2

  ALL values > 11.2 become exactly 11.2!

and it works the other way too:

request 1: K values range [-50.0, +50.0]
  k_scale = 50.0 / 200 = 0.25
  max storable = 448 × 0.25 = 112
  ✓ no clipping

request 2: K values range [-2.0, +2.0]
  k_scale STILL = 0.25 (locked!)
  value 2.0 → 2.0 / 0.25 = 8
  only uses 8 out of 448 FP8 levels!
  severe precision loss (most bits wasted)

the first scenario (clipping) is what causes our bug. the second scenario (precision loss) would cause subtle quality degradation but not catastrophic failure. let’s trace the catastrophic one.

the clipping: where things start going wrong

the clipping happens at the hardware level. here’s the actual CUDA code that does the FP8 conversion:

// from vllm/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh

// float -> fp8
template <>
__inline__ __device__ uint8_t scaled_vec_conversion<uint8_t, float>(
    const float& a, const float scale,
    const __nv_fp8_interpretation_t fp8_type) {
  __nv_fp8_storage_t res =
      __nv_cvt_float_to_fp8(a / scale, __NV_SATFINITE, fp8_type);
  return (uint8_t)res;
}

see that __NV_SATFINITE? that’s NVIDIA’s saturation mode - it automatically clips values exceeding the FP8 range to 448 (the max) instead of producing infinity or NaN. it’s a safety mechanism - but in this case, “safety” creates a much worse problem.

here’s what the error looks like at different value levels. the error formula is simple - how far off is the retrieved value from the original, as a percentage:

error = (original - retrieved) / original × 100

example: original = 50.0, retrieved = 11.2
  error = (50.0 - 11.2) / 50.0 × 100 = 77.6%

original value	scaled (÷0.025)	after clipping	retrieved (×0.025)	error
5.0	200	200	5.0	0%
11.2	448	448	11.2	0%
15.0	600	448	11.2	(15-11.2)/15 = 25.3%
20.0	800	448	11.2	(20-11.2)/20 = 44.0%
50.0	2000	448	11.2	(50-11.2)/50 = 77.6%

anything above 11.2 gets clamped to exactly 11.2. now imagine you have an image with rich, diverse feature values and a bunch of them are above 11.2. they ALL become the same number. this is where it gets really bad.

the identical values problem

when multiple tokens (especially image tokens in multimodal models) have K/V values above the clipping threshold, they ALL collapse to the same value:

original image token K values (diverse):
  [15.3, 18.7, 12.1, 19.5, 16.8, 20.0, 14.2, 17.9]

after FP8 quantization + clipping:
  [11.2, 11.2, 11.2, 11.2, 11.2, 11.2, 11.2, 11.2]

that’s… that’s really bad. all diversity gone. all information gone. just a flat line of identical values. and this feeds right into the next failure.

layernorm variance collapse: the NaN factory

remember layernorm from my attention post? let me refresh. layernorm normalizes each value in a token’s embedding individually, using this formula:

LayerNorm(x) = (x - μ) / sqrt(σ² + ε)

where:
  μ = mean of x
  σ² = variance of x
  ε = small constant (typically 1e-12)

so for a token embedding like [a, b, c, d], it computes μ and σ² across those values, then normalizes each one individually:

output[0] = (a - μ) / sqrt(σ² + ε)
output[1] = (b - μ) / sqrt(σ² + ε)
output[2] = (c - μ) / sqrt(σ² + ε)
output[3] = (d - μ) / sqrt(σ² + ε)

every element gets the same μ and σ², but each one is normalized separately. this matters because if σ² goes bad, every single element in that token’s embedding blows up.

now watch what happens when you feed it those identical values:

x = [11.2, 11.2, 11.2, 11.2, 11.2, 11.2, 11.2, 11.2]
μ = 11.2
σ² = Σ(x - μ)² / n = [0 + 0 + 0 + ...] / 8 = 0

variance is zero. okay, that’s what the ε is for, right? to prevent division by zero? well… not exactly. here’s where floating point strikes back.

in practice, GPUs use the efficient variance formula:

σ² = E[x²] - E[x]²

mathematically equivalent to the standard formula, but with floating point arithmetic, E[x²] and E[x]² are computed separately, and each has tiny rounding errors. when the true variance is zero (identical values), these rounding errors can actually make the result slightly negative:

11.2 is stored as 11.199999809265137 in FP32

with different accumulation order/rounding:
  σ² = -0.000000317835426  ← NEGATIVE!

and then:

denominator = sqrt(σ² + ε)
            = sqrt(-0.000000317835426 + 0.000000000001)
            = sqrt(-0.000000317834426)
            = NaN  ← square root of negative number!

boom. NaN. and once you have NaN, it’s game over. because NaN is the most infectious value in all of computing.

NaN propagation: one bad apple spoils the batch

this is the part that blew my mind when i first traced through it. a single NaN token corrupts EVERYTHING in one attention layer. let me show you exactly how.

starting point: let’s say token 2 (an image token) got the identical values and produced NaN after layernorm. the other tokens are fine.

hidden states after LayerNorm:
  token 0: [0.5, -0.3, 0.8, 0.2, ...]     ← normal
  token 1: [0.1, 0.7, -0.4, 0.3, ...]     ← normal
  token 2: [NaN, NaN, NaN, NaN, ...]       ← corrupted!
  token 3: [-0.2, 0.4, 0.9, -0.5, ...]    ← normal

Q/K/V projections: token 2’s NaN hidden states get multiplied by the weight matrices to compute Q, K, and V (remember from my attention post, each token’s hidden state gets projected through learned weight matrices). here’s what that looks like for the query projection:

Q₂ = hidden₂ @ W_q

[NaN, NaN, NaN, NaN] @ [[w₁₁, w₁₂, w₁₃],     [NaN, NaN, NaN]
                        [w₂₁, w₂₂, w₂₃],  =   
                        [w₃₁, w₃₂, w₃₃],      NaN × w₁₁ + NaN × w₂₁ + ...
                        [w₄₁, w₄₂, w₄₃]]      = NaN for EVERY output element

every output element involves summing products where at least one factor is NaN, so the entire result is NaN. same thing happens for K and V projections. so token 2’s Q, K, and V are all completely NaN.

attention scores (QK^T): this is where it spreads. remember, attention scores are computed as Q @ K^T - every query dot-products with every key:

        k₀    k₁    k₂    k₃
                    ↑
                (all NaN)
    ┌────────────────────────
q₀  │  s₀₀   s₀₁   NaN   s₀₃ │
q₁  │  s₁₀   s₁₁   NaN   s₁₃ │
q₂  │  NaN   NaN   NaN   NaN │  ← entire row is NaN
q₃  │  s₃₀   s₃₁   NaN   s₃₃ │
    └────────────────────────

column 2 is NaN because k₂ is NaN. row 2 is NaN because q₂ is NaN. every other row has at least one NaN (from column 2).

softmax: the point of no return

softmax(x)_i = exp(x_i) / Σⱼ exp(x_j)

for row 0: [s₀₀, s₀₁, NaN, s₀₃]

exp([s₀₀, s₀₁, NaN, s₀₃]) = [exp(s₀₀), exp(s₀₁), NaN, exp(s₀₃)]

denominator = exp(s₀₀) + exp(s₀₁) + NaN + exp(s₀₃)
            = (some value) + NaN
            = NaN

softmax = [exp(s₀₀)/NaN, exp(s₀₁)/NaN, NaN/NaN, exp(s₀₃)/NaN]
        = [NaN, NaN, NaN, NaN]

every single row has a NaN in column 2, which means every row’s softmax denominator is NaN, which means every attention weight becomes NaN. one corrupted token just made ALL attention weights NaN.

attention output: NaN weights × values = NaN for ALL tokens.

residual connection: even the tokens that had perfectly normal inputs are now dead:

token 0: [0.5, -0.3, 0.8, ...] + [NaN, NaN, NaN, ...]
       = [NaN, NaN, NaN, ...]  ← normal + NaN = NaN!

the whole thing is corrupted. in ONE layer. within a single attention computation, we went from 1/4 tokens having NaN (25%) to 4/4 tokens (100%). the remaining layers just propagate 100% NaN forward.

start of layer:     1/4 tokens NaN (25%) - just token 2
after QK^T:         row 2 and column 2 are NaN
after softmax:      ALL rows become NaN
after attention:    ALL tokens' outputs are NaN
after residual:     4/4 tokens NaN (100%)

corruption complete in ONE layer!

why ALL future requests fail too

here’s the truly nasty part. once NaN gets into the KV cache, it stays there. and because of how attention works - every new query attends to ALL cached keys - any new request will compute attention scores against those NaN keys.

new request arrives (perfectly valid text):
1. new_query dot NaN_cached_key = NaN
2. NaN in softmax denominator = NaN
3. NaN attention weights = NaN output

result: "!!!!!!!!!!!!!!!!!!!"

the model is effectively poisoned until you restart the whole inference server. both text requests and image requests fail. it doesn’t matter how simple or valid the new input is - the cached NaN values will corrupt everything through the attention mechanism.

but why “!” specifically? this part is actually pretty interesting. two things line up perfectly to produce exclamation marks.

thing 1: NaN breaks argmax. by this point, every single logit (the score the model gives to each possible next token) is NaN. the model needs to pick a token to output, so it runs argmax to find the highest-scoring one.

but NaN has a weird property - any comparison with it returns false. always.

NaN > 5.0?      false
NaN > -1000.0?  false
NaN > NaN?      also false

so when argmax scans through the list looking for the biggest value, it goes element by element comparing each one to the current best. since every comparison returns false, nothing ever “wins.” the current best never gets updated. and since argmax starts at index 0, it just… stays there. returns 0. done.

logits = [NaN, NaN, NaN, NaN, NaN, ...]

argmax scanning:
  start: best = index 0 (NaN)
  index 1: NaN > NaN? false. best stays 0
  index 2: NaN > NaN? false. best stays 0
  ...
  result: 0

thing 2: token ID 0 is “!” in Qwen’s vocabulary. okay but why is “!” at index 0? it’s not random - it comes from how the tokenizer was built. and to understand that, we need a quick detour into how tokenizers actually work. i promise this connects back.

bytes: the foundation. a byte is 8 bits - eight slots that can each be 0 or 1. how many combinations is that? 2 × 2 × 2 × 2 × 2 × 2 × 2 × 2 = 2⁸ = 256. so a byte can be any value from 0 to 255. every single thing on your computer - text, images, code - is stored as a sequence of bytes. when you type “hello”, your computer stores 5 bytes:

h → byte 104
e → byte 101
l → byte 108
l → byte 108
o → byte 111

the first 128 of these (bytes 0-127) follow a standard called ASCII - a lookup table from the 1960s where somebody just decided: byte 33 = !, byte 34 = ", byte 65 = A, byte 97 = a, etc. there’s nothing mathematical about it - byte 33 is ! for the same reason your phone number is your phone number. someone assigned it. the remaining 128 byte values (128-255) were never standardized the same way, but they exist and show up in text, so the tokenizer needs to handle them too.

the tokenizer’s starting point. a tokenizer needs to convert any text into numbers (token IDs) that the model can process. the simplest approach: make each of the 256 possible byte (buckets) values a token.

but wait - 256 tokens for every language on earth? how? because every character, in any language, is ultimately stored as a sequence of bytes via UTF-8 encoding:

"h"    = 1 byte:  [104]
"é"    = 2 bytes: [195, 169]
"🎉"  = 4 bytes: [240, 159, 142, 137]
"你"   = 3 bytes: [228, 189, 160]

every single one of those byte values falls between 0 and 255. always. so if you have a token for each of the 256 possible byte values, you can break any input into byte tokens:

"🎉" = bytes [240, 159, 142, 137]
     → token for 240, token for 159, token for 142, token for 137
     → 4 tokens. not efficient, but it works!

nothing is ever “unknown.” think about it - any text you throw at the tokenizer, whether it’s english, chinese, arabic, emoji, or some random binary garbage, is ultimately just a sequence of bytes. and every byte is a value between 0 and 255. since we have a token for each of those 256 values, we can always break the input down. worst case, every character becomes multiple tokens (like 🎉 = 4 tokens). but it always works. no “out of vocabulary” errors, ever.

is 4 tokens for one emoji efficient? no. but that’s where BPE comes in later - it learns that bytes [240, 159, 142, 137] always appear together, so it merges them into a single “🎉” token. the 256 byte tokens are the safety net. BPE is the optimization on top.

but how does UTF-8 squeeze ~1.1 million possible characters into byte values that only go up to 255? the trick is that not all 256 byte values work the same way. UTF-8 splits the range into roles:

bytes 0-127:    standalone characters (ASCII)
                each one IS a character by itself.
                "h" = byte 104. done. one byte, one character.

bytes 128-191:  continuation bytes (10xxxxxx)
                NEVER start a sequence - they're always
                the 2nd, 3rd, or 4th byte in a multi-byte character.

bytes 192-255:  leading bytes (11xxxxxx)
                START a multi-byte sequence. their prefix tells
                you how many continuation bytes follow:
                110xxxxx (192-223) → "1 more byte coming"  → 2-byte char
                1110xxxx (224-239) → "2 more bytes coming" → 3-byte char
                11110xxx (240-247) → "3 more bytes coming" → 4-byte char

so when you see byte 195 followed by byte 169, UTF-8 reads it as:

195 = 11000011 → starts with 110, so "1 continuation byte follows"
169 = 10101001 → starts with 10, confirmed continuation byte

extract the free bits: 00011 + 101001 = 00011101001 = 233 → that's "é"

it’s not “256 values for 256 characters” - it’s “128 values for ASCII characters + 128 values for building multi-byte sequences that cover the rest of the world.” the upper half are puzzle pieces that only work in combinations.

and those structure bits in each byte are what limit the total. the first byte’s leading bits signal how many bytes follow, and each continuation byte starts with 10:

1-byte char:  0xxxxxxx                            → 7 free bits  → 128 chars
2-byte char:  110xxxxx 10xxxxxx                   → 11 free bits → 1,920 chars
3-byte char:  1110xxxx 10xxxxxx 10xxxxxx          → 16 free bits → 61,440 chars
4-byte char:  11110xxx 10xxxxxx 10xxxxxx 10xxxxxx → 21 free bits → 1,048,576 chars

see the 110, 1110, 11110 prefixes? those eat up bits to say “this is a 2/3/4-byte sequence.” and every continuation byte wastes 2 bits on the 10 prefix. so in a 4-byte character, out of 32 total bits, only 21 are free to encode the actual character. 2²¹ ≈ 2 million, minus reserved ranges gives you ~1.1 million usable slots.

and this means byte 240 (binary 11110000) will never appear alone in valid UTF-8 - its leading 1111 tells you “3 more bytes are coming.” so while it has its own base token, that token only ever appears as part of a 4-byte sequence like [240, 159, 142, 137] (🎉).

this also means the system never breaks. when a new emoji like 🫠 gets added to Unicode in 2024, the tokenizer doesn’t need updating - 🫠 is just bytes [240, 159, 171, 160], and those byte tokens already exist. initially it’s 4 tokens (one per byte), but if the model gets retrained on data full of 🫠, BPE will merge those bytes into a single token. new characters are just new combinations of the same 256 bytes.

but there’s a catch. bytes 0 through 32 are “control characters” - invisible things:

byte 0  = NULL (literally nothing)
byte 9  = TAB
byte 10 = NEWLINE
byte 32 = SPACE

you can’t see these. if you try to print them, they either do nothing or move your cursor around. you wouldn’t want these as token names in your vocabulary because you couldn’t even look at them to debug things.

the invisible byte problem. these invisible bytes absolutely show up in real text - newlines and tabs are everywhere. so the tokenizer will encounter them. the problem isn’t processing them, it’s displaying them. if token 10 in your vocabulary is labeled with a newline character… you literally can’t see it in a vocabulary list.

GPT-2’s fix: take those ~68 invisible bytes and remap them to obscure unicode characters that weren’t already used. they picked characters starting at unicode code point 256 (Ā):

invisible bytes get visible stand-ins:

byte 0   (NULL)    → "Ā"  (unicode 256)
byte 9   (TAB)     → "ĉ"  (unicode 265)
byte 10  (NEWLINE) → "Ċ"  (unicode 266)
byte 32  (SPACE)   → "Ġ"  (unicode 288)

so if you ever see Ġ in a tokenizer dump and get confused - that’s just a space wearing a costume. nothing is lost. all 256 bytes are still represented, just with visible labels.

BPE: building words from bytes. those 256 base tokens are your starting alphabet. now BPE (Byte Pair Encoding) builds the rest of the vocabulary on top. it’s surprisingly simple - it always merges two tokens at a time, pairing the most frequent neighbors.

take your entire training dataset (billions of words) and split every word into individual byte tokens:

"the"   → ["t", "h", "e"]
"cat"   → ["c", "a", "t"]
"that"  → ["t", "h", "a", "t"]
"this"  → ["t", "h", "i", "s"]
...millions more words, all split into bytes

now scan through all of it and count every adjacent pair. not in one word - across the entire dataset:

"t" + "h" → 9,847,231 times  (from "the", "that", "this", "there", "them"...)
"h" + "e" → 7,234,102 times  (from "the", "he", "her", "here"...)
"i" + "n" → 5,891,445 times  (from "in", "ing", "into", "inside"...)
"a" + "t" → 4,200,000 times  (from "at", "cat", "that", "what"...)
"f" + "o" → 1,300,000 times  (from "for", "food", "before"...)

whichever pair has the highest count wins. that’s the entire decision. no AI, no heuristics - just “which two neighbors appeared together the most?” "t"+"h" at 9.8 million? it wins round 1.

and here’s the important bit - after each merge, BPE recounts everything because the data has changed:

before merging "t"+"h":
  "the" was ["t", "h", "e"]  →  pairs: "t"+"h", "h"+"e"

after merging "t"+"h" into "th":
  "the" is now ["th", "e"]   →  pairs: "th"+"e"

  "h"+"e" no longer exists in "the"!
  "th"+"e" is a brand new pair to count.

so round 2 counts all pairs again with the updated tokens. now "th"+"e" might have the highest count (because “the” is the most common english word), so it wins:

round 1: "t" + "h" appear together 9.8M times → merge into "th" (token 256)
  "the"  → ["th", "e"]       ← merged!
  "that" → ["th", "a", "t"]  ← merged!
  "cat"  → ["c", "a", "t"]   ← no "t"+"h", unchanged

round 2: "th" + "e" is now most frequent → merge into "the" (token 257)
  "the"  → ["the"]           ← one token!
  "that" → ["th", "a", "t"]  ← "th"+"a" not frequent enough yet

round 3: "i" + "n" → merge into "in" (token 258)

round 4: "in" + "g" → merge into "ing" (token 259)

notice - it’s always two things merging into one. never three at once. but because “th” is already a token from round 1, merging “th” + “e” in round 2 creates a 3-character token. it’s like LEGO - you snap two pieces together to make a bigger piece, then snap that bigger piece with another piece:

building "therefore":

bytes:   t  h  e  r  e  f  o  r  e
           ↘↙
merge:    th  e  r  e  f  o  r  e
            ↘↙            ↘↙
merge:    the    r  e  f  o  re
               ↘↙    ↘↙
merge:    the  re    fo   re
                     ↘↙
merge:    the  re   fore
            ↘↙
merge:    there  fore
             ↘↙
merge:    therefore

always two in, one out. the tokens just get bigger.

keep going for tens of thousands of merges and you end up with a vocabulary of ~150,000 tokens. common words like “the” and “and” are single tokens. less common words get split into pieces. rare words fall all the way back to individual bytes. nothing is ever “unknown” because those 256 byte tokens can spell out literally anything:

"the"          → [token 257]                    ← common, one token
"hello"        → [token 5000]                   ← fairly common, one token
"quantization" → [token 8821, token 1924]       ← split into two pieces
"vLLM"         → ["v", "L", "L", "M"]          ← rare, back to bytes

okay, back to why “!” is token 0. now you know the 256 base byte tokens are just a mapping table. the question is: what order do they go in? GPT-2’s designers started with the printable ASCII characters, beginning at byte 33 - which is "!":

# from the tokenizer source (bytes_to_unicode function)
bs = list(range(ord("!"), ord("~") + 1))  # starts at "!" (33)

# so the base vocabulary looks like:
# token 0 → "!"  (byte 33)
# token 1 → '"'  (byte 34)
# token 2 → "#"  (byte 35)
# token 3 → "$"  (byte 36)
# ...and so on

important thing to notice: token ID ≠ byte value. token 0 is byte 33, not byte 0. the ordering in the vocabulary is completely independent of the byte values themselves - the designers just chose to list printable characters first (starting at !), then other printable Latin characters, then the invisible bytes (remapped to unicode stand-ins) last. all 256 bytes are in the vocabulary, but their positions (token IDs) don’t match their byte values.

and remember how BPE merges frequent byte pairs into bigger tokens? those merged tokens get added after the 256 base bytes. so the full vocabulary ends up looking like this:

# token 0   → "!"           (base byte 33)
# token 1   → '"'           (base byte 34)
# ...
# token 255 → last base byte
# ─── BPE merges start here ───
# token 256 → "th"          (merged from "t" + "h")
# token 257 → "the"         (merged from "th" + "e")
# ...
# token 5000  → "hello"     (common word, merged over many rounds)
# token 8821  → "quant"     (part of "quantization")
# ...
# token 155000 → "🎉"       (merged from bytes 240+159+142+137)
# token 155001 → "🫠"        (merged from bytes 240+159+171+160)

so emoji tokens exist too! they’re just way up in the higher token IDs because they were created by BPE merging 4 base byte tokens into one. when someone types 🎉, the tokenizer doesn’t use 4 separate byte tokens - it recognizes the byte sequence [240, 159, 142, 137] and maps it to a single token (like token 155000). that’s the whole point of BPE: common sequences, whether they’re english words or popular emojis, get their own token ID.

so “!” isn’t special or meaningful - it’s just the first printable ASCII character, which makes it the first entry in the byte-level vocabulary, which makes it token ID 0. and when argmax defaults to index 0 because everything is NaN, you get “!”. every single time. forever.

that’s the full picture: NaN logits → argmax returns 0 → token 0 is “!” → "!!!!!!!!!!!!!!!!!!!"

(if this were a model with a different tokenizer that puts, say, "<pad>" at index 0, you’d get a different garbage output. but with GPT-2 style tokenizers - which Qwen, LLaMA, and many others use - it’s always “!”.)

oh, and it gets even worse if prefix caching is enabled. vLLM can cache KV entries for common prompt prefixes so they don’t need to be recomputed. if the NaN-corrupted KV entries get saved as a prefix cache entry, then any future request with a matching prefix will reuse those poisoned entries directly - without even going through the computation that would normally trigger the bug. the corruption essentially gets “baked in” to the cache. even without prefix caching though, the corruption persists within the running server’s memory until restart.

the complete causal chain

let me trace the whole thing end to end, because seeing the full chain really drives home how a tiny design decision cascades into a catastrophic failure:

1. --kv-cache-dtype fp8 + --calculate-kv-scales enabled
       ↓
2. first request arrives, scale computed and LOCKED
       ↓
3. subsequent request has K/V values outside calibrated range
       ↓
4. values clipped to FP8 max (448) during quantization
       ↓
5. many tokens get identical dequantized values
       ↓
6. LayerNorm variance → zero (or slightly negative due to FP arithmetic)
       ↓
7. sqrt(negative) = NaN
       ↓
8. NaN propagates through Q/K/V projections
       ↓
9. NaN in attention scores (QK^T column from corrupted key)
       ↓
10. softmax produces all-NaN weights (one NaN poisons entire sum)
       ↓
11. ALL tokens become NaN after residual connection
       ↓
12. NaN stored in KV cache
       ↓
13. ALL future requests attend to NaN keys → permanent corruption
       ↓
14. output: "!!!!!!!!!!!!!!!!!!!"

14 steps from a configuration flag to garbage output. and the fix is embarrassingly simple.

the fix

# instead of:
args: [
  '--kv-cache-dtype', 'fp8',
  '--calculate-kv-scales',
]

# use:
args: [
  '--kv-cache-dtype', 'auto',
  # remove --calculate-kv-scales entirely
]

--kv-cache-dtype auto just stores the KV cache in its original precision (FP16 or BF16) without trying to squeeze it into FP8. no scales, no clipping, no NaN. yeah you use more memory (roughly 2× for the cache), but the model actually works correctly. for multimodal models where image tokens and text tokens have very different value ranges, this is really the only safe option right now.

could vLLM make FP8 work properly here? totally. a few ideas:

recalculate scales per request - instead of locking them after the first pass, update them as new value ranges come in
per-token dynamic quantization - give each token its own scale instead of one global scale for everything
outlier-aware quantization - detect values that would clip and handle them separately

but as of now, with --calculate-kv-scales, the scale is computed once on the first request and locked forever. and that “forever” is what breaks everything.

takeaways

a few things i want you to walk away with:

FP8 E4M3 maxes out at 448 - that’s 256 × 1.75. the 256 comes from 2⁸, where 8 is the max actual exponent (stored exponent 15, minus bias 7, gives 15 - 7 = 8). and 1.75 is the largest usable mantissa (1.110 in binary - remember, 111 is reserved for NaN).
the scale is everything - the scale is basically the ruler you use to fit your numbers into FP8’s tiny range. if you set the ruler based on small numbers and then big numbers show up later, they all get chopped down to the same max value. that’s clipping. and clipping means different numbers become identical, which is where everything falls apart.
NaN is incredibly infectious - one NaN value in one token propagates through the attention mechanism to corrupt ALL tokens in ONE layer. softmax is the amplifier - it can’t handle even a single NaN in its input because the denominator sum becomes NaN.
floating point arithmetic has subtle failure modes - the variance of identical values should be zero mathematically, but floating point rounding can make it slightly negative, and sqrt(negative) = NaN. the gap between math and computation is where bugs hide.
the KV cache persists across requests - this is normally a feature (faster inference!), but when corruption gets in, it becomes a bug amplifier. NaN values in the cache poison all future requests until restart.

so yeah, that’s how two command line flags, a one-time scale calculation, and the fundamental limitations of 8-bit floating point combine to produce "!!!!!!!!!!!!!!!!!!!". pretty wild chain of events for something that started as a memory optimization.

next time you see garbage output from a quantized model, you’ll know exactly where to start looking. until next time, adios!

Attention is all you need guys

2026-02-22T00:00:00+00:00

heads up! this is a long one - around 6,200 words and 37k characters. grab a coffee (or two), find a quiet spot, and give it your undivided attention (pun intended), prefer to read it on a laptop or a big screen!

yeah, I know what you’re thinking - “I’ll just throw this at Claude and get the summary.” sure, you could do that. but if you actually want to understand attention from first principles, there’s no shortcut. trust me, the 30 minutes you spend here will save you hours of confusion later. your call though!

reading in two sittings? if you’re short on time, here’s a natural break point: read until the end of “positional embeddings” section (the first ~10-12 minutes). that covers why transformers exist and the foundations. then come back fresh for the second half which dives into how attention actually computes - the Q/K/V mechanics, multi-head attention, and all the juicy details. both parts build on each other, but taking a break there won’t break your flow.

here we go!

i know there are ton of explanations of the transformers on the web but after reading it for 18th time (yeah, I read that “attention” paper back in my 2nd year) and until recently, I quite got all the gist of it, so i mean i think it’s a good time to put that into words here.

tldr : this might be the only explanation you will ever need to actually understand attention from first principles, and i hope you have some idea of rnns, lstms and a bit of how the gradient descent works, yeah the og algo, if not, would suggest you to just have a quick read (spin your claude now!). ok so..

the sequence problem or why rnns kinda suck

okay so let’s rewind to like 2015-ish?! back then, if you wanted to do anything with text or sequential data, rnns were pretty much your only real option. i mean the core idea was actually super intuitive when you think about it.

so imagine you’re reading a really long paragraph word by word, and you’re trying to remember everything as you go along. you start with something like “the cat that was sitting on the mat that john bought last week…” and by the time you get to the end of the sentence, you’ve kinda forgotten what was at the beginning, right? well that’s basically what happened with rnns (recurrent neural networks) - the old school way of processing text before transformers.

rnns would process sentences one word at a time, maintaining this “mental note” about everything they’d seen so far. sounds reasonable, right? except they had this massive problem: the further back in the sentence you went, the more the information would just fade away. it’s because of this thing called vanishing gradients, yeah the og problem.

now to be fair, vanishing gradients isn’t just an rnn thing - it’s a fundamental problem in deep learning. even simple convnets or stacked linear layers suffer from this when networks get deep enough. that’s exactly why resnets introduced skip connections. but rnns made it worse because the sequential nature meant gradients had to flow through potentially hundreds of timesteps, one after another, that just made it worse!

basically during backpropagation (when the model weights are updated), the error signal has to travel backwards through time to update all those weights. the issue is that the gradients get multiplied at each timestep, and since these values are typically fractions, they shrink crazy fast.

just to make it more clear, imagine you have a sequence of 100 words. now at each step back, your gradient gets multiplied by something like a very small number, let’s say 0.9. by the time you reach the beginning, you’re looking at 0.9^100 ≈ 0.000027 - that’s practically zero dude! now if you want to update the weights by this much amount, it’s basically no update at all.

this is why rnns struggled with longer contexts. if something important happened 50 steps ago, the model barely updates those early weights because the gradient signal is too weak by the time it travels back. the connections between those distant parts of the sequence become nearly impossible to learn. and that’s where the og thing came, the transformers baby.

well what really is a transformer?!

keeping things straight, transformer is basically a machine learning model which does sequence modelling. what it means is when you have a sequence of things and you feed the model some first x-1 units of that sequence, it tries to predict what the xth unit should be. this is it. this works for literally any type of sequence data. you can feed a transformer text, music notes, protein structures, code - anything where order matters and patterns exist.

like music chords, where after seeing [C, F, G, Am], the model might predict the next chord is F again because that’s a common progression.

or in something like “I am from India and I am…” and the model predicted “Indian”, where all the words before “Indian” are those x-1 number of units and the xth unit was “Indian” which the model predicted! unlike rnns, transformers can easily look at the entire context at once, so they don’t forget what happened way back at the beginning.

now if you’re reading this as someone who doesn’t know ml (dude, learn ml bro, duh!), let me quickly explain how words actually get fed into these models. so there’s this process called tokenization, which is just a fancy way of saying “converting words to numbers.” a sentence like “I am from India and I am” gets converted into something like [12, 24, 45, 11, 19, 10, 2], where each word gets its own number. these mappings aren’t special - they’re just like assigning id numbers to each word in a vocabulary.

but why do this? because computers are good with numbers, not text. they can do math with numbers, but they can’t directly do math with the word “India.” so we convert everything to numbers, do all our fancy math operations, and then convert back to words at the end.

ok so as we have numbers now, we can go ahead and actually understand how transformers really work.

attention baby!

alright so here’s where things get interesting. the big breakthrough with transformers was basically saying: fuck sequential processing, what if we just look at everything at once?

quick detour: what are encoders and decoders?

in the original transformer paper, they had this encoder-decoder setup for translation. but what does it really mean?! well ok think of two boxes that are connected together - each with a completely different job. the encoder box takes your input sentence (say, english) and processes the whole thing at once to build up this rich understanding of what it means. it’s like someone reading and really getting what you wrote - the full context, the relationships between words, everything. and then it passes this deep understanding to the second box.

what does the encoder actually output? it produces a sequence of contextual embeddings - one vector per input token. but these aren’t just regular word embeddings. each vector now captures not just what that word means, but how it relates to every other word in the sentence. so the embedding for “cat” in “the cat sat” contains information about sitting, while “cat” in “the cat meowed” would have different contextual information. the encoder is essentially a powerful embedding model that understands context.

now the decoder box has the tougher job guys - it generates the output sentence (say, french) one word at a time, using both what it’s already generated AND constantly checking back with the encoder’s contextual embeddings. it’s like having a translator who first completely understands your english sentence, then carefully crafts the perfect french translation word by word.

so encoder = understanding everything at once and producing rich contextual embeddings, decoder = generating piece by piece while referring back to those embeddings. let’s say if you’re translating something like “the cat sat on the mat,” the encoder processes the whole english sentence in one shot, outputting contextualized vectors that capture that there’s a cat, it’s sitting, and it’s on a specific mat. then the decoder produces “le chat s’est assis sur le tapis” one word after another, attending to the encoder’s embeddings to make sure it’s capturing everything correctly.

and before we dive into encoder vs decoder drama, here’s the thing - they’re basically the same. like, the core building block is the attention mechanism, and both encoder and decoder use it. the only real difference is that decoders have this “causal masking” thing where each word can only look at previous words (because you’re generating text left to right), and this is something called “autoregressive nature” while encoders let you look at the full sentence in both directions. but the actual attention computation? identical. so instead of explaining them separately and repeating myself, let’s just understand how attention works, and then the encoder/decoder distinction becomes obvious.. let’s goo.

the core idea (why attention is actually genius)

so imagine you’re translating “the cat sat on the mat” to french. when you’re working on translating “sat”, you probably want to look at “cat” (who’s doing the sitting?) and “mat” (where are they sitting?). you don’t really care about “the” at that moment.

that’s attention. instead of forcing the network to remember everything in some compressed hidden state, you just let it look back at all the words and decide which ones matter right now. the network learns what to pay attention to.

here’s what made it click for me, just hear me out: rnns are like someone telling you a story, and you’re trying to remember everything they said. by the end, you’ve forgotten half of it. attention is like having the whole story written down in front of you, and you can just look at whichever part you need at any moment. at any time, you can just look back and observe the whole context again to make sense of the current moment.

okay but wait - if we’re looking at all words simultaneously, how does the model know that “cat sat on mat” is different from “mat sat on cat”? word order obviously matters, right?

this is where positional embeddings come in.

positional embeddings : gives you the wings

before anything else happens, we take each word and turn it into a vector (word embedding - pretty standard stuff). but then we add positional information directly to these embeddings. position 1 gets one pattern added, position 2 gets another, and so on. imagine it as something which is just basic as telling the model where a specific word really is, that’s what positional embeddings do.

the transformer uses these specific formulas to create position patterns:

PE(pos, 2i) = sin(pos / 10000^(2i/d))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d))

let me break it down so that it doesn’t look freaky. pos is the position number in the sequence (1st word, 2nd word, etc). i is the dimension index we’re calculating (0, 1, 2… up to d/2-1) as each i value creates two dimension values, 2i through sin and 2i+1 through cos. d is the total embedding size which is 512 for the original transformer. if d is 512, then i only needs to go from 0 to 255:

• i=0 gives dimensions 0 and 1

• i=1 gives dimensions 2 and 3

• …

• i=255 gives dimensions 510 and 511

let me calculate actual values for different positions with d=512. say we have a sequence “the cat sat on mat”, now position 1 is the word “the”, at position 2 we have “cat”:

for position 1 (pos=1):

for dimension i=0:
  PE(1, 0) = sin(1/10000^(0/512)) = sin(1/1) = sin(1) = 0.84
  PE(1, 1) = cos(1/10000^(0/512)) = cos(1/1) = cos(1) = 0.54

for dimension i=1:
  PE(1, 2) = sin(1/10000^(2/512)) = sin(1/1.0027) = sin(0.997) = 0.84
  PE(1, 3) = cos(1/10000^(2/512)) = cos(1/1.0027) = cos(0.997) = 0.55

for position 2 (pos=2):

for dimension i=0:
  PE(2, 0) = sin(2/10000^(0/512)) = sin(2/1) = sin(2) = 0.91
  PE(2, 1) = cos(2/10000^(0/512)) = cos(2/1) = cos(2) = -0.42

for dimension i=1:
  PE(2, 2) = sin(2/10000^(2/512)) = sin(2/1.0027) = sin(1.994) = 0.91
  PE(2, 3) = cos(2/10000^(2/512)) = cos(2/1.0027) = cos(1.994) = -0.41

as you can see, each position gets a completely unique vector of 512 values. position 1’s vector starts with [0.84, 0.54, 0.84, 0.55…] while position 2’s starts with [0.91, -0.42, 0.91, -0.41…].

there’s one important thing here - the position embedding for position 1 is always exactly the same, no matter if it’s the 1st word in “the cat sat on mat” or the 1st word in “transformers are awesome” or any other sequence. it’s just telling the model, “hey, this word is at position 1, this word is at position 2” and so on.

ok so now we have two different embeddings for each word:

the word embedding from our embedding model (what the word means)
the positional embedding we just calculated (where the word is)

what do we do with them? super simple - we just add them together! that’s literally it.

for each word, we take its word embedding vector, add the positional embedding vector for its position, and that combined vector becomes our input to the attention layers.

why does addition work here? well, both embeddings are vectors, and vector addition is a standard mathematical operation for combining two entities into a new intermediate representation. when you add two vectors, you create a new direction and magnitude that carries information from both sources. think of it geometrically - the word “cat” points in some direction in the embedding space (capturing its meaning), and position 3 points in another direction (capturing location info). adding them gives you a new vector that encodes both pieces of information.

during training, the model learns to interpret these combined vectors and somehow extract both the meaning of words and their positions when needed. it’s one of those “surprisingly simple but effective” tricks in deep learning that just works!

how it actually works (the full picture)

alright now let’s go through the whole attention mechanism step by step. fasten your seat belt guys, because there’s more to it than just “queries, keys, values”.

before getting into this, I just want to clarify - the original attention paper used 512 as the size of the embedding passed into the model, and there are 8 heads with each head size being 64. there are some other defaults too, but for now just keep these in mind and we’re good to go.

another thing - when I say word, I mean unigram token. there are certain architectures where the tokens can be bigram or maybe n-gram (depends on the architecture we use).

step 0: starting point

you have your input - let’s say “cat sat mat” (3 words). each word gets converted to an embedding vector (let’s say dimension 512), and positional encodings are added. so now you have three vectors, let’s call them x₁, x₂, x₃, each of size 512.

step 1: creating queries, keys, values through projection

here’s the thing people often skip - Q, K, V aren’t just magically there. you CREATE them by multiplying your input by learned weight matrices. so basically you have three matrices: W_Q, W_K, W_V (each of size 512 × 64)

for each position i:

q_i = x_i × W_Q  (this is your query)
k_i = x_i × W_K  (this is your key)
v_i = x_i × W_V  (this is your value)

these matrices are learned during training as different tasks need different kinds of attention, so the model learns what transformations produce useful queries, keys, and values through learning those weights.

step 2: compute attention scores (scaled dot product)

okay quick recap of where we are: we started with word embeddings, added positional encodings, and now we have our input matrix. each word is a vector. then we projected these through weight matrices (W_Q, W_K, W_V) to create queries, keys, and values for each token.

now here’s the thing - when we talk about multi-head attention (which we’ll get to in detail soon), each head works with a smaller dimension. so if your full embedding is 512 dimensions and you have 8 heads, each head gets 512/8 = 64 dimensions. that’s what d_k represents - the dimension per head. for now we’re basically zooming in on one attention head. we’ll come to multi head later, but for now, just consider there is one head with 64 different features it can see for a given sequence.

so let’s say we’re focusing on position 2 (word “sat”) in one of these heads.

its query is q₂ (a 64-dimensional vector). we compute dot products with all keys:

score₁ = q₂ · k₁ (how much should "sat" attend to "cat")
score₂ = q₂ · k₂ (how much should "sat" attend to itself)
score₃ = q₂ · k₃ (how much should "sat" attend to "mat")

and now you get some scores.

these scores give an impression of how much one word is important to another in terms of attention. higher score = more relevant. but here’s the critical part which often doesn’t get enough recognition - we scale these scores by dividing by √d_k, where d_k is the dimension of the keys (64 in our example, since we’re looking at one head).

why? because without scaling, when d_k is large (think about the big architectures which have more dimensions per head), the dot products get really big. think about it - you’re adding up 64 products of numbers. the more numbers you add, the bigger the sum gets on average. and when you put really big numbers into softmax, it becomes super peaky - basically all attention goes to one position (like 0.99, 0.005, 0.005). in that case, it’s like saying only this position matters in the whole sequence, which is kinda wrong right? that’s where dividing by √d_k keeps the variance stable and the softmax reasonable. this is a bit tricky at first glance but when you think about it for some time, you will actually give more respect to this small operation. i posted more about this on my x, would be helpful if you want to dig a bit deeper.

ok so now the actual scores are:

score₁ = (q₂ · k₁) / √64 = (q₂ · k₁) / 8
score₂ = (q₂ · k₂) / 8
score₃ = (q₂ · k₃) / 8

this scaled dot-product attention is the core mechanism. simple but super effective.

step 3: softmax to get attention weights

now we run these scaled scores through softmax:

attention_weights = softmax([score₁, score₂, score₃])

what softmax does is turn any set of numbers into a probability distribution - all the scores after passed through softmax will sum to 1. so if our scaled scores were [1.2, 0.5, 1.1] for q2 which was “sat”, then softmax might give us [0.42, 0.16, 0.42]. this means “sat” should pay 42% attention to “cat”, 16% to itself, and 42% to “mat”. but just to put a note here - softmax is kinda dumb, guys!

step 4: weighted sum of values

alright, now comes the payoff time. we have attention weights that tell us how much to focus on each position. and we have value vectors from each position that contain the actual information to pass along.

think about it this way - the query and key vectors were just for figuring out relevance. query says “what am i looking for?” and key says “what kind of info do i represent?”. they just compute the attention scores. but the value vector? that’s the actual content. that’s “here’s what i can actually give you”.

like if you’re searching for “best pizza in brooklyn”, the query is your search, the keys are like those titles/descriptions that help you find relevant results, but the values are the actual articles/reviews you end up reading. the keys helped you find what’s relevant, but the values are the real information you take away.

now combining this, we compute the output as a weighted average of these value vectors using the attention weights we got earlier:

output₂ = 0.42 × v₁ + 0.16 × v₂ + 0.42 × v₃

what’s happening here? well to be honest we’re just mixing information from different positions based on what we calculated earlier. it’s like “sat” is grabbing 42% of the information from “cat”, 42% from “mat”, and only 16% from itself.

this output₂ is the new representation for position 2 after attention! look closely at how it started as just the word “sat”, but now it’s enriched with contextual information from “cat” and “mat”. that’s the magic - each word’s representation gets updated based on relevant context.

now just imagine this whole process happening for all positions at the same time. yeah, position 1’s query is attending to all keys and getting its weighted values, position 2 is doing the same, position 100 is doing it - everyone in parallel. each position gets a new contextual representation based on how much that word should attend to all the other words. just read it once more! this is literally it, the whole gist of attention.

this is huge for speed btw. remember how rnns had to process word 1, then word 2, then word 3 sequentially? well here, every position computes its attention simultaneously. you can throw this at a gpu and it just parallelizes everything. that’s why transformers are so much faster too.

the full picture in one formula:

time to see everything at once:

Attention(Q, K, V) = softmax(QK^T / √d_k)V

now as i said earlier, imagine this attention thing happening for all positions at the same time. these are just matrices after all. Q is all your queries stacked as rows (one row per position), K is all your keys stacked as rows, V is all your values stacked as rows.

when you do QK^T, you’re computing every query with every key in one shot. it’s a matrix multiplication that gives you a grid of scores - position 1’s query against all keys, position 2’s query against all keys, and so on. you scale everything by √d_k, then apply softmax row-wise (each row becomes probabilities that sum to 1), then multiply by V to get all the weighted sums at once.

boom, you’ve computed attention for all positions in parallel. no sequential processing, just big matrix multiplications that gpus are really good at. you can process a 1000-word sequence and every position attends to every other position simultaneously. how cool is that!

multi-head attention (because one perspective isn’t enough)

okay so here’s where it gets even cooler. what if we had multiple attention mechanisms running in parallel, each learning to focus on different aspects? remember earlier I said something about the number of heads?

it’s like understanding a sequence through different lenses. one head might capture grammatical relationships between words (like subject-verb agreement), another head might be better at understanding semantic connections (like “bank” relating to “river” vs “bank” relating to “money”), another might focus on long-range dependencies, and yet another might specialize in local context.

just think about how you read a sentence yourself. you’re simultaneously tracking multiple things - sometimes it’s grammar or meaning, sometimes it’s references to earlier parts, tone, relationships between entities. you’re not just looking at it one way. multi-head attention does the same thing - it processes the sequence from multiple angles at once.

that’s multi-head attention. instead of one set of W_Q, W_K, W_V matrices, you have h sets (the paper uses h=8). so you have:

head 1: W_Q¹, W_K¹, W_V¹ → produces attention output₁
head 2: W_Q², W_K², W_V² → produces attention output₂
...
head 8: W_Q⁸, W_K⁸, W_V⁸ → produces attention output₈

each head is smaller (if total dimension is 512 and you have 8 heads, each head works with 64 dimensions). they all run the same attention mechanism, just with different learned projections.

just don’t confuse here - whatever we understood earlier about how attention works? that’s exactly what’s going on in all the other heads too. same queries, keys, values, same dot products, same softmax, same weighted sums. the computation is identical.

the only difference is instead of working with the whole 512-dimensional space, each head works on a smaller slice. if you have 8 heads, each one gets 64 dimensions (512/8). so head 1 might be working with dimensions 0-63, head 2 with dimensions 64-127, and so on. same attention mechanism, just on different parts of the representation.

and here’s the cool part - even though they’re all doing the same computation, they learn to specialize during training. maybe head 1 learns to focus on subject-verb relationships, head 2 picks up on spatial/location words, head 3 captures temporal markers like “yesterday” or “will”. the model figures out these specializations on its own just from the training objective. you don’t tell it what to look for - it discovers useful patterns.

combining the heads

after you’ve computed all 8 attention outputs, you concatenate them:

concat = [output₁, output₂, ..., output₈]

but you can’t just leave it like that - you need to mix the information from different heads together. so there’s one more learned matrix W_O that you multiply the concatenated heads by:

final_output = concat × W_O

this W_O matrix is really crucial guys - it lets the model learn how to combine the different perspectives from each head into one coherent representation. without it, the heads would just be sitting next to each other not really talking to each other. it’s like having different juices in different containers, but we want to drink a mix of them. that’s what W_O is for.

feed-forward network (the other half of the layer)

alright so attention is cool and all, but there’s another component in each transformer layer that often gets less spotlight - the feed-forward network.

but before we get there, we need to talk about two things that happen between attention and ffn: residual connections and layer normalization. these might sound boring but they’re actually crucial for making deep networks trainable.

residual connections (aka skip connections)

okay so remember the vanishing gradient problem from rnns? turns out when you stack many layers of any neural network, you can run into similar issues. when i say many layers here, i’m talking about the entire transformer block (attention + ffn and other linear projections) getting repeated multiple times. (more on ffn in a bit) the original paper stacks 6 of these blocks, GPT-3 has 96, and we’ll get into why and how this stacking works later. but for now, just know that if you have many of these blocks stacked on top of each other, during backprop the gradient gets weaker and weaker as it travels backwards.

because when you backpropagate, you’re using the chain rule. the gradient at block 1 depends on the gradient at block 2, which depends on block 3, and so on. so you end up multiplying a bunch of gradients together, one from each block. and if most of these multiplications involve numbers less than 1 (which they often do), you’re multiplying 0.8 × 0.7 × 0.9 × … and after a couple of blocks, that product gets really tiny.

it’s like the telephone game, but with math. each block adds a bit of “loss” to the signal. by the time the gradient reaches the early blocks, it’s so weak that those layers barely learn anything as the weights don’t update meaningfully because the gradient is basically zero.

well the fix is surprisingly simple: instead of just passing the attention output to the next step, you add it to the original input. so if x was your input to the attention layer, and attention(x) is the output, you actually use x + attention(x). that’s it, literally.

residual connections are like emergency exits for your neural network. when information passes through attention layers, some of it might get lost or distorted because there’s just too much dependency on layers, one after another.

but with residual connections, we basically say “hey, let’s also keep a copy of the original input and just add it back after the transformation.” this creates a direct highway for both information and gradients to flow through the network. so during training, even if the gradient signal gets weak going through the attention block, it can still travel through this shortcut path without any disturbance. that’s why we can stack these blocks 12+ layers deep without the network becoming untrainable.

and btw the math behind residual connections is beautifully simple (i know it’s too much but just stick with it…)

y = x + F(x)

where:
  x = your input (like the word embeddings + positional embeddings)
  F(x) = the transformation (attention layer)
  y = your output

during backpropagation, we calculate ∂L/∂x - the gradient of loss with respect to input:

∂L/∂x = ∂L/∂y × (1 + ∂F(x)/∂x)

this is where the magic happens - that “+1” term means even if ∂F(x)/∂x becomes tiny (like 0.0001) through many layers, we still have that constant 1 guaranteeing gradient flow.

with standard networks (y = F(x)), your gradient is just ∂L/∂y × ∂F(x)/∂x, which can shrink to almost nothing after many layers. but that +1 from residual connections creates an express lane for gradients that never gets blocked, no matter how deep you go.

layer normalization

after each sub-layer (attention or feed-forward), we apply something called layer normalization.

but what exactly is it doing?

here’s a simple way to think about layer normalization.

imagine you’re trying to teach a group of students. some are really loud talkers and others whisper. the loud ones dominate every conversation while the quiet ones get ignored. that’s basically what happens in neural networks without normalization.

in a neural network, different neurons can have wildly different “volumes” - some might output huge numbers like 1000, others tiny ones like 0.01. when this happens, the big numbers bulldoze through the network and the small ones get lost in the noise. it’s like having one person shouting in a room full of whispers. that’s a problem as this basically nullifies the relevance of the quieter neurons.

layer normalization is like giving everyone the same volume control. it takes all the outputs from a layer and says “hey, let’s make sure you’re all speaking at roughly the same level.” it doesn’t change what each neuron is trying to say, just normalizes how loudly they’re saying it.

why does this matter?! well, when all neurons are “speaking” at similar volumes, the network can actually learn from everyone. the gradients flow better during training, convergence happens faster, and the whole system becomes more stable. it’s like finally being able to hear that quiet person in the back who actually had great ideas all along…

without normalization, training can be painfully slow or just break entirely. with it, your network learns more efficiently and generalizes better. pretty neat for something so conceptually simple, right?

in our case, layer norm looks at each position’s vector (512 dimensions) and normalizes it across the feature dimension. it’s like saying “let’s make sure no single position has wildly different activation patterns than the others.”

mathematically, for each position’s vector, we:

calculate the mean (μ) and standard deviation (σ) across all feature dimensions
subtract the mean from each value: (x - μ)
divide by standard deviation: (x - μ)/σ
apply learnable scale (γ) and shift (β) parameters: γ × ((x - μ)/σ) + β

this is it. doing this makes sure the mean is 0 and the variance is 1 across the feature dimension.

in transformers specifically, layer norm is critical because attention weights can cause values to grow unpredictably. normalization ensures that even after 12+ layers, the activations don’t explode or vanish.

now to the feed-forward network

alright, so after multi-head attention → residual connection → layer norm, you’ve got a normalized output which is still a matrix of shape (sequence_length × 512).

now this goes into the feed-forward network. and here’s the key thing - the ffn processes each position independently. it’s applied to each row of this matrix separately. no mixing across positions here, just transforming each position’s representation on its own.

the structure is super simple:

ffn(x) = max(0, xW₁ + b₁)W₂ + b₂

let’s break down what’s actually happening with the dimensions.

input to ffn: one position's vector, size 512
W₁: weight matrix of size 512 × 2048 (expands to 4x larger)
after W₁ + bias + ReLU: you get a 2048-dimensional vector
W₂: weight matrix of size 2048 × 512 (compresses back down)
output: back to 512 dimensions

so you’re expanding each position’s representation to a much higher dimension (2048), applying non-linearity with ReLU, then compressing back to the original size. this happens identically for every position in the sequence. expanding to higher dimensions is like giving the model a bigger “workspace” to compute complex transformations that wouldn’t be possible in the original 512-dim space.

the higher dimensional space lets the network learn really cool patterns and relationships - like detecting subtle linguistic features, syntactic structures, or semantic associations that require multiple intermediate computations. it’s kind of like how you might need scratch paper to solve a complex math problem, right?! the 2048 dimensions give you that scratch space, then compress back to 512 so it can be passed to later layers for the next set of computations.

and yeah, W₁ and W₂ are learned during training just like the attention matrices. in fact, most of the parameters in a transformer are actually in these feed-forward layers, not the attention! if you count parameters, the FFN is usually where most of them live.

the full layer structure

so putting it all together, one transformer layer looks like:

multi-head attention (mix information across positions)
add & norm (residual connection + layer normalization)
feed-forward network (process each position independently)
add & norm again (another residual + layer norm)

then you stack like 6 or 12 or 96 of these layers depending on how big you want your model. more layers means more parameters to train, and each layer refines the representations further. the first layer might capture basic patterns, middle layers get more abstract relationships, and later layers handle really high-level understanding.

that’s it, that’s your attention thing!

the autoregressive thing (why decoders are slightly different)

okay so remember i said encoders and decoders are basically the same? here’s the one difference: when you’re generating text (like in GPT or any modern LLM), you can’t look at future words because they don’t exist yet. you’re predicting the next word based only on previous words.

so in decoder attention, when you’re at position 3 (say, the word “sat”), you can look at positions 1 and 2 (“the” and “cat”) and yourself (position 3), but not at positions 4 and 5 (“on” and “mat”) because those haven’t been generated yet.

you implement this with a mask. before the softmax, you set the attention scores for all future positions to negative infinity. after softmax, these become exactly zero. so those positions contribute nothing to the weighted sum.

same attention mechanism, just with this causal mask. this is why GPT-style models only use decoders - for text generation, you only need the causal version anyway.

encoders (like in BERT) let you look at the whole sentence in both directions, which is great for understanding tasks like “is this review positive or negative?” but can’t generate text word by word. decoders are one-directional, perfect for generation.

and honestly? most modern LLMs just use decoders now. turns out you don’t need the encoder for most things. the decoder-only architecture is simpler and scales better.

a quick note on kv-cache (the inference trick)

now here’s something interesting that becomes super important when you actually run these models in production. when generating text token by token, you’d normally have to recompute the keys and values for all previous tokens at every step. think about it - if you’ve generated 500 tokens and want to predict the 501st, you’d recompute K and V for all 500 tokens again. that’s wasteful as hell, right?

kv-cache is the simple but brilliant fix. since the keys and values for already generated tokens don’t change (they only depend on the input, not future tokens), you just… cache them. store them in memory. so when generating token 501, you only compute the new query for position 501, use the cached keys and values from positions 1-500, and compute just one new key-value pair to add to the cache.

this makes inference way faster, especially for long sequences. the tradeoff is memory - you’re storing these cached tensors which can get pretty large for long contexts. there’s a whole rabbit hole here about different caching strategies, memory optimization, and how this interacts with things like flash attention. but that’s a story for another blog. more on kv-cache and inference optimizations coming soon!

so that’s it. if you have come along this far reading everything, trust me, your attention is og (just like the attention paper, lol, bad joke, yeah!) nevermind, i hope this actually made you understand everything about attention. until next time, adios!

Create a custom retriever

2025-07-02T00:00:00+00:00

As retrieval-augmented generation (RAG) systems continue to evolve, the need for custom, domain-specific retrievers is becoming more and more obvious. Sure, traditional vector databases are great for basic similarity search—but the moment you throw in more complex, context-heavy queries, they start to fall short. Especially when you’re working with real-world data that needs richer filtering or semantic understanding.

In this post, we’ll walk through how to build a custom retriever by combining Superlinked—a high-performance vector compute engine—with LlamaIndex’s flexible retrieval framework. Think of this as a hands-on guide: we’ll show you how to wire things up end-to-end using a Steam game recommendation system as the example. We’ll tap into semantic similarity, multi-field indexing, and advanced query logic to make sure recommendations aren’t just relevant—they actually make sense.

Why Custom Retrievers Matter

Before we dive into the implementation, let’s talk about why building your own retriever is often the better choice in real-world RAG setups.

Tuned for Your Domain – Generic retrievers are fine for general use, but they tend to miss the subtle stuff. Think about jargon, shorthand, or domain-specific phrasing—those don’t usually get picked up unless your retriever knows what to look for. That’s where custom ones shine: you can hardwire in that context.
Works Beyond Just Text – Most real-world data isn’t just plain text. You’ll often have metadata, tags, maybe even embeddings from images or audio. For example, in a game recommendation system, we don’t just care about the game description—we also want to factor in genres, tags, user ratings, and more. Think about this logic: someone searching for a “strategy co-op game with sci-fi elements” won’t get far with text-only matching.
Custom Filtering and Ranking Logic – Sometimes you want to apply your own rules to how things are scored or filtered. Maybe you want to prioritize newer content, or penalize results that don’t meet certain quality thresholds. I mean, having that kind of control is like giving your retriever an actual brain—it can reason through relevance instead of just relying on vector distances.
Performance Gains – Let’s be real: general-purpose solutions are built to work “okay” for everyone, not great for you. If you know your data and your access patterns, you can fine-tune your retriever to run faster, rank better, and reduce unnecessary noise in the results.

Implementation Breakdown

Part 1: Core Dependencies and Imports

import pandas as pd
from typing import List
from llama_index.core.retrievers import BaseRetriever
from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.response_synthesizers import get_response_synthesizer
from llama_index.core import Settings
from llama_index.llms.openai import OpenAI
import superlinked.framework as sl

The import structure reveals our hybrid approach:

LlamaIndex Core: Provides the retrieval abstraction layer
Superlinked Framework: Handles vector computation and semantic search
Pandas: Manages data preprocessing and manipulation

Part 2: Understanding LlamaIndex Custom Retrievers

Before diving into our Superlinked implementation, it’s crucial to understand how LlamaIndex’s custom retriever architecture works and why it’s so powerful for building domain-specific RAG applications.

BaseRetriever Abstraction

LlamaIndex provides an abstract BaseRetriever class that serves as the foundation for all retrieval operations. The beauty of this design lies in its simplicity—any custom retriever only needs to implement one core method:

from abc import abstractmethod
from llama_index.core.retrievers import BaseRetriever
from llama_index.core.schema import NodeWithScore, QueryBundle

class BaseRetriever:
    @abstractmethod
    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """Retrieve nodes given query."""
        pass

The moat here is the presence of the Retrieval Protocol from the LlamaIndex. As this “retrieval protocol” makes it easy to plug in different backends or strategies without having to touch the rest of your system. Let’s break it down on what’s exactly is going on :

Input: QueryBundle

This is the query object passed into your retriever. At minimum, it contains the user’s raw query string (e.g., “sci-fi strategy games”). But it can also include extra metadata—like filters, embeddings, or user preferences. Basically, anything that might help shape a more relevant response.
Output: List[NodeWithScore]

The retriever returns a list of nodes—these are your chunks of content, documents, or data entries—each paired with a relevance score. The higher the score, the more relevant the node is to the query. This list is what gets passed downstream to the LLM or other post-processing steps. As in our case, we are plugging on the
Processing: Backend-Agnostic

Here’s the cool part: how you get from query to result is totally up to you. You can use a vector database, a traditional search engine, a REST API, or even something handcrafted for your specific use case. This decouples logic and gives you full control over the retrieval stack.

But Why This Matters?

This abstraction isn’t just clean—it’s powerful. It means you can:

Combine multiple strategies – Use dense vector search and keyword filtering together if needed.
Run A/B tests easily – Compare different retrievers to see what gives better results for your users.
Plug into any agent or tool – Whether you’re building a chatbot, a search UI, or a full-blown agent system, this retriever interface slots in easily.

Think of the retrieval protocol as the API contract between your “retrieval brain” and everything else. Once you follow it, you’re free to innovate however you want behind the scenes.

Plugging Superlinked into LlamaIndex

Now let’s see how we bridge Superlinked’s advanced vector computation capabilities with LlamaIndex’s retrieval abstraction:

class SuperlinkedSteamGamesRetriever(BaseRetriever):
    """A custom LlamaIndex retriever using Superlinked for Steam games data."""
    
    def __init__(self, csv_file: str, top_k: int = 10):
        self.top_k = top_k
        self.df = pd.read_csv(csv_file)
        
        # Critical: Text field combination for semantic richness
        self.df['combined_text'] = (
            self.df['name'].astype(str) + " " +
            self.df['desc_snippet'].astype(str) + " " +
            self.df['genre'].astype(str) + " " +
            self.df['game_details'].astype(str) + " " +
            self.df['game_description'].astype(str)
        )
        
        self._setup_superlinked()

Integration Architecture Deep Dive:

Part 3: Superlinked Schema Definition and Setup

def _setup_superlinked(self):
    """Set up Superlinked schema, space, index, and executor."""

    # Schema Definition - The Foundation of Type Safety
    class GameSchema(sl.Schema):
        game_number: sl.IdField
        name: sl.String
        desc_snippet: sl.String
        game_details: sl.String
        languages: sl.String
        genre: sl.String
        game_description: sl.String
        original_price: sl.Float
        discount_price: sl.Float
        combined_text: sl.String  # Our semantic search target

In Superlinked, the schema isn’t just about defining data types— it’s more like a formal definition between our data and the underlying vector compute engine. This schema determines how our data gets parsed, indexed, and queried—so getting it right is crucial.

In our SuperlinkedSteamGamesRetriever, the schema is defined like this:

class GameSchema(sl.Schema):
    game_number: sl.IdField
    name: sl.String
    desc_snippet: sl.String
    game_details: sl.String
    languages: sl.String
    genre: sl.String
    game_description: sl.String
    original_price: sl.Float
    discount_price: sl.Float
    combined_text: sl.String  # Key field for semantic search

Let’s break down what each of these elements actually does:

sl.IdField (→ game_number)

Think of this as our primary key. It gives each game a unique identity and allows Superlinked to index and retrieve items efficiently, I mean basically it’s about how we are telling the Superlinked to segregate the unique identify of the games, and btw it’s especially important when you’re dealing with thousands of records.
sl.String and sl.Float

Now these aren’t just type hints—they enable Superlinked to optimize operations differently depending on the field. For instance, sl.String fields can be embedded and compared semantically, while sl.Float fields can support numeric filtering or sorting.
combined_text

This is the semantic anchor of our retriever. It’s a synthetic field where we concatenate the game name, description, genre, and other relevant attributes into a single block of text. This lets us build a single text similarity space using sentence-transformer embeddings:
```
  self.text_space = sl.TextSimilaritySpace(
      text=self.game.combined_text,
      model="sentence-transformers/all-mpnet-base-v2"
  )
    
```
Why do this? Because users don’t just search by genre or name—they describe what they’re looking for. By embedding all the important signals into combined_text, we can better match fuzzy, natural-language queries with the most relevant games.

This schema definition is what allows Superlinked to parse our dataset, build the index, and execute semantic queries—all while retaining full control over how data is structured and ranked.

Part 4: Vector Space Configuration

    # Create text similarity space using sentence transformers
    self.text_space = sl.TextSimilaritySpace(
        text=self.game.combined_text,
        model="sentence-transformers/all-mpnet-base-v2"
    )

    # Create index
    self.index = sl.Index([self.text_space])

To power the semantic search over our Steam games dataset, we made two intentional design choices that balance performance, simplicity, and flexibility.

First, for the embedding model, we selected all-mpnet-base-v2 from the Sentence Transformers library. This model produces 768-dimensional embeddings that strike a solid middle ground: they’re expressive enough to capture rich semantic meaning, yet lightweight enough to be fast in production. I mean it’s a reliable general-purpose model, known to perform well across diverse text types — which matters a lot when your data ranges from short genre tags to long-form game descriptions. In our case, we needed a model that wouldn’t choke on either end of that spectrum, and all-mpnet-base-v2 handled it cleanly.

Next, although Superlinked supports multi-space indexing — where you can combine multiple fields or even modalities (like text + images) — we deliberately kept things simple with a single TextSimilaritySpace. I would have included the RecencySpace in here too but we don’t have the information on the release date for the games. But just to put this out here, if we have the release date information, we could plug in the RecencySpace here, and we can even sort the games with the TextSimilaritySpace along with the Recency of the games. Cool..

Now for the current TextSimilaritySpace details, this was built on a custom field we created by concatenating the name, genre, description, and other textual fields for each game. The idea was straightforward: keep the query logic minimal, and let the model’s semantic power do the heavy lifting.

Part 5: Data Pipeline and Executor Setup

    # Map DataFrame columns to schema - Critical for data integrity
    parser = sl.DataFrameParser(
        self.game,
        mapping={
            self.game.game_number: "game_number",
            self.game.name: "name",
            self.game.desc_snippet: "desc_snippet",
            self.game.game_details: "game_details",
            self.game.languages: "languages",
            self.game.genre: "genre",
            self.game.game_description: "game_description",
            self.game.original_price: "original_price",
            self.game.discount_price: "discount_price",
            self.game.combined_text: "combined_text"
        }
    )

    # Set up in-memory execution environment
    source = sl.InMemorySource(self.game, parser=parser)
    self.executor = sl.InMemoryExecutor(sources=[source], indices=[self.index])
    self.app = self.executor.run()

    # Load and process data
    source.put([self.df])

At the heart of our retrieval system is a streamlined pipeline built for both clarity and speed. We start with the DataFrameParser, which serves as our ETL layer. It ensures that each field in the dataset is correctly typed and consistently mapped to our schema — essentially acting as the contract between our raw CSV data and the Superlinked indexing layer.

Once the data is structured, we feed it into an InMemorySource, which is ideal for datasets that comfortably fit in memory . This approach keeps everything lightning-fast without introducing storage overhead or network latency. Finally, the queries are handled by an InMemoryExecutor, which is optimised for sub-millisecond latency. This is what makes Superlinked suitable for real-time applications like interactive recommendation systems, where speed directly impacts user experience.

Part 6: The Retrieval Engine

def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
    """Core retrieval logic with semantic similarity scoring."""

    query_text = query_bundle.query_str

    # Construct Superlinked query with explicit field selection
    query = (
        sl.Query(self.index)
        .find(self.game)
        .similar(self.text_space, query_text)
        .select([
            self.game.game_number,
            self.game.name,
            self.game.desc_snippet,
            self.game.game_details,
            self.game.languages,
            self.game.genre,
            self.game.game_description,
            self.game.original_price,
            self.game.discount_price
        ])
        .limit(self.top_k)
    )

    # Execute and convert to pandas for processing
    result = self.app.query(query)
    df_result = sl.PandasConverter.to_pandas(result)

One of the things that makes Superlinked genuinely enjoyable to work with is its fluent-style query builder. If you’ve used libraries like SQLAlchemy or Django ORM, the pattern will feel familiar — each method in the chain adds clarity instead of clutter. In our case, the query starts by selecting the relevant index and defining the similarity search using the .similar() method, which computes cosine similarity in the embedding space. This is what allows us to retrieve semantically close games based on the user’s natural language query.

Another thoughtful design decision we made was to explicitly select the fields we care about in the result set — rather than doing something like SELECT *. This might sound minor, but it keeps the data lean, reduces processing overhead, and ensures we’re not passing around unnecessary payload during post-processing. Think of it as precision over bulk — especially important when you’re moving data between components in a latency-sensitive pipeline.

Part 7: Result Processing and Node Creation

    # Convert to LlamaIndex NodeWithScore format
    nodes_with_scores = []
    for i, row in df_result.iterrows():
        text = f"{row['name']}: {row['desc_snippet']}"
        metadata = {
            "game_number": row["id"],
            "name": row["name"],
            "desc_snippet": row["desc_snippet"],
            "game_details": row["game_details"],
            "languages": row["languages"],
            "genre": row["genre"],
            "game_description": row["game_description"],
            "original_price": row["original_price"],
            "discount_price": row["discount_price"]
        }

        # Simple ranking score based on result position
        score = 1.0 - (i / self.top_k)
        node = TextNode(text=text, metadata=metadata)
        nodes_with_scores.append(NodeWithScore(node=node, score=score))

    return nodes_with_scores

Now once we receive the results from Superlinked, we transform them into a format that plays well with LlamaIndex. First, we construct a human-readable text string by combining the game’s name with its short description — this becomes the content of each node, making it easier for the language model to reason about. It’s a small touch, but it really improves how relevant and understandable the retrieved data is when passed to the LLM.

Next, we make sure that all original fields from the dataset — including things like genre, pricing, and game details — are retained in the metadata. This is crucial because downstream processes might want to filter, display, or rank results based on this information. We don’t want to lose any useful context once we start working with the retrieved nodes.

Finally, we apply a lightweight score normalisation strategy. Instead of relying on raw similarity scores, we assign scores based on the position of the result in the ranked list. This keeps things simple and consistent — the top result always has the highest score, and the rest follow in descending order. It’s not fancy, but it gives us a stable and interpretable scoring system that works well across different queries.

Show Time: Executing the pipeline

Now that all components are in place, it’s time to bring our Retrieval-Augmented Generation (RAG) system to life. Below is the end-to-end integration of Superlinked and LlamaIndex in action.

# Initialize the RAG pipeline
Settings.llm = OpenAI(api_key="your-api-key-here")
retriever = SuperlinkedSteamGamesRetriever("sampled_games.csv", top_k=5)
response_synthesizer = get_response_synthesizer()

query_engine = RetrieverQueryEngine(
    retriever=retriever,
    response_synthesizer=response_synthesizer
)

# Execute a query
query = "I want to find a magic game with spells?"
response = query_engine.query(query)
print(response)

This setup combines our custom semantic retriever with a response synthesizer backed by an LLM. The query is seamlessly passed through the retrieval pipeline and returned with a meaningful response.

Real Estate Natural Language Querying Agent

2025-06-07T00:00:00+00:00

Real estate search is straight-up broken. I mean we all know the drill right, It’s just too much dropdowns, rigid filters, and somehow we still end up scrolling through properties that miss the mark completely. Want a “family-friendly home near good schools”? Good luck translating that into checkboxes and price sliders. I mean imagine there is a system which can just translate these queries to something meaningful reference and give us better results.

Taking into that account, I created an agentic Natural Language Querying (NLQ) system that actually understands what people want. Just tell it “Find me a 3-bedroom home in Toronto under $600,000 with 2 bedrooms and 3 bathrooms” and it figures out the rest. I mean this is one of the coolest thing right..

Here’s how I combined vector search with AI agents to create a real estate assistant that actually gets it. If you want to follow along, refer to this colab

The Problem with Current Real Estate Search

Every real estate platform treats search like you’re querying a database. It’s all structured filters and rigid categories. Need a family home at a place? Here’s your homework:

Pick a price range from a dropdown
Select bedrooms and bathrooms
Choose your city from a list
Cross your fingers that “good schools” somehow correlates with your selected neighborhood

This approach fails because real estate decisions are nuanced as hell. It’s super hard to translate the relevant requirements to the dropdowns and rigid filters. When someone says “affordable family home,” they could mean totally different things based on their income, family size, and what they consider “affordable.” I mean we need something where we can just put out the queries and get the relevant results. That’s where NLQ comes into play.

What’s Natural Language Querying All About?

NLQ flips the script. Instead of forcing users to translate their thoughts into database queries, it lets them just… talk normally.

Old way: Price: $400K-$600K, Bedrooms: 3+, Location: Toronto

New way: "Find me a 3-bedroom home in Toronto under $600,000"

The system handles the translation from human language to structured queries while actually understanding context and intent.

What it took to made it..

Well there are couple of the different components that I glued together which provides NLQ capabilities, semantic search, and contextual reasoning.

Superlinked: Well this is the most important aspect of making this happen, I mean this is the one which provides the core vector search and NLQ interfaces for retrieving relevant results based on user queries.
LLM: Well for this project, I used OpenAI GPT-4 which powers both the intent classification and generative outputs such as recommendations, insight narratives, and query reformulations.
For the embedding generation I used sentence-transformers/all-mpnet-base-v2

You can use any LLM or any embedding model of your choice based on the requirements.

Data Requirements

To demonstrate how the NLQ system and agentic flow work in practice, I used a dataset that reflects what you’d typically find in real-world property listings. This includes property-level attributes like price, number of bedrooms and bathrooms, and address, along with geographic fields such as city and province.

To support more context-aware reasoning, the data also includes demographic details like population size and median household income—useful for queries that imply lifestyle or investment considerations.

How I Built This Thing

Step 1: Vector Spaces for Property Magic

I created specialised vector spaces for different property characteristics. Think of these as smart ways to understand fuzzy concepts like “near downtown” or “family-friendly.”

# Text similarity space for addresses
address_space = sl.TextSimilaritySpace(
    text=real_estate.address,
    model="sentence-transformers/all-mpnet-base-v2"
)

# Numerical similarity spaces
price_space = sl.NumberSpace(
    number=real_estate.price,
    min_value=0,
    max_value=int(max_price),
    mode=sl.Mode.MINIMUM
)

beds_space = sl.NumberSpace(
    number=real_estate.number_beds,
    min_value=0,
    max_value=int(max_beds),
    mode=sl.Mode.MINIMUM
)

baths_space = sl.NumberSpace(
    number=real_estate.number_baths,
    min_value=0,
    max_value=int(max_baths),
    mode=sl.Mode.MINIMUM
)

population_space = sl.NumberSpace(
    number=real_estate.population,
    min_value=0,
    max_value=int(max_population),
    mode=sl.Mode.MINIMUM
)

income_space = sl.NumberSpace(
    number=real_estate.median_family_income,
    min_value=0,
    max_value=int(max_median_family_income),
    mode=sl.Mode.MINIMUM
)

# Categorical similarity spaces
city_space = sl.CategoricalSimilaritySpace(
    category_input=real_estate.city,
    categories=unique_city_categories
)

province_space = sl.CategoricalSimilaritySpace(
    category_input=real_estate.province,
    categories=unique_province_categories,
    negative_filter=-1.0,
    uncategorized_as_category=False
)

Why vector spaces? They enable fuzzy matching, so when someone searches “near downtown” the system intelligently finds relevant addresses without needing exact geolocation data. Numerical spaces compare property details like price, bedrooms, bathrooms, population, and income, prioritising user preferences. Categorical spaces ensure precise matching for cities and provinces, filtering out irrelevant locations. It’s all about semantic understanding and flexible, accurate comparisons.

Step 2: Weighted Query Construction

Weighted queries prioritize what matters most in searches. As weights defines how much importance that parameter have when we fetch the relevant results. Higher weights, like for price, ensure results match user priorities, like sticking to a budget.

query = sl.Query(index, weights={
    address_space: sl.Param("address_weight", default=1.0),
    price_space: sl.Param("price_weight", default=100.0),
    city_space: sl.Param("city_weight", default=1.0),
    beds_space: sl.Param("beds_weight", default=1.0),
    baths_space: sl.Param("baths_weight", default=1.0),
    population_space: sl.Param("population_weight", default=0.2),
    income_space: sl.Param("income_weight", default=1.0),
})

Notice price gets 100x weight? That’s because budget is usually the make-or-break factor. No point showing someone a $2M mansion when they said “affordable.” Crazy laugh here…

Why Agents + NLQ for Real Estate Search?

Before we get into the tools and how they simplify everything, let’s take a quick look at why pairing agents with Natural Language Querying (NLQ) leads to such a powerful search experience.

NLQ helps machines speak our language.

When someone says, “affordable family homes under $250,000,” NLQ knows how to break that down into structured search parameters and get the relevant results for us.
But on its own, NLQ hits a ceiling.

Sure, it’s great at filtering by the criteria as I mention, but it tends to fall short when:
- The results are too few or too literal, without understanding the intent behind the query.
- You’re looking for personalised suggestions that go beyond the filters. I mean just think when you have let’s say 10 different properties which are completing your criteria but you still need some kind of personalised recommendations.
- You want to compare multiple results in context.
That’s where agents step in.

Agents bring reasoning to the table. They can:
- Recognise when a search is too narrow and offer smarter alternatives.
- Read between the lines and pick up on what the user might actually want. So instead of just throwing the results, we can utilize these agents to make more of those findings.
- Handle more complex workflows—like comparing options or building narratives from raw data.

So when you blend NLQ (which listens and interprets) with agents (which think and act), you get something far more intelligent: a system that doesn’t just understand your request—it also knows how to respond meaningfully. That’s the sweet spot.

Step 3: The Agentic Layer - Where the Magic Happens

Ok, Here’s where it starts to get fun. Rather than building a one-size-fits-all (and honestly, kind of dumb) search tool, I designed an agentic system made up of five specialised tools—each one kicks in based on what the user is actually trying to do.

class RealEstateAgent:
    def __init__(self, df, app, master_query, openai_client):
        self.tools = [
            PropertyRetrievalTool(),      # Basic "show me homes" requests
            PropertyRecommendationTool(), # "What's good for me?" queries
            QueryRefinementTool(),        # When searches come up empty
            NarrativeInsightTool(),       # Market analysis stuff
            MultiStepQueryTool()          # "Compare Toronto vs Vancouver"
        ]

Each tool has a specific job. Let me break them down all of them one step at a time..

PropertyRetrievalTool

This one’s is used for the straightforward lookups. So if we type a query something like, “3-bedroom homes in Oshawa under $600K.” This tool parses that to our Superlinked Search engine, runs it through and—bam—we get relevant listings back. It’s quick, clean, and no fuss.

Behind the scenes, what it really do is it wraps the user query, hands it to the NLQ-powered engine, and formats the results for display. If nothing shows up, it gracefully pass the actions to another tool called QueryRefinementTool which takes care of the situation. More on this a bit later..

class PropertyRetrievalTool(Tool):
    def name(self) -> str:
        return "PropertyRetrievalTool"

    def description(self) -> str:
        return "Retrieves properties based on natural language queries, leveraging NLQ intelligence for ranking"

    def use(self, query: str, df: pd.DataFrame, app: Any, master_query: Any, openai_client: Any) -> str:
        try:
            results = app.query(master_query, natural_query=query, limit=10)
            df_results = results_to_dataframe(results, df)
            if len(df_results) == 0:
                # If there are no results found, pass it to the QueryRefinementTool to get the better response
                refine_tool = QueryRefinementTool()
                return refine_tool.use(query, df, app, master_query, openai_client)
            return f"Search Results for '{query}':\\n\\n{format_property_display(df_results, query, detailed=True)}"
        except Exception as e:
            return f"Error in property retrieval: {str(e)}"

PropertyRecommendationTool

This tool isn’t just matching filters—it’s more like a reasoning. When someone enters a query, it tries to understand their intent and preferences. For example, it might see that “affordable housing for a growing family” suggests a budget, space needs, and a preference for family-friendly areas.

It:

Uses GPT-4 to extract a lightweight user profile from the query.
Fetches relevant results with NLQ.
Then it generates top 3 recommendations with pros and cons, using both the data and inferred preferences.

It’s like having a real estate advisor who reads between the lines.

class PropertyRecommendationTool(Tool):
    def name(self) -> str:
        return "PropertyRecommendationTool"

    def description(self) -> str:
        return "Provides personalized property recommendations based on inferred user preferences"

    def use(self, query: str, df: pd.DataFrame, app: Any, master_query: Any, openai_client: Any) -> str:
        # Infer user profile from query
        profile_prompt = f"""
        Extract user preferences from this query, such as budget, family size, commute preferences, or lifestyle (e.g., urban, family-friendly).
        Query: "{query}"
        Return a JSON object with inferred preferences (e.g., {{"budget": 800000, "family_size": 4}}). If none, return {{}}.
        """
        try:
            profile_response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": profile_prompt}],
                temperature=0.3,
                max_tokens=100
            )
            user_profile = eval(profile_response.choices[0].message.content.strip()) or {}
        except Exception as e:
            user_profile = {}
            print(f"Error inferring user profile: {str(e)}")

        # Fetch properties using NLQ
        results = app.query(master_query, natural_query=query, limit=15)
        df_results = results_to_dataframe(results, df)
        if len(df_results) == 0:
            return "No properties found to recommend."

        # Prepare property data (top properties from NLQ ranking)
        properties_summary = [
            {
                'address': row['address'],
                'city': row['city'],
                'province': row['province'],
                'price': int(row['price']),
                'beds': int(row['number_beds']),
                'baths': int(row['number_baths']),
                'population': int(row['population']),
                'median_income': int(row['median_family_income']),
                'relevance': round(row['relevance_score'], 3)
            }
            for _, row in df_results.head(10).iterrows()
        ]

        # Generate recommendations
        prompt = f"""
        You are a real estate advisor. Provide recommendations based on:
        Query: "{query}"
        Inferred Profile: {user_profile}
        Properties (ranked by NLQ relevance): {properties_summary}

        The properties are already ranked by our intelligent NLQ system considering all factors including price relevance.
        Provide:
        1. Top 3 recommendations with reasons (include address) - respect the NLQ ranking
        2. Pros and cons for each
        3. Additional considerations (e.g., neighborhood, lifestyle fit)
        """
        try:
            response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=800
            )
            recommendations = response.choices[0].message.content.strip()

            # Use top properties from NLQ ranking (no additional sorting)
            formatted_properties = format_property_display(df_results.head(5), query, detailed=True)
            return f"Recommendations for '{query}':\\n\\n{recommendations}\\n\\nProperty Listings (NLQ Ranked):\\n{formatted_properties}"
        except Exception as e:
            return f"Error generating recommendations: {str(e)}"

QueryRefinementTool

Sometimes people are too specific and get little (or nothing) back. That’s where this tool steps in. If a query returns too few results, it suggests a smart alternative—maybe relaxing price, broadening the location, or tweaking filters.

It then reruns the new query and returns the refined results. It’s especially useful when users aren’t sure how to phrase their search in the most productive way. That’s where this comes into play..

class QueryRefinementTool(Tool):
    def name(self) -> str:
        return "QueryRefinementTool"

    def description(self) -> str:
        return "Refines queries when results are sparse or irrelevant"

    def use(self, query: str, df: pd.DataFrame, app: Any, master_query: Any, openai_client: Any) -> str:
        # Try initial query
        results = app.query(master_query, natural_query=query, limit=10)
        df_results = results_to_dataframe(results, df)

        if len(df_results) >= 3:
            return f"Search Results for '{query}':\\n\\n{format_property_display(df_results, query, detailed=True)}"

        # Refine query if fewer than 3 results
        refine_prompt = f"""
        The query "{query}" returned {len(df_results)} results. Suggest an alternative query to get more relevant results.
        Consider relaxing filters (e.g., price, location) or expanding scope (e.g., nearby cities).
        Return a single alternative query as a string.
        """
        try:
            response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": refine_prompt}],
                temperature=0.3,
                max_tokens=50
            )
            new_query = response.choices[0].message.content.strip()

            # Try refined query
            results = app.query(master_query, natural_query=new_query, limit=10)
            df_new_results = results_to_dataframe(results, df)

            if len(df_new_results) == 0:
                return f"No properties found for '{query}'. Suggested query '{new_query}' also returned no results."

            return f"""
Original Query: '{query}' returned {len(df_results)} results.
Suggested Query: '{new_query}'
Results for '{new_query}' (NLQ Ranked):
{format_property_display(df_new_results, new_query, detailed=True)}
            """
        except Exception as e:
            return f"Error refining query: {str(e)}"

NarrativeInsightTool

This tool’s for the data-savvy users—the investors, the analysts, the curious minds.It takes a query and returns a story: average prices, price ranges, bedroom counts, top cities, median incomes, Then it wraps all that data into a readable narrative: market trends, investment potential, and key takeaways based on the context.

You get both the numbers and the story behind them. So instead of just seeing the relevant results in the tabular format, this tool gives the results a more readable format.

class NarrativeInsightTool(Tool):
    def name(self) -> str:
        return "NarrativeInsightTool"

    def description(self) -> str:
        return "Provides narrative insights for complex queries (e.g., investment potential)"

    def use(self, query: str, df: pd.DataFrame, app: Any, master_query: Any, openai_client: Any) -> str:
        results = app.query(master_query, natural_query=query, limit=10)
        df_results = results_to_dataframe(results, df)
        if len(df_results) == 0:
            return "No properties found for analysis."

        stats = {
            'avg_price': int(df_results['price'].mean()),
            'median_price': int(df_results['price'].median()),
            'price_range': (int(df_results['price'].min()), int(df_results['price'].max())),
            'avg_beds': round(df_results['number_beds'].mean(), 1),
            'avg_baths': round(df_results['number_baths'].mean(), 1),
            'cities': df_results['city'].value_counts().head(3).to_dict(),
            'avg_income': int(df_results['median_family_income'].mean())
        }

        prompt = f"""
        Provide a narrative analysis for the query: "{query}"
        Data: {stats}
        Properties (NLQ ranked by relevance): {[{k: v for k, v in row.items() if k in ['address', 'city', 'price', 'number_beds', 'number_baths', 'relevance_score']} for _, row in df_results.head(5).iterrows()]}

        Note: Properties are already intelligently ranked by our NLQ system considering price and other factors.
        Include:
        1. Market overview
        2. Investment potential (if relevant)
        3. Key considerations (e.g., location, price trends)
        """
        try:
            response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=600
            )
            insights = response.choices[0].message.content.strip()
            formatted_properties = format_property_display(df_results, query, detailed=True)
            return f"Insights for '{query}':\\n\\n{insights}\\n\\nProperty Listings (NLQ Ranked):\\n{formatted_properties}"
        except Exception as e:
            return f"Error generating insights: {str(e)}"

MultiStepQueryTool

This one’s built for the power users—the folks asking things like, “Compare properties in Toronto vs. family homes in Ontario under $300K.” That’s not a simple query; it’s actually two or more requests bundled into one.

Instead of throwing its hands up, this tool does it :

It just break down the question into the manageable parts, something like more doable.
Now it runs those sub-query in the independent fashion using the correct retrieval logic.
Then comes the comparison where it just compares them side by side

It’s what makes the whole system feel flexible—like it understands layered intent instead of forcing users to rephrase things over and over.

class MultiStepQueryTool(Tool):
    def name(self) -> str:
        return "MultiStepQueryTool"

    def description(self) -> str:
        return "Handles queries requiring multiple NLQ calls (e.g., comparisons)"

    def use(self, query: str, df: pd.DataFrame, app: Any, master_query: Any, openai_client: Any) -> str:
        # Detect if query involves comparison
        compare_pattern = re.compile(r'\\b(compare|versus|vs\\.?)\\b', re.IGNORECASE)
        if not compare_pattern.search(query):
            return "Query does not require comparison."

        # Extract sub-queries
        prompt = f"""
        Split the query into sub-queries for comparison.
        Query: "{query}"
        Return a JSON list of sub-queries (e.g., ["3 bedroom homes in Toronto", "3 bedroom homes in Kitchener"]).
        """
        try:
            response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.3,
                max_tokens=100
            )
            sub_queries = eval(response.choices[0].message.content.strip())
            if not isinstance(sub_queries, list) or len(sub_queries) < 2:
                return "Unable to split query for comparison."

            # Run sub-queries (NLQ handles intelligent ranking)
            results_dict = {}
            for sub_query in sub_queries[:2]:  # Limit to 2 for simplicity
                results = app.query(master_query, natural_query=sub_query, limit=5)
                df_sub_results = results_to_dataframe(results, df)
                results_dict[sub_query] = df_sub_results

            # Generate comparison
            data_str = ""
            for sub_query, df_sub_results in results_dict.items():
                data_str += f"\\n{sub_query} (NLQ ranked):\\n"
                for _, row in df_sub_results.head(3).iterrows():
                    data_str += f"  - Address: {row['address']}, City: {row['city']}, Price: ${row['price']:,.0f}, Beds: {row['number_beds']}, Baths: {row['number_baths']}, Relevance: {row['relevance_score']:.3f}\\n"

            compare_prompt = f"""
            Compare properties from these sub-queries (each ranked by NLQ intelligence):
            {list(results_dict.keys())}
            Data:
            {data_str}

            Note: Each property list is already ranked by our NLQ system considering all factors including price relevance.
            Provide:
            1. Side-by-side comparison (price, beds, baths, NLQ relevance)
            2. Key differences (e.g., location, value, market positioning)
            3. Recommendations for different buyer types
            """
            response = openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": compare_prompt}],
                temperature=0.7,
                max_tokens=800
            )
            comparison = response.choices[0].message.content.strip()

            # Format results
            output = [f"Comparison for '{query}':\\n\\n{comparison}"]
            for sub_query, df_sub_results in results_dict.items():
                if not df_sub_results.empty:
                    output.append(f"\\nResults for '{sub_query}' (NLQ Ranked):\\n{format_property_display(df_sub_results, sub_query, detailed=True)}")

            return "\\n".join(output)
        except Exception as e:
            return f"Error processing comparison: {str(e)}"

Step 4: Intent Classification – Figuring Out What the User Really Wants

At this point, the system needs to make a smart decision on which tool we should invoke? Which tool is the best fit for this query?

To do that, it relies on GPT-4 to act like a traffic director. I mean more like a router which takes the query and classifies the intent behind it. Based on the outcome, it routes the query to the right tool:

def classify_query(self, query: str) -> str:
        valid_categories = {'retrieval', 'recommendation', 'refinement', 'insight', 'multistep'}
        prompt = f"""
Classify the query into one category based on its intent, returning only the category name ('retrieval' if unsure):
- retrieval: Search for properties by criteria (e.g., location, price, bedrooms).
- recommendation: Personalized property suggestions for user needs.
- refinement: Strict criteria likely yielding few results, needing adjustment.
- insight: Analysis, investment advice, or market trends.
- multistep: Comparison or multiple criteria across locations.
Query: "{query}"
        """
        try:
            response = self.openai_client.chat.completions.create(
                model="gpt-4",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.3,
                max_tokens=20
            )
            classification = response.choices[0].message.content.strip().lower()
            return classification if classification in valid_categories else 'retrieval'
        except Exception as e:
            print(f"Classification error: {str(e)}")
            return 'retrieval'

To understand how the system operates in practice, let’s walk through a few real-world queries and how they are handled end-to-end.

Example 1: Basic Search

Consider the query: “3 bedroom homes in Oshawa under $600,000.” The system classifies this as a retrieval task. It extracts structured parameters—location (Oshawa), number of bedrooms (3+), and price ceiling ($600,000)—and executes a vector-based search using these constraints.

The results are then ranked and returned based on semantic relevance and price alignment. Eassy pisssy..

Processing query: 3 bedroom homes in Oshawa under $600,000
Query classified as: retrieval
Search Results for '3 bedroom homes in Oshawa under $600,000':

  city  id province              address    price  number_beds  number_baths relevance_score
Oshawa 171  Ontario #70 -53 TAUNTON RD E $520,000            3             3           0.418
Oshawa 603  Ontario      145 BANTING AVE $499,900            6             5           0.417

Example 2: Personalised Recommendations

In this case, the query is: “I live in Toronto, show me the most expensive option that would be good for me.” The system classifies this under the recommendation category. It infers that the user is located in Toronto and likely has a preference for high-end listings. Now if this thing goes to the NLQ directly, the NLQ might not be able to make a ton of sense as we haven’t putted out too much of the filtration here, that’s where the agentic tooling helps.

So now the query pass through the router, then the router directs that to the relevant tool which is RecommendationTool in this case, and then that tool extracts additional implicit preferences (e.g., lifestyle, city familiarity, price range), fetches the top relevant properties using the user profile, and then generates recommendations with a short rationale, including pros and cons for each option.

THIS IS TOO MUCH HANDY MANNN

Processing query: I want to know which will be a good option for me if I live in Toronto, most expensive one
Query classified as: recommendation
Recommendations for 'I want to know which will be a good option for me if I live in Toronto, most expensive one':

1. Top 3 Recommendations:

   1.1. Property: #LPH19 -2095 LAKE SHORE BLVD W
   Reasons: This property is the most expensive on the list, priced at $4,336,900. It has 3 bedrooms and 4 bathrooms, which might be suitable if you're looking for a spacious property. It is also situated in a city with a high median income, which indicates a prosperous neighborhood.

   1.2. Property: 209 BELLWOODS AVE
   Reasons: This property is priced at $2,500,000. It has 4 bedrooms and 4 bathrooms, providing ample space. It's the second most expensive property on the list, which might be desirable if you're looking for a high-end property.

   1.3. Property: 43 WALMSLEY BLVD
   Reasons: This property is priced at $1,998,000. It has 3 bedrooms and 1 bathroom, which might be a good fit if you have a smaller family or prefer fewer bathrooms to maintain.

2. Pros and Cons:

   2.1. #LPH19 -2095 LAKE SHORE BLVD W
   Pros: It's the most expensive property available, indicating luxury and high-standard living. It has 3 bedrooms and 4 bathrooms, offering comfort and convenience.
   Cons: The high price might be beyond the budget of some people.

   2.2. 209 BELLWOODS AVE
   Pros: The property has 4 bedrooms and 4 bathrooms, offering ample space for a family. The price is also relatively high, suggesting a quality property.
   Cons: May be expensive for some buyers.

   2.3. 43 WALMSLEY BLVD
   Pros: The property is reasonably priced for its offerings. It has 3 bedrooms, which can be suitable for a family.
   Cons: The property has only 1 bathroom, which might not be enough if you have a larger family or frequently have guests.

3. Additional Considerations:

   3.1. Neighborhood: All properties are situated in Toronto, a vibrant city with a high median income. You might want to consider the specific neighborhood's safety, amenities, and community when making your decision.

   3.2. Lifestyle fit: Consider your lifestyle needs. If you frequently host guests, a property with more bedrooms and bathrooms might be preferred. Similarly, if you enjoy luxury living, the most expensive property might be more appealing.

   3.3. Future Resale Value: If you're considering the property as an investment, you might want to consider the potential resale value. Higher-end properties might have a better resale value.

Property Listings (NLQ Ranked):
   city  id province                        address      price  number_beds  number_baths relevance_score
Toronto   9  Ontario #LPH19 -2095 LAKE SHORE BLVD W $4,336,900            3             4          -0.624
Toronto 489  Ontario              209 BELLWOODS AVE $2,500,000            4             4          -0.677
Toronto 190  Ontario               43 WALMSLEY BLVD $1,998,000            3             1          -0.687
Toronto 195  Ontario                411 CONCORD AVE $1,859,000            3             3          -0.690
Toronto 357  Ontario              212 EDENBRIDGE DR $1,800,000            4             3          -0.690

The query “3 bedroom homes in Oshawa under $300,000” is classified as a retrieval request, but the result set is empty due to overly strict constraints. Instead of returning nothing, the system triggers a refinement process.

It automatically adjusts the query—suggesting a more realistic variant like “3 bedroom homes in Oshawa under $400,000”—and displays both the original and refined results. This helps the user quickly iterate without needing to rephrase their search.

Processing query: 3 bedroom homes in Oshawa under $300,000
Query classified as: retrieval

Original Query: '3 bedroom homes in Oshawa under $300,000' returned 0 results.
Suggested Query: '"3 bedroom homes in Durham Region under $400,000"'
Results for '"3 bedroom homes in Durham Region under $400,000"' (NLQ Ranked):
       city  id                  province                 address    price  number_beds  number_baths relevance_score
    Calgary 263                   Alberta   84 Whitehaven Road NE $329,000            3             2           0.793
 St. John's 544 Newfoundland and Labrador     28 Chapman Crescent $199,900            3             1           0.790
 St. John's 148 Newfoundland and Labrador 193 Cumberland Crescent $190,000            3             2           0.789
 St. John's  30 Newfoundland and Labrador       162 Bennetts Road $369,900            3             2           0.788
 St. John's 658 Newfoundland and Labrador 127 Commonwealth Avenue $295,000            3             2           0.788
Thunder Bay 723                   Ontario          71 Hull Avenue $329,900            3             1           0.787
     Regina 315              Saskatchewan    2914 Avonhurst DRIVE $265,000            3             2           0.787
     Regina 683              Saskatchewan          35 Hind STREET $159,900            3             1           0.786
     Regina 547              Saskatchewan  4609 Green Rock ROAD E $395,000            3             3           0.784
 Lethbridge 596                   Alberta          12 Lynx Road N $388,000            3             3           0.783

Example 4: Market Insight Generation

When a user asks: “What’s a good investment in Toronto?” the system identifies this as an insight-oriented query. It analyses the dataset to produce a narrative summary, including average prices, market trends, top-performing neighbourhoods, and potential investment signals.

The output combines raw statistics with interpreted insights to guide decision-making.

Processing query: What's a good investment in Toronto?
Query classified as: insight
Insights for 'What's a good investment in Toronto?':

Market Overview:
The Toronto property market is characterized by an average property price of $577,158, with a median price slightly higher at $589,000. The price range for properties in the market spans from $415,000 to $695,000. The average property consists of approximately 1.6 bedrooms and 1.1 bathrooms.

Investment Potential:
Given the average income in Toronto is $97,000, the property prices indicate a relatively high cost of living. This presents an opportunity for investors as there is a substantial market for both property purchasing and rental if buying to let. The properties in Toronto have been ranked by our NLQ system based on price and other factors, giving a relevance score. This score indicates the property's potential for return on investment considering these factors.

Key Considerations:
The top-ranked property by our NLQ system is located at '#1808 -1 KING ST W' with a price of $415,000. This property has 1 bedroom and 1 bathroom and it is the most affordable among the top-ranked properties. Other good investment options include properties located at '#201 -115 RICHMOND ST E', '#1608 -101 ERSKINE AVE', '#906 -1190 DUNDAS ST E', and '#2009 -24 WELLESLEY ST W' with prices ranging from $479,000 to $599,000.

It is important for potential investors to consider the location of the property, as properties in certain neighborhoods may offer higher potential returns due to demand and price trends. It is also crucial to consider the number of bedrooms and bathrooms, as these factors can affect both the rental income potential and the resale value.

Property Listings (NLQ Ranked):
   city  id province                  address    price  number_beds  number_baths relevance_score
Toronto  25  Ontario       #1808 -1 KING ST W $415,000            1             1           0.418
Toronto 914  Ontario  #201 -115 RICHMOND ST E $479,000            1             1           0.418
Toronto  91  Ontario   #1608 -101 ERSKINE AVE $539,900            1             1           0.418
Toronto 699  Ontario   #906 -1190 DUNDAS ST E $579,000            1             1           0.418
Toronto 592  Ontario #2009 -24 WELLESLEY ST W $599,000            2             1           0.418
Toronto 695  Ontario       #1606 -36 ZORRA ST $519,888            2             1           0.418
Toronto 492  Ontario #907 -185 ROEHAMPTON AVE $609,900            2             1           0.418
Toronto 164  Ontario     #325 -361 FRONT ST W $659,900            2             1           0.417
Toronto 649  Ontario     #213 -525 WILSON AVE $675,000            2             2           0.417
Toronto 706  Ontario      #2112 -101 PETER ST $695,000            2             1           0.417

Example 5: Comparative Search

In the query: “Compare 3 bedroom homes in Toronto and Kitchener,” the system recognizes the need for multi-step reasoning. It decomposes the query into two independent sub-queries—one for Toronto, one for Kitchener—runs each using the retrieval tool, and then formats the results into a side-by-side comparison.

Processing query: Compare 3 bedroom homes in Toronto and Kitchener
Query classified as: multistep
Comparison for 'Compare 3 bedroom homes in Toronto and Kitchener':

1. Side-by-side comparison:

   - 3 bedroom homes in Toronto:
     - Prices: $929,000, $848,000, $1,149,999
     - Beds: 3 for all
     - Baths: 2 for all
     - NLQ Relevance: 0.417, 0.417, 0.415

   - 3 bedroom homes in Kitchener:
     - Prices: $589,000, $649,900, $624,900
     - Beds: 3 for all
     - Baths: 1, 2, 2
     - NLQ Relevance: 0.419 for all

2. Key Differences:
   - Location: The key difference between these two sets of properties is their location, with one set located in Toronto and the other in Kitchener.
   - Value: The properties in Toronto are significantly more expensive than those in Kitchener, reflecting the higher real estate prices in Toronto.
   - Market Positioning: The properties in Toronto might be positioned towards buyers looking for homes in a bustling, cosmopolitan city, while the homes in Kitchener might attract those looking for a more relaxed, suburban lifestyle.

3. Recommendations for Different Buyer Types:
   - Buyers with a higher budget and a preference for city living might find the properties in Toronto more attractive. These properties offer the convenience of city living, with easy access to amenities and are ideal for those who work in the city.
   - On the other hand, buyers looking for more affordable options, or those who prefer a quieter, suburban lifestyle might prefer the properties in Kitchener. These homes offer great value for money and could also be ideal for families with children, given the lower cost of living.
   - For investors, both cities could offer good opportunities depending on their strategy. Toronto's high property prices could lead to high rental income, while Kitchener's lower prices could result in higher yield.

Results for '3 bedroom homes in Toronto' (NLQ Ranked):
   city  id province              address      price  number_beds  number_baths relevance_score
Toronto 508  Ontario    166 HIAWATHA RD W   $929,000            3             2           0.417
Toronto 869  Ontario         299 OSLER ST   $848,000            3             2           0.417
Toronto  54  Ontario         35 FOURTH ST $1,149,999            3             2           0.415
Toronto 813  Ontario 1042* ST CLARENS AVE $1,139,000            3             2           0.415
Toronto 659  Ontario    108 ARMSTRONG AVE $1,189,000            3             2           0.415

Results for '3 bedroom homes in Kitchener' (NLQ Ranked):
     city  id province                        address    price  number_beds  number_baths relevance_score
Kitchener 471  Ontario           204 SCHLUETER Street $589,000            3             1           0.419
Kitchener 918  Ontario 355 FISHER MILLS Road Unit# 75 $649,900            3             2           0.419
Kitchener 156  Ontario         24 TIMBERLANE Crescent $624,900            3             2           0.419
Kitchener 906  Ontario     400 WILSON Avenue Unit# 49 $599,000            3             3           0.419
Kitchener 371  Ontario      110 FERGUS Avenue Unit# 1 $790,000            3             2           0.418

Here is the concluded look of how the whole system looks under the hood..

Well before ending this one, there are couple of points I want your attention. So doing this end to end , we took care of a few core architectural choices that shaped how this system really works.

The most important one here is that I intentionally gave more weight to the Price rather than other filters when we were creating the query, as It acts as a hard constraint, while other attributes like location or amenities adjust the relevance within that boundary. I mean the Price is the most important factor I guess when we are looking out for the real estate.

And as I said earlier, our NLQ engine doesn’t depends on the structured filters, instead it understand the natural language directly.

An agent to read research papers better

2025-03-10T00:00:00+00:00

I’ve always wanted an AI assistant that could make it easier to go through a collection of research papers, answer my questions, summarize the content, and give me the important points. In short, I wanted a research helper that would make the whole process simpler. That’s when I decided to create it.

What’s the Plan?

I wanted to build an agentic system that could do three main things:

Find Papers: Search for research papers by topic (e.g., “quantum computing”) and rank them by relevance and recency.
Summarize: Crunch those papers into bite-sized insights fast when I asked it to do so.
Answer Questions: Pull answers from the specific papers which I tell it to do and based on that it answer those questions

In all of it, Superlinked is the backbone here, As it improves the vector search relevance by encoding metadata together with your unstructured data into vectors. This powers a vector database that understands both the text’s meaning and the time factor. I mean when we will stitch the schema design later on, we will be using RecencySpace which will literally help us to add the relevancy of the papers during the retrieval (ps: without using any rerankers) that means during the retrieval, if the two papers have the same kind of relevancy, it will prioritize paper which is more recent. More on this later.

For the LLM reasoning, I plan to use OpenAI’s models. However, you can easily replace it with any LLM of your choice. And last but not least, I made a Kernel Agent to handle the queries. Too much talking, let’s go now.. If you want to follow along with me, use this colab for the reference.

Step 1 : Setting up the toolbox

pip3 install openai pandas superlinked==19.21.1 sentence-transformers transformers 

To make things easier and more modular, I created an Abstract Tool class. This will simplify the process of building and adding tools

import pandas as pd
import superlinked.framework as sl
from datetime import timedelta
from openai import OpenAI
import os
from abc import ABC, abstractmethod
from tqdm import tqdm

# Abstract Tool Class for clean modularity
class Tool(ABC):
    @abstractmethod
    def name(self) -> str:
        pass

    @abstractmethod
    def description(self) -> str:
        pass

    @abstractmethod
    def use(self, *args, **kwargs):
        pass

# OpenAI setup
api_key = os.environ.get("OPENAI_API_KEY", "your-key-here")
if not api_key:
    raise ValueError("Set your OPENAI_API_KEY!")
client = OpenAI(api_key=api_key)
model = "gpt-3.5-turbo"

Step 2 : Data needs

What does the data look like? For this experiment, I used a dataset consisting of 10,000 AI research papers from Kaggle. To make it easy, simply run the cell below, and it will automatically download the dataset to your working directory. You can also replace this with your own data sources, whether they are research papers or other academic content. If you decide to do so, all you need to do is adjust the schema design slightly and update the column names.

import pandas as pd

!wget --no-check-certificate 'https://drive.google.com/uc?export=download&id=1FCR3TW5yLjGhEmm-Uclw0_5PWVEaLk1j' -O arxiv_ai_data.csv

For now, I’m going to work with just 1,000 papers—mainly to speed things up for the demo. SHHHHHH :)

There’s one more thing: I’ll convert the published data from string timestamps (like ‘1993-08-01 00:00:00+00:00’) into pandas datetime objects. This conversion is necessary because it allows us to perform date/time operations.

df = pd.read_csv('arxiv_ai.csv').head(1000)

# Convert to datetime but keep it as datetime (more readable and usable)
df['published'] = pd.to_datetime(df['published'])

# Ensure summary is a string
df['summary'] = df['summary'].astype(str)

# Add 'text' column for similarity search
df['text'] = df['title'] + " " + df['summary']

# Debug: Confirm columns in original DataFrame
print("Debug: Columns in original DataFrame:", df.columns.tolist())

Debug: Columns in original DataFrame: ['authors', 'categories', 'comment', 'doi', 'entry_id', 'journal_ref' 'pdf_url', 'primary_category', 'published', 'summary', 'title', 'updated']

Let me briefly explain the columns and what they contain, as this will be helpful later on:

published: The date the paper was published.
summary: The abstract of the paper.
entry_id: The arXiv entry ID for the paper.

For our experiment, the only columns we need are entry_id, published, title, and summary. To make things more effective, we can combine the title and the summary of each paper to create a comprehensive text, which will be the core of our work. So text column will contains the title and the summary

Just a bit on the Superlinked’s in-memory indexer. So basically I think it’s a game-changer for our small dataset of 1,000 papers. I mean it stores data in the computer’s memory, which means queries are lightning-fast—no waiting for disk access. This is perfect for real-time applications and prototyping, like our proof-of-concept.

Step 3 : Defining the Superlinked Schema

To move ahead, we will need a schema to map our data. I set up PaperSchema with key fields:

class PaperSchema(sl.Schema):
    text: sl.String
    published: sl.Timestamp
    entry_id: sl.IdField
    title: sl.String
    summary: sl.String

paper = PaperSchema()

Now, the most important part of the process is creating the spaces. There are two spaces we need to create for our schema. The first one is TextSimilaritySpace. This space plays a crucial role in organizing and searching through the research papers. Essentially, it converts the text data we provide into vectors, making it much easier to search through them. The TextSimilaritySpace is specifically designed to efficiently encode longer text sequences using specialized models, enabling effective similarity comparisons between different pieces of text

text_space = sl.TextSimilaritySpace(
    text=sl.chunk(paper.text, chunk_size=200, chunk_overlap=50),
    model="sentence-transformers/all-mpnet-base-v2"
)

On the other hand, RecencySpace focuses on the importance of time in the context of research papers. It encodes timestamps, prioritizing the recency of the papers. This ensures that more recent papers are given greater weight in the similarity calculations, allowing for a balance between the relevance of the content and the publication time of the papers.

recency_space = sl.RecencySpace(
    timestamp=paper.published,
    period_time_list=[
        sl.PeriodTime(timedelta(days=365)),  # 1 year
        sl.PeriodTime(timedelta(days=2*365)),  # 2 years
        sl.PeriodTime(timedelta(days=3*365)),  # 3 years
    ],
    negative_filter=-0.25  # Penalize ancient papers
)

Think of RecencySpace as a time-based filter, similar to sorting your emails by date or viewing Instagram posts with the newest ones first. It helps answer the question, ‘How fresh is this paper?’

Smaller timedeltas (like 365 days) allow for more granular, yearly time-based rankings.
Larger timedeltas (like 1095 days) create broader time periods.

The negative_filter penalizes very old papers. To explain it more clearly, consider the following example where two papers have identical content relevance, but their rankings will depend on their publication dates.

Paper A: Published in 1996 
Paper B: Published in 1993

Scoring example:
- Text similarity score: Both papers get 0.8
- Recency score:
  - Paper A: Receives the full recency boost (1.0)
  - Paper B: Gets penalized (-0.25 due to negative_filter)

Final combined scores:
- Paper A: Higher final rank
- Paper B: Lower final rank

These spaces are key to making the dataset more accessible and effective. They allow for both content-based and time-based searches, and really helpful in understanding the relevance and recency of research papers. This provides a powerful way to organize and search through the dataset based on both the content and the publication time.

Step 4 : Building the index

Next, I fused the spaces into an index, A.K.A the search engine’s core:

paper_index = sl.Index([text_space, recency_space])

Then, I mapped the DataFrame to the schema and loaded it in batches (10 papers at a time) into in-memory store:

parser = sl.DataFrameParser(
    paper,
    mapping={
        paper.entry_id: "entry_id",
        paper.published: "published",
        paper.text: "text",
        paper.title: "title",
        paper.summary: "summary",
    }
)

source = sl.InMemorySource(paper, parser=parser)
executor = sl.InMemoryExecutor(sources=[source], indices=[paper_index])
app = executor.run()

batch_size = 10
data_batches = [df[i:i + batch_size] for i in range(0, len(df), batch_size)]
for batch in tqdm(data_batches, desc="Loading Data"):
    source.put([batch])

The in-memory executor is why Superlinked shines here—1,000 papers fit snugly in RAM, and queries fly without disk I/O bottlenecks.

Step 5 : Crafting the query

Next, we move on to query creation. This is where we set up a template for crafting queries. To manage this, we need a query template that can balance both relevance and recency. Here’s what it looks like:

knowledgebase_query = (
    sl.Query(
        paper_index,
        weights={
            text_space: sl.Param("relevance_weight"),
            recency_space: sl.Param("recency_weight"),
        }
    )
    .find(paper)
    .similar(text_space, sl.Param("search_query"))
    .select(paper.entry_id, paper.published, paper.text, paper.title, paper.summary)
    .limit(sl.Param("limit"))
)

This lets me puts on how much I care about content (relevance_weight) versus freshness (recency_weight) — a killer combo for our agentic needs.

Step 6 : Building tools

Now comes the tooling part, here we go lads.

Ok , So We’ll be creating three tools for our agent to choose from. These tools will help find papers, summarize them if needed, and answer relevant questions.

Retrieval Tool : This tool is crafted by hooking into Superlinked’s index, letting it pull the top 5 papers based on a query. It balances relevance (1.0 weight) and recency (0.5 weight) to nail the “find papers” goal. I mean what we wanted is to find the papers which are relevant to the query. So if the query is something like, “What are the Quantum computers papers are published in b/w 1993 and 1994” , then the retrieval tool will retrieve those papers, will bring out the summary of those papers one by one, and will return the result.

class RetrievalTool(Tool):
    def __init__(self, df, app, knowledgebase_query, client, model):
        self.df = df
        self.app = app
        self.knowledgebase_query = knowledgebase_query
        self.client = client
        self.model = model

    def name(self) -> str:
        return "RetrievalTool"

    def description(self) -> str:
        return "Fetches papers with Superlinked’s vector search—built by wiring up the Superlinked index to query for topic matches and recent papers."

    def use(self, query: str) -> pd.DataFrame:
        result = self.app.query(
            self.knowledgebase_query,
            relevance_weight=1.0,
            recency_weight=0.5,
            search_query=query,
            limit=5
        )
        df_result = sl.PandasConverter.to_pandas(result)
        df_result['summary'] = df_result['summary'].astype(str)
        return df_result

Next up is the Summarization Tool. This tool is designed for cases where we need a concise summary of a paper. To use it, we’ll provide the paper_id, which is the ID of the paper we want summarized. Keep in mind that if you don’t provide the paper_id, the tool won’t work, as it requires these IDs to look up the corresponding papers in the dataset.

class SummarizationTool(Tool):
    def __init__(self, df, client, model):
        self.df = df
        self.client = client
        self.model = model

    def name(self) -> str:
        return "SummarizationTool"

    def description(self) -> str:
        return "Summarizes papers using GPT-3.5-turbo—put together with OpenAI’s API to crunch paper summaries by ID into quick insights."

    def use(self, query: str, paper_ids: list) -> str:
        papers = self.df[self.df['entry_id'].isin(paper_ids)]
        if papers.empty:
            return "No papers found."
        summaries = "\n\n".join(papers['summary'].tolist())
        prompt = f"Summarize these:\n\n{summaries}\n\nKeep it tight."
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.7,
            max_tokens=500
        )
        return response.choices[0].message.content.strip()

Finally, we have the QuestionAnsweringTool. This tool chains the RetrievalTool to fetch the relevant papers and then uses them to answer the questions. If no relevant papers are found to answer the questions, it will provide an answer based on general knowledge

class QuestionAnsweringTool(Tool):
    def __init__(self, retrieval_tool, client, model):
        self.retrieval_tool = retrieval_tool
        self.client = client
        self.model = model

    def name(self) -> str:
        return "QuestionAnsweringTool"

    def description(self) -> str:
        return "Answers questions about research topics using retrieved paper summaries or general knowledge if no specific context is available."

    def use(self, query: str) -> str:
        df_result = self.retrieval_tool.use(query)
        if 'summary' not in df_result.columns:
            # Tag as a general question if summary is missing
            prompt = f"""
            You are a knowledgeable research assistant. This is a general question tagged as [GENERAL]. Answer based on your broad knowledge, not limited to specific paper summaries. If you don't know the answer, provide a brief explanation of why.

            User's question: {query}
            """
        else:
            # Use paper summaries for specific context
            contexts = df_result['summary'].tolist()
            context_str = "\n\n".join(contexts)
            prompt = f"""
            You are a research assistant. Use the following paper summaries to answer the user's question. If you don't know the answer based on the summaries, say 'I don't know.'

            Paper summaries:
            {context_str}

            User's question: {query}
            """
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.7,
            max_tokens=500
        )
        return response.choices[0].message.content.strip()

Step 7 : Building the Kernel Agent

Next, we have the Kernel Agent. Whenever I build an agent-based system, I always include the Kernel Agent—it acts like the ‘boss.’ It really makes everything run smoothly. The Kernel Agent is the centerpiece of the system. If you have multiple agents working in parallel, it acts as a router, directing queries based on the intent. In a single-agent system like this one, you simply use the Kernel Agent with the relevant tools.

class KernelAgent:
    def __init__(self, retrieval_tool: RetrievalTool, summarization_tool: SummarizationTool, question_answering_tool: QuestionAnsweringTool, client, model):
        self.retrieval_tool = retrieval_tool
        self.summarization_tool = summarization_tool
        self.question_answering_tool = question_answering_tool
        self.client = client
        self.model = model

    def classify_query(self, query: str) -> str:
        prompt = f"""
        Classify the following user prompt into one of the three categories:
        - retrieval: The user wants to find a list of papers based on some criteria (e.g., 'Find papers on AI ethics from 2020').
        - summarization: The user wants to summarize a list of papers (e.g., 'Summarize papers with entry_id 123, 456, 789').
        - question_answering: The user wants to ask a question about research topics and get an answer (e.g., 'What is the latest development in AI ethics?').

        User prompt: {query}

        Respond with only the category name (retrieval, summarization, question_answering).
        If unsure, respond with 'unknown'.
        """
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.7,
            max_tokens=10
        )
        classification = response.choices[0].message.content.strip().lower()
        print(f"Query type: {classification}")
        return classification

    def process_query(self, query: str, params: Optional[Dict] = None) -> str:
        query_type = self.classify_query(query)
        if query_type == 'retrieval':
            df_result = self.retrieval_tool.use(query)
            response = "Here are the top papers:\n"
            for i, row in df_result.iterrows():
                # Ensure summary is a string and handle empty cases
                summary = str(row['summary']) if pd.notna(row['summary']) else ""
                response += f"{i+1}. {row['title']} \nSummary: {summary[:200]}...\n\n"
            return response
        elif query_type == 'summarization':
            if not params or 'paper_ids' not in params:
                return "Error: Summarization query requires a 'paper_ids' parameter with a list of entry_ids."
            return self.summarization_tool.use(query, params['paper_ids'])
        elif query_type == 'question_answering':
            return self.question_answering_tool.use(query)
        else:
            return "Error: Unable to classify query as 'retrieval', 'summarization', or 'question_answering'."

By this point, everything is set up for our Research Agent System. All you need to do is initialize the Kernel Agent with the tools, and voilà—our Research Agent system is ready to roll.

retrieval_tool = RetrievalTool(df, app, knowledgebase_query, client, model)
summarization_tool = SummarizationTool(df, client, model)
question_answering_tool = QuestionAnsweringTool(retrieval_tool, client, model)

# Initialize KernelAgent
kernel_agent = KernelAgent(retrieval_tool, summarization_tool, question_answering_tool, client, model)

Now let’s test the system..

# Test query
print(kernel_agent.process_query("Find papers on quantum computing in last 10 years"))

When you run this, the RetrievalTool will be activated. It will fetch the relevant papers based on both relevance and recency, and return the relevant columns. If the returned result includes the summary column (indicating the papers were retrieved from the dataset), it will use those summaries and return them to us.

Query type: retrieval
Here are the top papers:
1. Quantum Computing and Phase Transitions in Combinatorial Search 
Summary: We introduce an algorithm for combinatorial search on quantum computers that
is capable of significantly concentrating amplitude into solutions for some NP
search problems, on average. This is done by...

2. The Road to Quantum Artificial Intelligence 
Summary: This paper overviews the basic principles and recent advances in the emerging
field of Quantum Computation (QC), highlighting its potential application to
Artificial Intelligence (AI). The paper provi...

3. Solving Highly Constrained Search Problems with Quantum Computers 
Summary: A previously developed quantum search algorithm for solving 1-SAT problems in
a single step is generalized to apply to a range of highly constrained k-SAT
problems. We identify a bound on the number o...

4. The model of quantum evolution 
Summary: This paper has been withdrawn by the author due to extremely unscientific
errors....

5. Artificial and Biological Intelligence 
Summary: This article considers evidence from physical and biological sciences to show
machines are deficient compared to biological systems at incorporating
intelligence. Machines fall short on two counts: fi...

Let’s try one more query, this time, let’s do a summarization one..

print(kernel_agent.process_query("Summarize this paper", params={"paper_ids": ["http://arxiv.org/abs/cs/9311101v1"]}))

Query type: summarization
This paper discusses the challenges of learning logic programs that contain the cut predicate (!). Traditional learning methods cannot handle clauses with cut because it has a procedural meaning. The proposed approach is to first generate a candidate base program that covers positive examples, and then make it consistent by inserting cut where needed. Learning programs with cut is difficult due to the need for intensional evaluation, and current induction techniques may need to be limited to purely declarative logic languages.

I mean it’s so cool.

Adios Amigos, until next time.

Creating a Fintech AI agent without any framework

2025-02-23T00:00:00+00:00

AI agents are popping up everywhere, and they’re only going to become more common in the next few years. Think of them as a bunch of little digital assistants—agents—that keep an eye on things, make decisions, and get stuff done. It makes me wonder: how are we going to manage all these agents down the road? How do we trust them to handle decisions for us and let them run on their own?

These days, many frameworks are out there to help you with AI agents quickly. Honestly, it feels like a new “low-code” tool shows up every other day, promising to make building agents a breeze. They’re all over the place—simple drag-and-drop setups or text boxes where you plug in a few prompts and some data, and just like that, your agent’s ready to go.

Every AI agent has three main pieces at its core: the Model, the Tools, and the Reasoning Loop.

The Model: This is the Large Language Model (LLM) that drives the agent. It’s the brain behind it all, trained on piles of text to understand language and tackle tricky instructions. This is what lets the agent think and talk almost like a person.
The Tools: The agent uses specific skills to interact with the world and get things done. Whether it’s grabbing data, hitting up APIs, or working with other programs, the tools are what make the agent useful.
The Reasoning Loop: This is how the agent figures things out. It’s the process it uses to solve problems, pick the best path forward, and wrap up tasks.

In this post, I’m going to mix things up a little. Instead of leaning on a framework, I’ll walk you through building a fintech agent from scratch using Python. This one’s going to help with decisions about loans and insurance.

What’s the deal here?

So how does this fintech agent work? Well, picture yourself running a FinTech company. Customers come at you with stuff like, “Can I borrow $10K for a new kitchen? I’m 26 years old,” or “I crashed my car—will you cover it?” Haha, I’d probably shut that insurance claim down real quick.

The point is, that we need a system that’s fast, sharp, and reliable when it comes to deciding on loan requests and insurance claims. So, what I did is this: instead of one lone agent handling everything, I split it into three—Loan, Insurance, and a Kernel agent.

Loan Agent: Takes a query, figures out what they want (like “home improvement”), and uses an ML model to predict if they’re good for the loan based on age, income, and the amount.
Insurance Agent: Reads a claim, checks it against past cases, and decides if it’s worth paying out—kind of like a semantic search over a claims database.
Kernel Agent: The brains that tie it all together, sending queries to the right place.

Now the question always comes of why raw Python? Well, the most viable reason is flexibility. I’m just putting together a proof of concept here, and companies can easily tweak it to fit their own data and needs—no black-box framework nonsense. The Credit Risk Dataset and synthetic claims are just placeholders; real firms would swap in their own loan histories and claim records.

Big Picture

Ok so what we’ve got is a Kernel Agent directing traffic, a Loan Agent predicting eligibility with machine learning, and an Insurance Agent checking claims against a vector database.

I’m using the Credit Risk Dataset as a demo for loans. In reality, companies could use something like a CIBIL score to see if someone’s eligible. But for this POC, I grabbed that dataset, trained a random forest model, and used it for predictive analysis. For the insurance agent, I built a synthetic pipeline to whip up a dataset of fake insurance queries. The target’s simple: if the claim gets approved based on the conversation, it’s a one. These are just placeholders to show how it works; companies can plug in their own data instead. Saying that, let’s get started.

Step 1 : Setting up the stage - Libraries and Setup

We start by grabbing our tools—libraries that handle data, models, and vectors.

pip install pandas joblib pyarrow sentence-transformers lancedb mistralai

import pandas as pd
import joblib
import pyarrow as pa
from sentence_transformers import SentenceTransformer
import lancedb
from mistralai import Mistral
import os
from abc import ABC, abstractmethod
from typing import Any

# Abstract Tool Class
class Tool(ABC):
    @abstractmethod
    def name(self) -> str:
        pass

    @abstractmethod
    def description(self) -> str:
        pass

    @abstractmethod
    def use(self, *args, **kwargs) -> Any:
        pass

# Initialize Mistral and Embedder
api_key = os.environ.get("MISTRAL_API_KEY", "xxxxxxxxxxxxx")  
if not api_key:
    raise ValueError("Please set the MISTRAL_API_KEY environment variable.")

model = "mistral-large-latest"
client = Mistral(api_key=api_key)
hf_token = "hf_xxxxxxx"
embedder = SentenceTransformer('all-MiniLM-L6-v2', token=hf_token)

# Connect to LanceDB
db = lancedb.connect('./lancedb_data')

For the quick demonstration, I am going to use mistral-large-latest a model from the Mistral for LLM tasks and the intent classification in the Loan queries. The Tool class is our way to keep the agents modular - every tool gets a clear job.

Step 2 : Loan Agent - Predictive analysis with ML

As soon as a query comes in, the Loan Agent uses its ML tool to spin up the predictions. I am using the Credit Risk Dataset here.

Loading the Demo Data

loan_data = pd.read_csv('./credit_risk_dataset.csv')
loan_data.head()

   person_age  person_income  loan_amnt  loan_intent  loan_status
0          22          59000      35000     PERSONAL            1
1          21           9600       1000    EDUCATION            0

Why Credit Risk Dataset? : It’s public, well-structured, and mimics real loan data with features like loan_status (0 = approved, 1 = rejected). For demo purposes, it’s spot-on. Companies would swap this for their own credit records—same columns, different rows, or maybe other relevant benchmarks that can be utilized for setting up the target in terms of the loan status.

Building the ML model

For the sake of the demonstration, I have used a Random Forest Classifier here. You can literally play around with different sets of these models, ensemble them, and get better in terms of accuracy and data processing.

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestClassifier
from sklearn.pipeline import Pipeline
import joblib

features = ['person_age', 'person_income', 'loan_amnt', 'loan_intent']
target = 'loan_status'

loan_data = loan_data.dropna(subset=features + [target])

preprocessor = ColumnTransformer(
    transformers=[('cat', OneHotEncoder(handle_unknown='ignore'), ['loan_intent'])],
    remainder='passthrough'
)

X = loan_data[features]
y = loan_data[target]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

model = Pipeline(steps=[
    ('preprocessor', preprocessor),
    ('classifier', RandomForestClassifier(random_state=42))
])
model.fit(X_train, y_train)

accuracy = model.score(X_test, y_test)
print(f"Model accuracy: {accuracy:.2f}")  # 0.84
joblib.dump(model, 'loan_approval_model.pkl')

Why Random Forest? Well as we know, Random forest comes under the ensemble category, and it’s robust in nature, handles mixed data (numbers and categories) well, and gives us probabilities—not just yes/no. For loans, that’s gold; banks love knowing the odds though.
How It Works: We preprocess loan_intent into numbers (e.g., “MEDICAL” becomes a categorical encoded number), split the data, train, and test. Got an 84% accuracy which is decent for a demo. Real data could boost or tweak that.
Swap It Out: As I said earlier, companies can use their own loan data here—same features, their numbers. Train it fresh, and it’s tailored to them.

Completing the integration for the Loan Agent

Now comes the agent itself:

class PredictiveMLTool:
    def __init__(self):
        self.model = joblib.load('loan_approval_model.pkl')
    
    def use(self, intent, age, income, loan_amount):
        input_data = pd.DataFrame(
            [[age, income, loan_amount, intent]],
            columns=['person_age', 'person_income', 'loan_amnt', 'loan_intent']
        )
        prob_approved = self.model.predict_proba(input_data)[0][0]
        print(f"Probability for the loan approval is: {prob_approved}")
        return "Eligible" if prob_approved > 0.5 else "Not Eligible"

class LoanAgent:
    def __init__(self):
        self.tools = [PredictiveMLTool()]
        self.llm = llm
    
    def process(self, query, params):
        intent = self.extract_intent(query)
        age = params.get('age')
        income = params.get('income')
        loan_amount = params.get('loan_amount')
        return self.tools[0].use(intent, age, income, loan_amount)
    
    def extract_intent(self, query):
        valid_intents = ['DEBTCONSOLIDATION', 'EDUCATION', 'HOMEIMPROVEMENT', 'MEDICAL', 'PERSONAL', 'VENTURE']
        prompt = f"Given the query: '{query}', classify the intent into one of: {', '.join(valid_intents)}. Respond with only the intent in uppercase (e.g., 'HOMEIMPROVEMENT'). If unsure, respond with 'PERSONAL'."
        response = self.llm(prompt)[0]['generated_text'].strip().upper()
        return response if response in valid_intents else 'PERSONAL'

Why ML Here? It’s predictive, which means it takes raw data and spits out a decision. LLM handles intent because it’s quick and decent for text, though companies might use something beefier or their own intent rules. I let the LLM decide what’s the intent here.- How It Fits: Query → Intent (e.g., “school fees” → EDUCATION) → ML prediction. It’s as simple as the tool’s the muscle, and the agent’s the coordinator.

If the companies need a different logic to get the target status in terms if the loan can be provided to an individual or not, as I said earlier, they can easily replace the predictive ml modeling with another set of benchmarks or maybe thresholding some different set of metrics to figure out if that person is made for the loan or not.

To make things more smoother, let’s introduce our Kernel Agent which will be used for redirecting the queries to the Loan and the Insurance agent based on the query type.

class KernelAgent:
    def __init__(self, loan_agent):
        self.loan_agent = loan_agent
    
    def process_query(self, query, params):
        return self.loan_agent.process(query, params)

Well, this is good, but let’s not forget we now need to work on our Insurance agent.

Step 3 : Insurance Agent - Vector Search with LanceDB

Before diving in, I want to walk you through my thought process for creating the Insurance Agent. Picture this: you walk into an insurance company, sit down with someone, and start explaining why you need to file a claim—why it’s legit and why you deserve it. That conversation is key because it’s the starting point for deciding whether the claim gets processed or not.

I tried hunting down a dataset with those kinds of conversations and clear yes or no labels for claims, but I came up empty. So, to keep things moving, I built a simple synthetic dataset pipeline. It generates a fake dataset with text queries on one side and a target on the other—showing if the person filing the claim is eligible or not.

Now, let’s talk about where the VectorDB comes in. Here’s the plan: we’ll embed those insurance claim queries first. Then, when a new claim pops up, we’ll pull the k-most similar queries from the database, look at their average eligibility scores, and use that to decide if the new query deserves a legitimate claim or not. With that in mind, let’s start by generating some synthetic claims for the automotive industry. This is just a quick demo—companies can swap in real conversations instead, and you can tweak the pipeline a bit to handle actual data.

import random
import csv
from mistralai import Mistral

api_key = "xxxxxxxxxxxxxxx"
model = "mistral-large-latest"
client = Mistral(api_key=api_key)

base_prompt = "Generate a detailed query for an auto insurance claim..."
def generate_query(denied=False):
    prompt = base_prompt
    if denied:
        prompt += " Include a reason the claim might be denied (e.g., drunk driving, uninsured vehicle)."
    response = client.chat.complete(model=model, messages=[{"role": "user", "content": prompt}], max_tokens=200)
    return response.choices[0].message.content.strip()

def assign_target(query):
    rejection_keywords = ["drunk", "influence", "racing", "uninsured"]
    return 0 if any(kw in query.lower() for kw in rejection_keywords) else 1

dataset = []
for i in range(100):
    denied = random.random() < 0.4
    query = generate_query(denied)
    target = assign_target(query) if not denied else 0
    dataset.append({"query": query, "target": target})

with open("auto_insurance_claims.csv", "w", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=["query", "target"])
    writer.writeheader()
    writer.writerows(dataset)

To make it a bit more realistic, I have included some of the rejection keywords like drunk, influence racing and the uninsured which is readily used to make the dataset better in terms of the realistic cases where the claims are not processed. Running this will generate a CSV with the relevant data.

Now comes the LanceDB setup..

import pyarrow as pa

df = pd.read_csv('auto_insurance_claims.csv')
embeddings = embedder.encode(df['query'].tolist())
embeddings_list = [embedding.tolist() for embedding in embeddings]

schema = pa.schema([
    pa.field("embedding", pa.list_(pa.float32(), list_size=384)),
    pa.field("target", pa.int32()),
    pa.field("query", pa.string())
])

table = db.create_table("insurance_queries", schema=schema)
df_lance = pd.DataFrame({"embedding": embeddings_list, "target": df['target'], "query": df['query']})
table.add(df_lance)

As I said earlier, companies could load their real claim logs here instead. Now comes the main part, we are going to integrate the semantic tooling capability into our Insurance agent.

# Insurance Agent Tools
class SemanticSearchTool(Tool):
    def __init__(self, table):
        self.table = table
        self.embedder = embedder

    def name(self):
        return "Semantic Search Tool"

    def description(self):
        return "Performs semantic search in LanceDB to assess claim approval based on similar past claims."

    def use(self, query, k=5):
        new_embedding = self.embedder.encode([query])[0].tolist()
        results = self.table.search(new_embedding).limit(k).to_pandas()
        approval_rate = results['target'].mean()
        similar_queries = results['query'].tolist()
        decision = "Approved" if approval_rate > 0.5 else "Not Approved"
        print(f"Approval rate among similar claims: {approval_rate*100:.1f}%")
        return {"decision": decision, "similar_queries": similar_queries}

# Insurance Agent
class InsuranceAgent:
    def __init__(self, table):
        self.tools = [SemanticSearchTool(table)]

    def process(self, query):
        return self.tools[0].use(query)

This is it, let’s integrate our Insurance agent within our Kernel agent along with the Loan Agent.

Step 4 : Completion of Kernel Agent

# Kernel Agent
class KernelAgent:
    def __init__(self, loan_agent, insurance_agent):
        self.loan_agent = loan_agent
        self.insurance_agent = insurance_agent
        self.client = client
        self.model = model

    def classify_query(self, query):
        prompt = f"""
        Given the query: '{query}', classify it as 'loan' or 'insurance'.
        Respond with only 'loan' or 'insurance'.
        If unsure, respond with 'unknown'.
        """
        response = self.client.chat.complete(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.7,
            max_tokens=10
        )
        print(f"Query type : {response.choices[0].message.content.strip().lower()}")
        return response.choices[0].message.content.strip().lower()

    def process_query(self, query, params=None):
        query_type = self.classify_query(query)
        if query_type == 'loan':
            if not params:
                return "Error: Loan query requires parameters (age, income, loan_amount)."
            return self.loan_agent.process(query, params)
        elif query_type == 'insurance':
            return self.insurance_agent.process(query)
        else:
            return "Error: Unable to classify query as 'loan' or 'insurance'."

Now all you need to do is toss in a query—could be anything about loans or insurance. It’ll hit the Kernel Agent first, and based on what the query’s about, it’ll get handed off to either the Loan Agent or the Insurance Agent. From there, they’ll handle the rest of the decisions.

# Test Cases
test_cases = [
    {
        "query": "I had a bad accident yesterday, but i think i was a bit drunk, need some insurance",
        "params": None
    },
]

# Run Tests
print("Testing Integrated System with Mistral:")
print("-" * 50)
for case in test_cases:
    query = case["query"]
    params = case.get("params")
    result = kernel_agent.process_query(query, params)
    print(f"Query: '{query}'")
    if params:
        print(f"Params: {params}")
    if isinstance(result, dict):
        print(f"Decision: {result['decision']}")
        print("Similar Queries:")
        for q in result['similar_queries']:
            print(f"- {q}")
    else:
        print(f"Result: {result}")
    print("-" * 50)

Think about the potential here: fintech companies can plug in their own tools and data to process loans and insurance claims. Picture how much manual work insurance folks put in, checking every little thing to figure out if a claim’s legit or not—or the same deal with loans, deciding if someone qualifies. This could cut out so much of that hassle.

Now if you made this far, dude, you are really into agents. To make your life more easier, here is the colab to run the whole game.

See you in the next one.

Python Package to convert image datasets to lance type

2024-12-04T00:00:00+00:00

A few months ago, I wrote two articles on how the Lance format can supercharge your machine learning workflows. In the first, I showed how Lance’s columnar storage can make handling large image datasets much more efficient for the ML workflows. Then, I followed up with a guide on converting datasets like cinic and mini-imagenet into Lance format using a custom Python script in Google Colab. While that worked well, it was a bit manual.

Well, Some of my friends are lazy as fuxcks but excited enough to run the colab and use the Lance formatted datatype for some of their experiments. Being a good friend, I’m excited to share a much easier solution: the lancify Python package.

I mean it’s literally just running a one command, and boom — your image datasets are in lance format, ready to go. And, just between us, it makes my life a lot easier too.

Installing the Package

Before diving into the conversion process, let’s start by installing the lancify package. You can easily install it via pip:

pip install lancify

Converting Your Image Dataset to Lance

Once you’ve installed the package, converting any image dataset to the Lance format is as simple as running the following Python code. The lancify package abstracts away the complexity of running the colab notebooks manually;

from lancify.converter import convert_dataset

# Define the path to your image dataset
image_dataset_path = 'cards-classification'
resize_dimensions = (256, 256)  # Example resize dimensions
splits = ["train", "test", "valid"]

convert_dataset(
    dataset_path=image_dataset_path,
    batch_size=10,  # You can adjust the batch size as needed
    resize=resize_dimensions,  # Pass resize dimensions if you want to resize images
    splits=splits
)

For this demonstration, I have used this dataset, which provides flexibility in terms of image resizing and dataset splits. The image resizing is optional; by default, the images are processed with their original dimensions. However, if needed, you can specify a target size, such as 256x256, by passing the desired dimensions. If you prefer to keep the original size, simply pass None for the resize parameter. Regarding dataset splits, if the dataset includes predefined divisions like training, testing, and validation sets, you can pass a list specifying the relevant splits.

For datasets that do not have predefined splits, the images are organized by classification labels. In such cases, you only need to provide the dataset path, and a single lance file will be generated, containing all the images with their corresponding labels. This makes sure that the various kinds of image datasets are handled properly whether they include splits or not.

from lancify.converter import convert_dataset

image_dataset_path = 'weather-classification-data'

convert_dataset(
    dataset_path=image_dataset_path,
    batch_size=10,  # You can adjust the batch size as needed
)

The convert_dataset function automatically handles the following:

Reading the image data – It reads image files and their metadata (filename, category, data split).
Converting to Lance – The images are converted into the Lance format with proper schema.
Saving the Lance files – Lance files are saved for each dataset split (train, test, validation) if there are splits in the dataset, if not then a single lance file is saved with the combined data with an adequate schema to segregate the data with the respective labels.

This method is far more concise than manually iterating over directories, creating schemas, and writing to lance files as we did in the previous version using raw colab.

CLI SDK

In addition to using the lancify package programmatically through the imported function, you can also leverage the CLI SDK to convert your image datasets. The SDK offers a CLI for the lancify.

To use the CLI, all you need to do is install the package with pip install lancify and then run the lancify command in your terminal and follow the args.

What’s Happened Behind the Scenes?

To give you a better understanding, here’s a brief overview of what happens when you use lancify:

Image Data: The package reads images from your dataset directory and converts them into a binary format.
Metadata Extraction: Metadata such as the image’s filename, category (label), and data split (train/test/validation) are automatically extracted.
PyArrow RecordBatch: The image data and metadata are packaged into a PyArrow RecordBatch for efficient columnar storage.
Lance Dataset Creation: These RecordBatch objects are then written to Lance datasets, which are optimized for performance and storage.

This process mirrors the manual steps we previously took but in a much more user-friendly manner, significantly reducing the boilerplate code that was necessary before when you have to manually handle the colab

Loading Your Dataset into Pandas

Once your image dataset has been converted into the lance format, you can seamlessly load it into pandas data frames for doing all kind of stuff. Here’s how to do it for the card-classification training lance file..

import lance
import pandas as pd

# Load Lance dataset
ds = lance.dataset('cards-classification/cards-classification_train.lance')
table = ds.to_table()

# Convert Lance table to Pandas dataframe
df = table.to_pandas()
print(df.head())

This is a simple and efficient way to convert your image datasets to the lance format using the lancify package, and it integrates smoothly into your deep learning projects.

Switching to the Lance format makes your data pipelines faster and more efficient, especially when dealing with large image datasets. All you need to do is install the package and run the conversion script on your datasets—no hassle.

Just this small change can really speed up your machine learning workflows—data loading and processing become much quicker, which means your models train faster. If you need a reference, this is a quickie on how to use the lance formatted image datasets for training you deep learning models. And if you’re looking for more ideas, there are plenty of other deep learning recipes built on lance.

Trust me, it’s worth it! 🤗

Improving RAG with RAPTOR

2024-10-30T00:00:00+00:00

Traditional RAG setups often split documents into fixed-size chunks, but this can lead to problems in maintaining the semantic coherence of the text. If a key concept spans multiple chunks, and we only retrieve one chunk, the LLM might lack the full understanding of the idea, leading to incomplete or misleading responses. As a result, crucial ideas and relationships between concepts may be overlooked, leading to incomplete or inaccurate responses.

Additionally, In a flat retrieval structure where all the retrieved chunks are treated equally, this can dilute the importance of critical sections. For example, if one section of the document has key insights but gets buried among less relevant chunks, the model won’t know which parts to prioritize unless we introduce more intelligent weighting or hierarchical structures. I mean it becomes really difficult during the retrieval to weigh which chunk is more important and might be better suitable as a context.

What is RAPTOR?

RAPTOR, which stands for Recursive Abstractive Processing for Tree Organized Retrieval, is a new technique which solves the problems mentioned before. Think of RAPTOR as a librarian who organizes information in a tree-like structure. Instead of simply stacking books in a pile, it clusters similar titles together, creating a hierarchy that narrows as you ascend. Each cluster of books represents a group of related documents, and at the top of each cluster, there’s a summary that encapsulates the key points from all the books below it. This process continues all the way to the top of the tree, providing a comprehensive view of the information—it’s like having both a magnifying glass and a telescope!

To visualize this further, think of the leaves of the tree as document chunks. These chunks are grouped into clusters to generate meaningful summaries, which then become the new leaves of the tree. This recursive process repeats until reaching the top.

Key terms to look out for

Before we dive in, let’s quickly review some key terms that will be useful as we explore RAPTOR tech. I just want to put it up here to make sure you are comfortable with the nitty tech details as we go along.

GMM Clustering: Gaussian Mixture Models (GMM) group data into clusters based on statistical probabilities. So instead of rigidly classifying each instance into one category like K-means, GMM generates K-Gaussian distributions that consider the entire training space. This means that each point can belong to one or more distributions.
Dimensionality Reduction: This process simplifies the data by reducing the number of variables while retaining essential features. It’s particularly important for understanding high-dimensional datasets like embeddings.
UMAP: Uniform Manifold Approximation and Projection (UMAP) is a powerful dimensionality reduction algorithm we’ll use to shrink the size of our data point embeddings. This reduction makes it easier for clustering algorithms like GMM to cluster high-dimensional embeddings.
BIC and Elbow Method: Both techniques help identify the optimal number of clusters in a dataset. The Bayesian Information Criterion (BIC) evaluates models based on their fit to the data while penalizing complexity. The Elbow Method plots explained variance against the number of clusters, helping to pinpoint where adding more clusters offers diminishing returns. For our purposes, we’ll leverage both methods to determine the best number of clusters.

How it actually works?

Now that you’re familiar with the key terms (and if not, no worries—you’ll catch on as we go!), let’s dive into how everything actually works under the hood of RAPTOR.

Starting Documents as Leaves: The leaves of the tree represent a set of initial documents, which are our text chunks.
Embedding and Clustering: The leaves are embedded and clustered. The authors utilize the UMAP dimensionality reduction algorithm to minimize the embedding size of these chunks. For clustering, Gaussian Mixture Models (GMM) are employed to ensure effective grouping, addressing the challenges posed by high-dimensional vector embeddings.
Summarizing Clusters: Once clustered, these groups of similar chunks are summarized into higher-level abstractions nodes. Each cluster acts like a basket for similar documents, and the individual summaries encapsulate the essence of all nodes within that cluster. This process builds from the bottom up, where nodes are clustered together to create summaries that are then passed up the hierarchy.
Recursive Process: This entire procedure is recursive, resulting in a tree structure that transitions from raw documents (the leaves) to more abstract summaries, with each summary derived from the clusters of various nodes.

Building the RAPTOR

Now that we’ve unpacked how it all works (and you’re still with me hopefully, right?), let’s shift gears and talk about how we actually build the RAPTOR tree.

Setup and Imports

pip install lancedb scikit-learn openai torch sentence_transformers tiktoken umap-learn PyPDF2

import os
import uuid
import tiktoken
import re
import numpy as np
import pandas as pd
import transformers
import torch
import umap.umap_ as umap
import matplotlib.pyplot as plt
from openai import OpenAI
from typing import List, Tuple, Optional, Dict
from sklearn.mixture import GaussianMixture
from sentence_transformers import SentenceTransformer

openai_api_key = "sk-XXXXXXXXXXXXXXX"
client = OpenAI(api_key=openai_api_key)

Creating the Chunks

Setting up RAPTOR is pretty straightforward and builds on what we’ve already covered. The first step is to break down our textual documents into smaller chunks. Once we have those, we can convert them into dense vector embeddings. To simplify things, our use case example will use just one financial document as a source. Download it and add to a directory. For now, I used a data directory where all the documents can be added. You can of course load any source documents (and alter the query below) in your own experiments!

Note: Make sure you use an unencrypted version of this file to avoid errors when PyPDF2 tries to read the file.

import os
import PyPDF2

# Function to extract text from a PDF file
def extract_pdf_text(file_path):
    with open(file_path, 'rb') as file:
        reader = PyPDF2.PdfReader(file)
        text = ""
        for page_num in range(len(reader.pages)):
            page = reader.pages[page_num]
            text += page.extract_text()
    return text

# Function to split text into chunks with overlap
def split_text(text, chunk_size=1000, chunk_overlap=50):
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunks.append(text[start:end])
        start += chunk_size - chunk_overlap
    return chunks

# Function to process all PDFs in a directory
def process_directory(directory_path, chunk_size=1000, chunk_overlap=50):
    all_chunks = []
    # Iterate over all PDF files in the directory
    for filename in os.listdir(directory_path):
        if filename.endswith(".pdf"):
            file_path = os.path.join(directory_path, filename)
            print(f"Processing file: {file_path}")
            
            # Step 1: Extract text from the PDF
            pdf_text = extract_pdf_text(file_path)
            
            # Step 2: Split the extracted text into chunks
            chunks = split_text(pdf_text, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
            
            # Append chunks from this file to the overall list
            all_chunks.extend(chunks)
    
    return all_chunks

directory_path = os.path.join(os.getcwd(), "data")  # Specify your directory path
chunk_size = 1000
chunk_overlap = 50

# Process all PDF files in the directory and get the chunks
chunks = process_directory(directory_path, chunk_size=chunk_size, chunk_overlap=chunk_overlap)

# Optional: Print the number of chunks and preview some of the chunks
print(f"Total number of chunks: {len(chunks)}")
for i, chunk in enumerate(chunks[:5]):  # Show first 5 chunks as a preview
    print(f"Chunk {i+1}:\n{chunk}\n")

Now that we have our chunks, it’s time to dive into the recursive processing to create summarized nodes. For the embedding part, I’ll be using the all-MiniLM-L6-v2 model from Sentence Transformers, but feel free to choose any embedding model that suits your needs—it’s entirely up to you!

embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

Clustering and Dimensionality Reduction

Now we have our embedded chunks, it’s time to step on for the next set of tasks. When diving into RAPTOR, one of the biggest hurdles we encounter is the high dimensionality of vector embeddings. Traditional clustering methods like Gaussian Mixture Models (GMM) often struggle with this complexity, making it tough to effectively cluster high-dimensional data chunks. To tackle this challenge, we turn to Uniform Manifold Approximation and Projection (UMAP). UMAP excels at simplifying data while preserving the essential structures that matter most.

A key factor in UMAP’s effectiveness is the n_neighbors parameter. This setting dictates how much of the data’s neighborhood UMAP considers during dimensionality reduction. In simpler terms, it helps you choose between zooming in on details or taking a broader view:

Higher n_neighbors: A higher value encourages UMAP to “look at many neighbors,” which helps maintain the global structure of the data. This results in larger, more general clusters.
Lower n_neighbors: Conversely, lowering n_neighbors prompts UMAP to “focus on close relationships,” enabling it to preserve the local structure and form smaller, more detailed clusters.

Think of it this way: Imagine you’re at a party. If you take a step back and look around (high n_neighbors), you can see the whole room—where the groups are forming, who’s mingling, and the general vibe. But if you lean in closer to a specific group (low n_neighbors), you can hear their conversation and pick up on the nuances, like inside jokes or shared interests. Both perspectives are valuable; it just depends on what you want to understand.

In RAPTOR, we leverage this flexibility in n_neighbors to create a hierarchical clustering structure. We first run UMAP with a higher n_neighbors to identify the global clusters—the broad categories. Then, we narrow the focus by lowering the value to uncover local clusters within those broader groups. This two-step approach ensures we capture both large-scale patterns and intricate details.

Well, TL;DR

Dimensionality Reduction helps manage high-dimensional data, and UMAP is our primary tool for that.
The n_neighbors parameter controls the balance between seeing the “big picture” and honing in on local details.
The clustering process begins with global clusters (using high n_neighbors), followed by a focus on local clusters with a lower setting of n_neighbors.

def dimensionality_reduction(
    embeddings: np.ndarray,
    target_dim: int,
    clustering_type: str,
    metric: str = "cosine",
) -> np.ndarray:
    if clustering_type == "local":
        n_neighbors = max(2, min(10, len(embeddings) - 1))
        min_dist = 0.01
    elif clustering_type == "global":
        n_neighbors = max(2, min(int((len(embeddings) - 1) ** 0.5), len(embeddings) // 10, len(embeddings) - 1))
        min_dist = 0.1
    else:
        raise ValueError("clustering_type must be either 'local' or 'global'")

    umap_model = umap.UMAP(
        n_neighbors=n_neighbors,
        min_dist=min_dist,
        n_components=target_dim,
        metric=metric,
    )
    return umap_model.fit_transform(embeddings)

I plan to leverage both the Elbow Method and the Bayesian Information Criterion (BIC) to pinpoint the optimal number of clusters for our analysis.

def compute_inertia(embeddings: np.ndarray, labels: np.ndarray, centroids: np.ndarray) -> float:
    return np.sum(np.min(np.sum((embeddings[:, np.newaxis] - centroids) ** 2, axis=2), axis=1))

def optimal_cluster_number(
    embeddings: np.ndarray,
    max_clusters: int = 50,
    random_state: int = SEED
) -> int:
    max_clusters = min(max_clusters, len(embeddings))
    number_of_clusters = np.arange(1, max_clusters + 1)
    inertias = []
    bic_scores = []
    
    for n in number_of_clusters:
        gmm = GaussianMixture(n_components=n, random_state=random_state)
        labels = gmm.fit_predict(embeddings)
        centroids = gmm.means_
        inertia = compute_inertia(embeddings, labels, centroids)
        inertias.append(inertia)
        bic_scores.append(gmm.bic(embeddings))
    
    inertia_changes = np.diff(inertias)
    elbow_optimal = number_of_clusters[np.argmin(inertia_changes) + 1]
    bic_optimal = number_of_clusters[np.argmin(bic_scores)]
    
    return max(elbow_optimal, bic_optimal)

def gmm_clustering(
    embeddings: np.ndarray, 
    threshold: float, 
    random_state: int = SEED
) -> Tuple[List[np.ndarray], int]:
    n_clusters = optimal_cluster_number(embeddings, random_state=random_state)
    gm = GaussianMixture(n_components=n_clusters, random_state=random_state, n_init=2)
    gm.fit(embeddings)
    probs = gm.predict_proba(embeddings)
    labels = [np.where(prob > threshold)[0] for prob in probs] 
    return labels, n_clusters  

Tree Construction

Now that we’ve wrapped up the clustering part, let’s talk about how we build our hierarchical tree. After several rounds of clustering and summarization (while keeping track of how deep we go), here’s what we have:

Leaf Nodes: These are our original text chunks, forming the base of the tree.
Summary Nodes: As we go up the tree, each node acts like a quick summary of its child nodes, capturing the main idea of the cluster.
Hierarchical Embeddings: The summary nodes can also become the new nodes at their level. Each of these nodes gets its own vector embedding, representing the summarized meaning. So, we’re essentially adding more nodes while enriching them with summaries.

The process flows nicely: we embed the chunks, reduce their dimensions using UMAP, cluster them with Gaussian Mixture Models, start with a broad overview, and then zoom in for more detailed clusters before summarizing.

def clustering_algorithm(
    embeddings: np.ndarray,
    target_dim: int,
    threshold: float,
    random_state: int = SEED
) -> Tuple[List[np.ndarray], int]:
    if len(embeddings) <= target_dim + 1:
        return [np.array([0]) for _ in range(len(embeddings))], 1
    
    # Global clustering
    reduced_global_embeddings = dimensionality_reduction(embeddings, target_dim, "global")
    global_clusters, n_global_clusters = gmm_clustering(reduced_global_embeddings, threshold, random_state=random_state)

    all_local_clusters = [np.array([]) for _ in range(len(embeddings))]
    total_clusters = 0

    # Local clustering within each global cluster
    for i in range(n_global_clusters):
        global_cluster_mask = np.array([i in gc for gc in global_clusters])
        global_cluster_embeddings = embeddings[global_cluster_mask]

        if len(global_cluster_embeddings) <= target_dim + 1:
            # Assign all points in this global cluster to a single local cluster
            for idx in np.where(global_cluster_mask)[0]:
                all_local_clusters[idx] = np.append(all_local_clusters[idx], total_clusters)
            total_clusters += 1
            continue

        try:
            reduced_local_embeddings = dimensionality_reduction(global_cluster_embeddings, target_dim, "local")
            local_clusters, n_local_clusters = gmm_clustering(reduced_local_embeddings, threshold, random_state=random_state)

            # Assign local cluster IDs
            for j in range(n_local_clusters):
                local_cluster_mask = np.array([j in lc for lc in local_clusters])
                global_indices = np.where(global_cluster_mask)[0]
                local_indices = global_indices[local_cluster_mask]
                for idx in local_indices:
                    all_local_clusters[idx] = np.append(all_local_clusters[idx], j + total_clusters)

            total_clusters += n_local_clusters
        except Exception as e:
            print(f"Error in local clustering for global cluster {i}: {str(e)}")
            # Assign all points in this global cluster to a single local cluster
            for idx in np.where(global_cluster_mask)[0]:
                all_local_clusters[idx] = np.append(all_local_clusters[idx], total_clusters)
            total_clusters += 1

    return all_local_clusters, total_clusters
    
 def generate_summary(context: str) -> str:
    prompt = f"""
    Provide the Summary for the given context. Here are some additional instructions for you:

    Instructions:
    1. Don't make things up, Just use the contexts and generate the relevant summary.
    2. Don't mix the numbers, Just use the numbers in the context.
    3. Don't try to use fancy words, stick to the basics of the language that is being used in the context.

    Context: {context}
    """
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a helpful assistant that summarizes text."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=200,
        n=1,
        stop=None,
        temperature=0.7
    )
    summary = response.choices[0].message.content.strip()
    return summary

def embed_clusters(
    texts: List[str],
    target_dim: int = 10,
    threshold: float = 0.1
) -> pd.DataFrame:
    textual_embeddings = np.array(embedding_model.encode(texts))
    clusters, number_of_clusters = clustering_algorithm(textual_embeddings, target_dim, threshold)
    print(f"Number of clusters: {number_of_clusters}")
    return pd.DataFrame({
        "texts": texts,
        "embedding": list(textual_embeddings),
        "clusters": clusters
    })

def embed_cluster_summaries(
    texts: List[str],
    level: int,
    target_dim: int = 10,
    threshold: float = 0.1
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    df_clusters = embed_clusters(texts, target_dim, threshold)
    main_list = []
    
    for _, row in df_clusters.iterrows():
        for cluster in row["clusters"]:
            main_list.append({
                "text": row["texts"],
                "embedding": row["embedding"],
                "clusters": cluster
            })
    
    main_df = pd.DataFrame(main_list)
    unique_clusters = main_df["clusters"].unique()
    if len(unique_clusters) == 0:
        return df_clusters, pd.DataFrame(columns=["summaries", "level", "clusters"])

    print(f"--Generated {len(unique_clusters)} clusters--")

    summaries = []
    for cluster in unique_clusters:
        text_in_df = main_df[main_df["clusters"] == cluster]
        unique_texts = text_in_df["text"].tolist()
        text = "------\n------".join(unique_texts)
        summary = generate_summary(text)
        summaries.append(summary)

    df_summaries = pd.DataFrame({
        "summaries": summaries,
        "level": [level] * len(summaries),
        "clusters": unique_clusters
    })

    return df_clusters, df_summaries


def recursive_embedding_with_cluster_summarization(
    texts: List[str],
    number_of_levels: int = 3,
    level: int = 1,
    target_dim: int = 10,
    threshold: float = 0.1
) -> Dict[int, Tuple[pd.DataFrame, pd.DataFrame]]:
    if level > number_of_levels:
        return {}
    
    results = {}
    df_clusters, df_summaries = embed_cluster_summaries(texts, level, target_dim, threshold)
    results[level] = (df_clusters, df_summaries)
    
    if df_summaries.empty or len(df_summaries['clusters'].unique()) == 1:
        print(f"No more unique clusters found at level {level}. Stopping recursion.")
        return results
    
    if level < number_of_levels:
        next_level_texts = df_summaries['summaries'].tolist()
        next_level_results = recursive_embedding_with_cluster_summarization(
            next_level_texts, 
            number_of_levels, 
            level + 1,
            target_dim,
            threshold
        )
        results.update(next_level_results)
    
    return results

Okay, the code might seem a bit daunting at first glance, but don’t worry! Just give it a couple of looks, and it will start to make sense. Essentially, we’re just following the flow I mentioned earlier.

def process_text_hierarchy(
    texts: List[str], 
    number_of_levels: int = 3,
    target_dim: int = 10,
    threshold: float = 0.1
) -> Dict[str, pd.DataFrame]:
    hierarchy_results = recursive_embedding_with_cluster_summarization(
        texts, number_of_levels, target_dim=target_dim, threshold=threshold
    )
    
    processed_results = {}
    for level, (df_clusters, df_summaries) in hierarchy_results.items():
        if df_clusters.empty or df_summaries.empty:
            print(f"No data for level {level}. Skipping.")
            continue
        processed_results[f"level_{level}_clusters"] = df_clusters
        processed_results[f"level_{level}_summaries"] = df_summaries
    
    return processed_results


results = process_text_hierarchy(chunks, number_of_levels=3)

Inference

Now that we have our tree structure with leaf nodes at the bottom and summarized nodes in between, it’s time to query the RAG. There are two main methods for navigating the RAPTOR tree: Tree Traversal and Collapsed Tree Retrieval.

Tree Traversal Retrieval: This method systematically explores the tree, starting from the root node. It first selects the top-k most relevant root nodes based on their cosine similarity to the query embedding. Then, for each selected root node, its children are considered in the next layer, where the top-k nodes are again selected based on their cosine similarity to the query vector. This process repeats until we reach the leaf nodes. Finally, the text from all the selected nodes is concatenated to form the retrieved context.
Collapsed Tree Retrieval: This approach simplifies things by viewing the tree as a single layer. Here, it directly compares the query embedding to the vector embeddings of all the leaf nodes (the original text chunks) and summary nodes. This method works best for factual, keyword-based queries where you need specific details.

In the collapsed tree retrieval, we flatten the tree into one layer, retrieving nodes based on cosine similarity until we reach a specified number of top k documents. In our code, we’ll gather the textual chunks from earlier, along with the summarized nodes at each level for all the clusters, to create one big list of texts that includes both the root documents and the summarized nodes.

To be honest, if you look closely, we’ve been essentially adding more data points (chunks) to our RAG setup all along. Using RAPTOR, we now have both the original chunks and the summarized chunks for each cluster. Now, we’ll simply embed all these new data points and store them in a vector database along with their embeddings and use them for RAG.

raptor_texts = []
for level, row in results.items():
    if level.endswith("clusters"):
        raptor_texts.extend(row["texts"])
    else:
        raptor_texts.extend(row["summaries"])
        
raptor_embeddings = embedding_model.encode(raptor_texts)
len(raptor_embeddings)

Setting up Vector Database and RAG

Now it’s smooth sailing! We’ll just set up a LanceDB vector database to store our embeddings and query our RAG setup.

raptor_embeddings = embedding_model.encode(raptor_texts)

raptor_dict = {"texts": [], "embeddings": []}
for texts, embeddings in zip(raptor_texts, raptor_embeddings):
    raptor_dict["texts"].append(texts)
    raptor_dict["embeddings"].append(embeddings.tolist())

import lancedb
import pyarrow as pa
from lancedb.pydantic import Vector, LanceModel

uri = "lancedb_database"
db = lancedb.connect(uri)

class RAG(LanceModel):
    texts : str
    embeddings : Vector(384)

table_name = "rag_with_raptor"
raptor_table = db.create_table(table_name, schema = RAG, mode="overwrite")
raptor_table.add(rag_raptor_df)
raptor_table.create_fts_index("texts", replace=True)

Time to generate the results..

def generate_results(
    query : str,
    context_text : str
) -> str:

    prompt = f"""
    Based on the context provided, use it to answer the query. 

    query : {query}

    Instructions:
    1. Don't make things up, Just use the contexts and generate the relevant answer.
    2. Don't mix the numbers, Just use the numbers in the context.
    3. Don't try to use fancy words, stick to the basics of the language that is being used in the context.
    
    {context_text}
    """
    response = client.chat.completions.create(
        model="gpt-4", 
        messages=[
            {"role": "system", "content": "You are a helpful assistant that answers query and give the answers."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=200,
        n=1,
        stop=None,
        temperature=0.7
    )
    answer = response.choices[0].message.content.strip()
    return answer

query = "NTT DATA's net income attributable to shareholders increased from ¥69,227 million in Q3 FY2020 to ¥110,191 million in Q3 FY2021. How does this growth align with their acquisition strategy, particularly considering their stated reasons for acquiring Nexient, LLC and the provisional goodwill recorded in this transaction?"

In our query, there are several key points that must be addressed when crafting the answers. First, we need to note the increase in net income from ¥69,227 million in Q3 FY2020 to ¥110,191 million in Q3 FY2021. Second, we should examine how this growth aligns with NTT DATA’s acquisition strategy, particularly their reasons for acquiring Nexient, LLC, and the provisional goodwill recorded in the transaction. With this context in mind, I created a VANILLA RAG to compare its results with those of RAPTOR RAG.

normal_embeddings = embedding_model.encode(chunks) # default chunks from our data

normal_dict = {"texts": [], "embeddings": []}
for texts, embeddings in zip(chunks, normal_embeddings):
    normal_dict["texts"].append(texts)
    normal_dict["embeddings"].append(embeddings.tolist())
    
rag_normal_df = pd.DataFrame(normal_dict)

table_name = "rag_without_raptor"
normal_table = db.create_table(table_name, schema = RAG, mode="overwrite")
normal_table.add(rag_normal_df)
normal_table.create_fts_index("texts", replace=True)

With RAPTOR, we now have an increased number of chunks due to the addition of cluster-level summary nodes alongside the default chunks we had earlier.

D-Day

raptor_contexts = raptor_table.search(query).limit(5).select(["texts"]).to_list()
raptor_context_text = "------\n\n".join([context["texts"] for context in raptor_contexts])
raptor_context_text = "------\n\n" + raptor_context_text

normal_contexts = normal_table.search(query).limit(5).select(["texts"]).to_list()
normal_context_text = "------\n\n".join([context["texts"] for context in normal_contexts])
normal_context_text = "------\n\n" + normal_context_text

raptor_answer = generate_results(query, raptor_context_text)
normal_answer = generate_results(query, normal_context_text)

When we are comparing RAPTOR RAG with Vanilla RAG, it’s clear that RAPTOR performs better. Not only does RAPTOR retrieve details about the financial growth, but it also effectively connects this growth to the broader acquisition strategy, pulling relevant context from multiple sources. It excels in situations like this, where the query requires insights from various pages, making it more adept at handling complex, layered information retrieval.

And that’s a wrap for this article! If you want to dig into the intricacies of how everything works, I’d suggest checking out the official RAPTOR GitHub repository for more info and resources. For an even deeper dive, the official paper is a great read and highly recommended! Here is the Google colab for your reference.

Multi Document Agentic RAG

2024-08-16T00:00:00+00:00

Agentic RAG (Retrieval-Augmented Generation) represents a significant leap in how we handle information. Traditional RAG systems are built to retrieve information and relevant context which are then passively send that to a language model (LLM) to generate responses. However, Agentic RAG takes this further by adding more independence. Now, the system can not only gather data but also make decisions and take actions on its own.

Think of it as a shift from simple tools to smarter, more capable systems. Agentic RAG transforms what was once a passive process into an active one, where AI can work towards specific goals without needing constant guidance.

How Does an Agentic RAG Work?

To understand Agentic RAG, let’s first break down what an “agent” is. Simply put, an agent is a smart system capable of making decisions on its own. When given a question or task, it figures out the best way to handle it by breaking the task into smaller steps and using the right tools to complete it.

Now, when we talk about Agentic RAG, we’re taking this concept further. Instead of just retrieving information like regular RAG, Agentic RAG uses intelligent strategies to ensure the system provides the best possible response. It doesn’t stop at giving a basic, context-aware reply. Instead, these agents think through the query, select the right approach, and deliver a more thoughtful and refined answer. It’s like having a team of smart agents working together to solve a problem, one step at a time.

Here’s how the process works:

Understanding the Query: The agent starts by analyzing the question to grasp its specifics. It looks at the context, purpose, and intent to determine what information is needed.
Using Memory: The agent checks its memory for relevant information from past tasks that might assist with the current query.
Choosing Tools: Agentic RAG agents are designed to be intuitive. After analyzing the query, they evaluate available tools and resources to select the best method for retrieving the precise information. It’s like having a savvy assistant who knows exactly where to look, even for the most challenging questions.

How to Use it?

Let’s dive into developing an automotive-themed RAG agent based on our current understanding.

Consider a scenario where you own a vehicle and require assistance with tasks ranging from diagnosing issues to planning routine maintenance. Now imagine a specialized agent designed specifically for this purpose. This agent should be capable of interpreting your car’s symptoms, analyzing the issue, and delivering a detailed diagnosis, including potential causes.

Furthermore, the bot should assist with identifying specific parts, estimating repair costs, or creating a personalized maintenance schedule based on your vehicle’s mileage and model. This agent must effectively manage a range of tasks by utilizing relevant context from various datasources, employing different tools, and reasoning with the available information to provide accurate and meaningful responses. Given the complexity of these tasks, which involves multiple retrieval and reasoning steps, I am going to use LanceDB to ensure fast retrieval by storing relevant embedded data chunks into it.

To meet our data needs, I will use six JSON files, each containing specific types of information for querying. You can get these JSON files here : Data

Here’s a brief overview of each file:

car_maintenance.json: Contains details about the car’s maintenance schedule, including relevant tasks based on mileage and estimated time for completion.
car_problems.json: Provides information on various car problems, including a brief description of each issue, required parts for resolution, estimated repair time, and other relevant metadata.
car_parts.json: Lists car parts used for maintenance and diagnosis, detailing their brands, categories, whether they are solid or liquid, and other relevant attributes.
car_diagnosis.json: Outlines the diagnosis process, potential causes based on symptoms, recommended actions, and related problems. The issues and parts mentioned should align with those in car_problems.json and car_parts.json to ensure the agent has relevant context for problem-solving.
car_cost_estimates.json: Provides cost estimates for addressing car problems, based on the issues listed in car_problems.json.
car_models.json: Contains information on common problems associated with specific car models, such as a 2017 Honda Accord with 190,000 kilometers, detailing typical issues that users might encounter in this range.

Please review the JSON files to see their structure and feel free to make any changes.

How to: Tech Stack

To build our ReAct-like agent, we’ll be using a few key tools to make things run smoothly:

LlamaIndex: Think of LlamaIndex as the backbone of our agent. This framework will be central to our implementation. It facilitates the abstraction of agentic logic. If the specifics seem unclear now, they will become more evident as we proceed with the implementation.
Memory Management: When we query the agent, it handles each question on its own without remembering past interactions or in isolation without maintaining state. To address this, we will use memory to retain conversational history. The agent stores chat history in a conversational memory buffer, which by default is a flat list managed by LlamaIndex. This ensures that the agent can refer to past as well as current conversation when deciding on the next set of actions.
Vector Databases: For the retrieval process, we will use VectorDBs. Queries will be embedded and matched semantically against the relevant VectorDB through our retrievers. We will employ LanceDB due to its exceptional retrieval speed and on-disk storage capabilities, which allows for local management of our database. Additionally, being open-source and free, it fits within our budget constraints.
LLM Integration: On the language model side, we’ll go with OpenAI’s GPT-4 for generating responses. For embeddings, we’re using Hugging face which provides seamless integration of local embedding models.

Ok I think this is enough, let’s dive in for our code part..

Environment Setup

import os
import tqdm
import json
import time
from typing import List, Optional, Dict, Any
from tqdm import tqdm
import logging
from llama_index.core import (
    SimpleDirectoryReader, 
    VectorStoreIndex, 
    StorageContext,
    Settings,
    Document,
)
from datetime import datetime, timedelta
from llama_index.core.node_parser import SentenceSplitter
from llama_index.vector_stores.lancedb import LanceDBVectorStore
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.core.tools import FunctionTool, ToolOutput
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.agent import AgentRunner

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Load environment variables
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
if not OPENAI_API_KEY:
    raise ValueError("Please set the OPENAI_API_KEY environment variable.")

# LLM setup
llm = OpenAI(model="gpt-4", api_key=OPENAI_API_KEY)

# Embedding model setup
embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Update the Settings with the new embedding model
Settings.embed_model = embed_model
Settings.chunk_size = 512

Make sure your .env file includes the OPENAI_API_KEY. You can adjust and use different LLMs and embedding models as needed. The key is to have an LLM for reasoning and an embedding model to handle data embedding. Feel free to experiment with various models to find the best fit for your needs.

Step 1 : Creating our DBs

Let’s setup our database which will be used to store our data.

# Vector store setup
problems_vector_store = LanceDBVectorStore(
    uri='./lancedb',
    table_name='problems_table',
    mode="overwrite",
)

parts_vector_store = LanceDBVectorStore(
    uri='./lancedb',
    table_name='parts_table',
    mode="overwrite",
)

diagnostics_vector_store = LanceDBVectorStore(
    uri='./lancedb',
    table_name='diagnostics_table',
    mode="overwrite",
)

cost_estimates_vector_store = LanceDBVectorStore(
    uri='./lancedb',
    table_name='cost_estimates_table',
    mode="overwrite",
)

maintenance_schedules_vector_store = LanceDBVectorStore(
    uri='./lancedb',
    table_name='maintenance_schedules_table',
    mode="overwrite",
)

cars_vector_store = LanceDBVectorStore(
    uri='./lancedb',
    table_name='car_maintenance_table',
    mode="overwrite",
)

Since we’re dealing with various types of information, we’ll need multiple tables to organize and retrieve our data effectively. The uri specifies that our data is stored in a database called lancedb, which contains different tables for each type of data. Let’s go ahead and load the data into the appropriate tables.

def load_and_index_documents(directory: str, vector_store: LanceDBVectorStore) -> VectorStoreIndex:
    """Load documents from a directory and index them."""
    documents = SimpleDirectoryReader(input_dir=directory).load_data()
    parser = SentenceSplitter(chunk_size=1024, chunk_overlap=200)
    nodes = parser.get_nodes_from_documents(documents)
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    return VectorStoreIndex(nodes, storage_context=storage_context)

def create_retriever(index: VectorStoreIndex) -> VectorIndexRetriever:
    """Create a retriever from the index."""
    return index.as_retriever(similarity_top_k=5)

# Load and index documents
problems_index = load_and_index_documents("../rag-agentic-system/problems", problems_vector_store)
parts_index = load_and_index_documents("../rag-agentic-system/parts", parts_vector_store)
cars_index = load_and_index_documents("../rag-agentic-system/cars_models", cars_vector_store)
diagnostics_index = load_and_index_documents("../rag-agentic-system/diagnostics", diagnostics_vector_store)
cost_estimates_index = load_and_index_documents("../rag-agentic-system/cost_estimates", cost_estimates_vector_store)
maintenance_schedules_index = load_and_index_documents("../rag-agentic-system/maintenance_schedules", maintenance_schedules_vector_store)


problems_retriever = create_retriever(problems_index)
parts_retriever = create_retriever(parts_index)
cars_retriever = create_retriever(cars_index)
diagnostics_retriever = create_retriever(diagnostics_index)
cost_estimates_retriever = create_retriever(cost_estimates_index)
maintenance_retriever = create_retriever(maintenance_schedules_index)

Each Vector DB will provide a retriever instance, a Python object that returns a list of documents matching a given query. For example, our problems_retriever will fetch documents related to car problems based on the query, while cars_retriever will help identify common issues faced by customers with their vehicles.

Keep in mind that if the bot misses some information or seems to hallucinate, it might be due to missing data in our JSON files. If you spot any inaccuracies or gaps, add the relevant data to the JSON files and re-index them to ensure everything stays up to date.

Now, let’s test our retrievers to ensure they’re working correctly.

Step 2 : Testing our retrievers

Let’s test our cost_estimates_retriever to ensure it’s working properly. We’ll use a query related to a brake problem and check if the retriever returns the correct documents about cost estimates. Additionally, I’ll verify if our query engine is accurately interpreting the query and if, by providing the relevant documents, we are receiving the correct response.

query = "My brake pad isn't working or I don't know, but the brakes are poor, and by the way, what's the cost for the solution?"
query_engine = cost_estimates_index.as_query_engine()
response = query_engine.query(query)
results = cost_estimates_retriever.retrieve(query)

print(f"Response: {response}")

for result in results:
    print(f"Result - Node ID: {result.node_id}")
    print(f"Relevant Text: {result.text[:100]}...")  
    print(f"Score: {result.score:.3f}")

This is the output I’m receiving in the cell

Response: The cost for the solution to address the issue with your brake pad would be around $100 to $300 for brake pad replacement.
Result - Node ID: 51b91e6e-1243-405b-a742-89e09d78616f
Relevant Text: [
    {
        "repair": "Brake pad replacement",
        "average_cost": 150,
        "cost_range": {
            "min": 100,
            "max": 300...
Score: 0.584

Everything looks good—our retriever is working as expected. We’re ready to move on to the next step.

Step 3 : Creating our agentic tools

LlamaIndex Agents are designed to process natural language input to execute actions rather than simply generating responses. The effectiveness of these agents relies on how well we abstract and utilize tools. So, what exactly does “tool” mean in this context? To clarify, imagine tools as weapons given to a warrior in battle. Just as a warrior might choose different weapons based on the opponent’s tactics, tools for our agent are like specialized API interfaces that help the agent interact with data sources or reason through queries to deliver the best possible responses.

In LlamaIndex, there are various types of tools. One important type is the FunctionTool, which transforms any user-defined function into a tool, capable of inferring the function’s schema and usage. These tools are essential for our agents, allowing them to reason about queries and perform actions effectively.

For each tool, it’s essential to provide a clear description of its purpose and functionality, as this helps the agent use the tool effectively. To start, we will create tools to leverage the retriever objects defined earlier.

max_context_information = 200

def retrieve_problems(query: str) -> str:
    """Searches the problem catalog to find relevant automotive problems for the query."""
    docs = problems_retriever.retrieve(query)
    information = str([doc.text[:max_context_information]for doc in docs])
    return information
    
def retrieve_parts(query: str) -> str:
    """Searches the parts catalog to find relevant parts for the query."""
    docs = parts_retriever.retrieve(query)
    information = str([doc.text[:max_context_information]for doc in docs])
    return information

def retrieve_car_details(make: str, model: str, year: int) -> str:
    """Retrieves the make, model, and year of the car."""
    docs = car_details_retriever.retrieve(make, model, year)
    information = str([doc.text[:max_context_information]for doc in docs])

def diagnose_car_problem(symptoms: str) -> str:
    """Uses the diagnostics database to find potential causes for given symptoms."""
    docs = diagnostics_retriever.retrieve(symptoms)
    information = str([doc.text[:max_context_information]for doc in docs])
    return information

def estimate_repair_cost(problem: str) -> str:
    """Provides a cost estimate for a given car problem or repair."""
    docs = cost_estimates_retriever.retrieve(problem)
    information = str([doc.text[:max_context_information]for doc in docs])
    return information

def get_maintenance_schedule(mileage: int) -> str:
    """Retrieves the recommended maintenance schedule based on mileage."""
    docs = maintenance_retriever.retrieve(str(mileage))
    information = str([doc.text[:max_context_information]for doc in docs])
    return information

retrieve_problems_tool = FunctionTool.from_defaults(fn=retrieve_problems)
retrieve_parts_tool = FunctionTool.from_defaults(fn=retrieve_parts)
diagnostic_tool = FunctionTool.from_defaults(fn=diagnose_car_problem)
cost_estimator_tool = FunctionTool.from_defaults(fn=estimate_repair_cost)
maintenance_retriever_tool = FunctionTool.from_defaults(fn=get_maintenance_schedule)

With the retriever tools now set up, our agent can effectively select the appropriate tool based on the query and fetch the relevant contexts. Next, we’ll create additional helper tools that will complement the existing ones, providing the agent with more context and enhancing its reasoning capabilities.

def comprehensive_diagnosis(symptoms: str) -> str:
    """
    Provides a comprehensive diagnosis including possible causes, estimated costs, and required parts.
    
    Args:
        symptoms: A string describing the car's symptoms.
    
    Returns:
        A string with a comprehensive diagnosis report.
    """
    # Use existing tools
    possible_causes = diagnose_car_problem(symptoms)
    
    # Extract the most likely cause (this is a simplification)
    likely_cause = possible_causes[0] if possible_causes else "Unknown issue"
    
    estimated_cost = estimate_repair_cost(likely_cause)
    required_parts = retrieve_parts(likely_cause)
    
    report = f"Comprehensive Diagnosis Report:\n\n"
    report += f"Symptoms: {symptoms}\n\n"
    report += f"Possible Causes:\n{possible_causes}\n\n"
    report += f"Most Likely Cause: {likely_cause}\n\n"
    report += f"Estimated Cost:\n{estimated_cost}\n\n"
    report += f"Required Parts:\n{required_parts}\n\n"
    report += "Please note that this is an initial diagnosis. For accurate results, please consult with our professional mechanic."
    
    return report

def plan_maintenance(mileage: int, car_make: str, car_model: str, car_year: int) -> str:
    """
    Creates a comprehensive maintenance plan based on the car's mileage and details.
    
    Args:
        mileage: The current mileage of the car.
        car_make: The make of the car.
        car_model: The model of the car.
        car_year: The year the car was manufactured.
    
    Returns:
        A string with a comprehensive maintenance plan.
    """
    car_details = retrieve_car_details(car_make, car_model, car_year)
    car_model_info = get_car_model_info(mileage, car_make, car_model, car_year)
    
    plan = f"Maintenance Plan for {car_year} {car_make} {car_model} at {mileage} miles:\n\n"
    plan += f"Car Details: {car_details}\n\n"
    
    if car_model_info:
        plan += f"Common Issues:\n"
        for issue in car_model_info['common_issues']:
            plan += f"- {issue}\n"
        
        plan += f"\nEstimated Time: {car_model_info['estimated_time']}\n\n"
    else:
        plan += "No specific maintenance tasks found for this car model and mileage.\n\n"
    
    plan += "Please consult with our certified mechanic for a more personalized maintenance plan."
    
    return plan

def create_calendar_invite(event_type: str, car_details: str, duration: int = 60) -> str:
    """
    Simulates creating a calendar invite for a car maintenance or repair event.
    
    Args:
        event_type: The type of event (e.g., "Oil Change", "Brake Inspection").
        car_details: Details of the car (make, model, year).
        duration: Duration of the event in minutes (default is 60).
    
    Returns:
        A string describing the calendar invite.
    """
    # Simulate scheduling the event for next week
    event_date = datetime.now() + timedelta(days=7)
    event_time = event_date.replace(hour=10, minute=0, second=0, microsecond=0)
    
    invite = f"Calendar Invite Created:\n\n"
    invite += f"Event: {event_type} for {car_details}\n"
    invite += f"Date: {event_time.strftime('%Y-%m-%d')}\n"
    invite += f"Time: {event_time.strftime('%I:%M %p')}\n"
    invite += f"Duration: {duration} minutes\n"
    invite += f"Location: Your Trusted Auto Shop, 123 Main St, Bengaluru, India\n\n"
    
    return invite

def coordinate_car_care(query: str, car_make: str, car_model: str, car_year: int, mileage: int) -> str:
    """
    Coordinates overall car care by integrating diagnosis, maintenance planning, and scheduling.
    
    Args:
        query: The user's query or description of the issue.
        car_make: The make of the car.
        car_model: The model of the car.
        car_year: The year the car was manufactured.
        mileage: The current mileage of the car.
    
    Returns:
        A string with a comprehensive car care plan.
    """
    car_details = retrieve_car_details(car_make, car_model, car_year)
    
    # Check if it's a problem or routine maintenance
    if "problem" in query.lower() or "issue" in query.lower():
        diagnosis = comprehensive_diagnosis(query)
        plan = f"Based on your query, here's a diagnosis:\n\n{diagnosis}\n\n"
        
        # Extract the most likely cause (this is a simplification)
        likely_cause = diagnosis.split("Most Likely Cause:")[1].split("\n")[0].strip()
        
        # Create a calendar invite for repair
        invite = create_calendar_invite(f"Repair: {likely_cause}", car_details)
        plan += f"I've prepared a calendar invite for the repair:\n\n{invite}\n\n"
    else:
        maintenance_plan = plan_maintenance(mileage, car_make, car_model, car_year)
        plan = f"Here's your maintenance plan:\n\n{maintenance_plan}\n\n"
        
        # Create a calendar invite for the next maintenance task
        next_task = maintenance_plan.split("Task:")[1].split("\n")[0].strip()
        invite = create_calendar_invite(f"Maintenance: {next_task}", car_details)
        plan += f"I've prepared a calendar invite for your next maintenance task:\n\n{invite}\n\n"
    
    plan += "Remember to consult with a professional mechanic for personalized advice and service."
    
    return plan

Additionally, we’ll implement some helper functions that, while not tools themselves, will be used internally within the tools to support the logic and enhance their functionality.

def get_car_model_info(mileage: int, car_make: str, car_model: str, car_year: int) -> dict:
    """Retrieve car model information from cars_models.json."""
    with open('cars_models/cars_models.json', 'r') as file:
        car_models = json.load(file)

    for car in car_models:        
        if (car['car_make'].lower() == car_make.lower() and car['car_model'].lower() == car_model.lower() and car['car_year'] == car_year):
            return car
    return {}

def retrieve_car_details(make: str, model: str, year: int) -> str:
    """Retrieves the make, model, and year of the car and return the common issues if any."""
    car_details = get_car_model_info(0, make, model, year)  # Using 0 for mileage to get general details
    if car_details:
        return f"{year} {make} {model} - Common Issues: {', '.join(car_details['common_issues'])}"
    return f"{year} {make} {model} - No common issues found."

Here are the additional tools in their complete form

comprehensive_diagnostic_tool = FunctionTool.from_defaults(fn=comprehensive_diagnosis)
maintenance_planner_tool = FunctionTool.from_defaults(fn=plan_maintenance)
calendar_invite_tool = FunctionTool.from_defaults(fn=create_calendar_invite)
car_care_coordinator_tool = FunctionTool.from_defaults(fn=coordinate_car_care)
retrieve_car_details_tool = FunctionTool.from_defaults(fn=retrieve_car_details)

Now, let’s combine all these tools into a comprehensive tools list, which we will pass to our agent to utilize.

tools = [
    retrieve_problems_tool,
    retrieve_parts_tool,
    diagnostic_tool,
    cost_estimator_tool,
    maintenance_schedule_tool,
    comprehensive_diagnostic_tool,
    maintenance_planner_tool,
    calendar_invite_tool,
    car_care_coordinator_tool,
    retrieve_car_details_tool
]

Step 4 : Creating the Agent

Now that we’ve defined the tools, we’re ready to create the agent. With LlamaIndex, this involves setting up an Agent reasoning loop. Basically, this loop allows our agent to handle complex questions that might require multiple steps or clarifications. Essentially, our agent can reason through tools and complete tasks across several stages.

LlamaIndex provides two main components for creating an agent: AgentRunner and AgentWorkers.

The AgentRunner acts as the orchestrator, like in a symphony, managing the overall process. It handles the current state, conversational memory, and tasks, and it runs steps for each task while providing a high-level user interface on what’s going on. On the other hand, AgentWorkers are responsible for the operational side. They select and use the tools and choose the LLM to interact with these tools effectively.

Now, let’s set up both the AgentRunner and AgentWorker to bring our agent to life.

# Function to reset the agent's memory
def reset_agent_memory():
    global agent_worker, agent
    agent_worker = FunctionCallingAgentWorker.from_tools(
        tools, 
        llm=llm, 
        verbose=True
    )
    agent = AgentRunner(agent_worker)

# Initialize the agent
reset_agent_memory()

Every time you call reset_agent_memory(), a new, fresh agent is created, ready to reason through and act on the user’s query.

With everything now set up—our tools, an agent for reasoning, and databases for retrieving relevant context—let’s test to see if our agent can handle simple questions effectively.

Step 5 : D-Day

Let’s ask the agent a straightforward question related to car maintenance based on the mileage count and see how well it handles it.

response = agent.query(
    "My car has 60,000 miles on it. What maintenance should I be doing now, and how much will it cost?"
)

and the response I got is

Added user message to memory: My car has 60,000 miles on it. What maintenance should I be doing now, and how much will it cost?
=== Calling Function ===
Calling function: get_maintenance_schedule with args: {"mileage": 60000}

=== Calling Function ===
Calling function: estimate_repair_cost with args: {"problem": "Oil and filter change"}

=== Calling Function ===
Calling function: estimate_repair_cost with args: {"problem": "Tire rotation"}

=== Calling Function ===
Calling function: estimate_repair_cost with args: {"problem": "Air filter replacement"}

=== Calling Function ===
Calling function: estimate_repair_cost with args: {"problem": "Brake inspection"}

=== LLM Response ===
At 60,000 miles, the recommended maintenance tasks for your car are:

1. Oil and filter change: This typically costs around $250.
2. Tire rotation: The average cost for this service is around $50.
3. Air filter replacement: This usually costs about $70.
4. Brake inspection: The cost for this can vary, but it's typically around $100.

Please note that these are average costs and can vary based on your location and the specific make and model of your car. It's always a good idea to get a few quotes from different service providers.

Well, this is amazing! The agent effectively understood the query and provided an excellent response. Notice how it first called the maintenance_schedule_tool, which utilized the get_maintenance_schedule retriever object to gather context on the relevant maintenance schedule, including different tasks based on the car’s mileage. This context was then used by the cost_estimator_tool.

The best part is that it passed the relevant parameters—problems extracted from the maintenance_schedule_tool—to the cost estimator tool, deciding on its own based on the user query. Finally, with all the gathered context, it produced a comprehensive response that perfectly addresses the user’s needs.

Btw, If you want the agent to retain the context of previous conversations, replace .query with .chat to ensure context is preserved. Keep in mind that the context size is limited by the information you provide when calling the retrievers. Watch out for the max_context_information parameter in the retrievers to avoid exceeding the token limits for the LLMs.

And that’s it! You’ve successfully created an agentic RAG that not only understands the user’s query but also delivers a well-reasoned and contextually accurate answer. Here is the colab for this example.

Zero Shot Image Classification with LanceDB and OpenAI's CLIP

2024-07-14T00:00:00+00:00

Imagine an AI having a conversation in a language it was never explicitly taught or suddenly playing a new game without any practice. In essence, if an AI can handle a task it hasn’t been directly trained for, that’s what we call zero-shot capability.

Zero-Shot classification

There are many state-of-the-art (SOTA) computer vision models that excel at various classification tasks, such as identifying animals, cars, fraud, and products in e-commerce. They can handle almost any image classification job. However, these models are often specialized and need fine-tuning for different use cases to be truly effective.

Fine-tuning can be challenging; it requires a well-labeled dataset, and if your use case is specific to an enterprise, it may also need significant computing power.

So, what does “Zero-Shot image classification” really means? Imagine a deep learning model trained only to distinguish between cats and dogs. Now, if you show it a picture of a person lounging on the couch playing video games, and the model identifies it as a “corporate employee enjoying a Sunday afternoon,” that’s zero-shot image classification. It means the model can correctly identify something it was never specifically trained to recognize. To help you follow along, here is the complete architecture..

Fundamentals

To make this work, we need a multimodal embedding model and a vector database. Let’s start with something called CLIP, which stands for Contrastive Language-Image Pre-Training. Think of CLIP as a smart box that can understand different types of files. Whether you give it an image or text, it can grasp the context behind them all.

But how it’s working behind the scenes? Consider there are two smaller boxes in that box: a Text Encoder and an Image Encoder. When OpenAI trained CLIP, they made sure these two encoders understand text and images in the same vector space.

They achieved this by training the model to place similar image-text pairs close together in vector space while separating the vectors of non-pairs. Although OpenAI hasn’t specified the exact data used, the CLIP paper mentions that the model was trained on 400 million image-text pairs collected from the internet. This extensive training gives the model an impressive ability to understand relevant image-text pairs.

So, here’s what we get from using CLIP:

Instead of datasets with specific class labels, CLIP only needs image-text pairs, where the text describes the image.
Instead of training a CNN to get features from an image, CLIP uses more expressive text descriptions, which can provide additional features.

The authors of CLIP demonstrated its superior zero-shot classification performance by comparing it to the ResNet-101 model trained specifically on ImageNet. When both models were tested on other datasets derived from ImageNet, CLIP outperformed the state-of-the-art ResNet-101, showing a better understanding of the dataset than the fine-tuned version of ResNet-101 trained on ImageNet data.

Reasoning of CLIP

So, the implementation is quite straightforward. But before going into that, Let’s just quickly understand how a CNN works.

Initially, each image in a traditional classification model has assigned class labels. We input these images into the model along with their respective class labels as the expected outputs. Through training, the model’s weights are adjusted based on calculated losses. Over time, the model learns to distinguish between various images by recognizing distinct features.

However, zero-shot classification takes this concept further by utilizing two key components: a Text Encoder and an Image Encoder. Yes those two small boxes that I described earlier, Now these encoders produce n-dimensional vectors for both images and text, mapping them to the same vector space. This means the n-dimensional vector of an image of a “cat” would be semantically similar to the vector of a text description like “a photo of a cat”.

By leveraging this shared vector space, zero-shot classification enables the model to classify images into categories it hasn’t explicitly seen during training. Instead of relying solely on predefined class labels, the model can compare the vector representation of a new image to vector representations of textual descriptions of various categories.

To enhance the effectiveness of our zero-shot classification, we should transform our class labels from simple words like “cat,” “dog,” and “horse” into more descriptive phrases such as “a photo of a cat,” “a photo of a dog,” or “a photo of a horse.” This transformation is crucial because it mirrors the text-image pairs used during the model’s pretraining phase. OpenAI used prompts like "a photo of a {label}" paired with each label to create these image-text pairs.[1]

By adopting a similar approach, our classification task aligns more closely with the model’s pretrained understanding of how images relate to their textual descriptions.

Final thoughts

Let’s take a step back and solidify our understanding before implementation. The CLIP model is pre-trained on a massive dataset of image-text pairs, learning that “a photo of a cat” corresponds to an actual image of a cat, and vice versa. This means whenever we feed an image or text into CLIP, we can expect it to grasp the relevance between the two.

Now, if you want to get into the nitty-gritty of the algorithm, it’s not overly complex. At its core, CLIP encodes each image and text as a n-dimensional embedding vector. Let’s say T1 is the vector for “a photo of a cat”, T2 for “a photo of a bird”, and T3 for “a photo of a horse”. If we have an image of a cat with embedding V1, the similarity score between V1 and T1 should be the highest among all text embeddings. This high similarity tells us that the V1 vector indeed represents “a photo of a cat”.

So, when we pass an image of a cat to our CLIP model, it should reason like “this is a cat, I know this already”. Or if we input an image of bananas on a table, it might get the nerve and put up something like “I think this image shows bananas placed on a table”. Pretty cool, right?

We’ve achieved our goal of classifying images without explicitly training a model on specific categories. And this is how CLIP does the heavy lifting for us, leveraging its pre-training to generalize to a wide range of concepts and enable zero-shot classification.

Using LanceDB

To bring our zero-shot classification system to life, we need a robust Vector Database to store our label embeddings. The process is straightforward: we’ll transform our simple text labels like “cat” into more descriptive phrases such as “a photo of a cat”, fetch their CLIP embeddings, and store these in our database. When it comes time to classify a new image, we’ll retrieve its embedding from CLIP and perform a cosine similarity calculation against all our stored label embeddings in our DB. The label with the closest match becomes our predicted class.

For this crucial task, I’ve opted for LanceDB, an impressive open-source vector database that’s like a super-smart data lake for managing complex information. LanceDB shines when we are handling complex data like our vector embeddings with an exceptional performance in fetching and storage, and the best part? It won’t cost you a dime.

But LanceDB’s appeal goes beyond just being free and open-source. Its unparalleled scalability, efficient on-disk storage, and serverless capabilities make it a standout choice. These features are part of a broader trend of columnar databases that are rapidly transforming ML workflows. I’ve actually written an in-depth article exploring the game-changing capabilities of these kind of databases. If you’re curious about how they’re revolutionizing the field, I highly recommend giving it a read!

Implementation

With all the tools at our disposal, let’s move on to a practical example of using CLIP for zero-shot image classification with the LanceDB vector database. For this demonstration, I’ll use the uoft-cs/cifar100 dataset from Hugging Face Datasets.

from datasets import load_dataset

imagedata = load_dataset(
    'uoft-cs/cifar100',
    split="test"
)

imagedata

Let’s see original label names

# labels names 
labels = imagedata.info.features['fine_label'].names
print(len(labels))
labels

100

['apple',
 'aquarium_fish',
 'baby',
 'bear',
 'beaver',
 'bed',
 'bee',
 'beetle',
 'bicycle',
 'bottle',
 'bowl',
 'boy',
 'bridge',
 'bus',
 'butterfly',
 'camel',
 'can',
 'castle',
 'caterpillar',
 'cattle',
 'chair',
 'chimpanzee',
 'clock',
 'cloud',
 'cockroach',
...
 'whale',
 'willow_tree',
 'wolf',
 'woman',
 'worm']

Looks good! We have 100 classes to classify images from, which would require a lot of computing power if you go for traditional CNN. However, let’s proceed with our zero-shot image classification approach.

Let’s generate the relevant textual descriptions for our labels

# generate sentences
clip_labels = [f"a photo of a {label}" for label in labels]
clip_labels

['a photo of a apple',
 'a photo of a aquarium_fish',
 'a photo of a baby',
 'a photo of a bear',
 'a photo of a beaver',
 'a photo of a bed',
 'a photo of a bee',
 'a photo of a beetle',
 'a photo of a bicycle',
 'a photo of a bottle',
 'a photo of a bowl',
 'a photo of a boy',
 'a photo of a bridge',
 'a photo of a bus',
 'a photo of a butterfly',
 'a photo of a camel',
 'a photo of a can',
 'a photo of a castle',
 'a photo of a caterpillar',
 'a photo of a cattle',
 'a photo of a chair',
 'a photo of a chimpanzee',
 'a photo of a clock',
 'a photo of a cloud',
 'a photo of a cockroach',
...
 'a photo of a whale',
 'a photo of a willow_tree',
 'a photo of a wolf',
 'a photo of a woman',
 'a photo of a worm']

Now let’s initialize our CLIP embedding model, I will use the CLIP implementation from hugginface.

# initialization
from transformers import CLIPProcessor, CLIPModel

model_id = "openai/clip-vit-large-patch14"

processor = CLIPProcessor.from_pretrained(model_id)
model = CLIPModel.from_pretrained(model_id)

import torch

# if you have CUDA set it to the active device like this
device = "cuda" if torch.cuda.is_available() else "cpu"
# move the model to the device
model.to(device)

If you’re new to Transformers, remember that computers understand numbers, not text. We’ll convert our text descriptions into integer representations called input IDs, where each number stands for a word or subword, more formally tokens. We’ll also need an attention mask to help the transformer focus on relevant parts of the input.

For more details, you can read about transformers here.

# create label tokens
label_tokens = processor(
    text=clip_labels,
    padding=True,
    return_tensors='pt'
).to(device)

# Print the label tokens with the corresponding text 
for i in range(5):
    token_ids = label_tokens['input_ids'][i]
    print(f"Token ID : {token_ids}, Text : {processor.decode(token_ids, skip_special_tokens=False)}")

Token ID : tensor([49406,   320,  1125,   539,   320,  3055, 49407, 49407, 49407]), Text : <|startoftext|>a photo of a apple <|endoftext|><|endoftext|><|endoftext|>
Token ID : tensor([49406,   320,  1125,   539,   320, 16814,   318,  2759, 49407]), Text : <|startoftext|>a photo of a aquarium _ fish <|endoftext|>
Token ID : tensor([49406,   320,  1125,   539,   320,  1794, 49407, 49407, 49407]), Text : <|startoftext|>a photo of a baby <|endoftext|><|endoftext|><|endoftext|>
Token ID : tensor([49406,   320,  1125,   539,   320,  4298, 49407, 49407, 49407]), Text : <|startoftext|>a photo of a bear <|endoftext|><|endoftext|><|endoftext|>
Token ID : tensor([49406,   320,  1125,   539,   320, 22874, 49407, 49407, 49407]), Text : <|startoftext|>a photo of a beaver <|endoftext|><|endoftext|><|endoftext|>

Now let’s get our CLIP embeddings for our text labels

# encode tokens to sentence embeddings from CLIP

with torch.no_grad():
    label_emb = model.get_text_features(**label_tokens) # passing the label text as in "a photo of a cat" to get it's relevant embedding from clip model

# Move embeddings to CPU and convert to numpy array
label_emb = label_emb.detach().cpu().numpy()
label_emb.shape

(100, 768)

We now have a 768-dimensional vector for each of our 100 text class sentences. However, to improve our results when calculating similarities, we need to normalize these embeddings.

Normalization helps ensure that all vectors are on the same scale, preventing longer vectors from dominating the similarity calculations simply due to their magnitude. We achieve this by dividing each vector by the square root of the sum of the squares of its elements. This process, known as L2 normalization, adjusts the length of our vectors while preserving their directional information, making our similarity comparisons more accurate and reliable.

import numpy as np

# normalization
label_emb = label_emb / np.linalg.norm(label_emb, axis=0)
label_emb.min(), label_emb.max()

Ok, let’s see a random image from our dataset

import random

index = random.randint(0, len(imagedata)-1)
selected_image = imagedata[index]['img']
selected_image

When you execute this code, you’ll be presented with a visual representation of a data point from our dataset. In my case, the output displayed a pixelated image of a whale.

Before we can analyze our image with CLIP, we need to preprocess it properly. First, we’ll run the image through our CLIP processor. This step ensures the image is resized first, then the pixels are normalized, then converting it into the tensor and finally adding a batch dimension. All of these things are settled up for the model.

image = processor(
    text=None,
    images=imagedata[index]['img'],
    return_tensors='pt'
)['pixel_values'].to(device)
image.shape

torch.Size([1, 3, 224, 224])

Now here this shape represents a 4-dimensional tensor:

1: Batch size (1 image in this case)
3: Number of color channels (Red, Green, Blue)
224: Height of the image in pixels
224: Width of the image in pixels

So, we have one image, with 3 color channels, and dimensions of 224x224 pixels. Now we’ll use CLIP to generate an embedding - a numerical representation of our image’s features. This embedding is what we’ll use for our classification task.

img_emb = model.get_image_features(image)
img_emb.shape

torch.Size([1, 768])

This gives us 768 dimensional embedding to us, that’s our Image Embedding. Only thing that is left for now is to use LanceDB to store our labels, with their corresponding embeddings and do the vector search for our Image Embedding on that database.. Here how it looks in the whole go

import lancedb
import numpy as np

data = []
for label_name, embedding in zip(labels, label_emb):
    data.append({"label": label_name, "vector": embedding})

db = lancedb.connect("./.lancedb")
table = db.create_table("zero_shot_table", data, mode="Overwrite")

# Prepare the query embedding
query_embedding = img_emb.squeeze().detach().cpu().numpy()
# Perform the search
results = (table.search(query_embedding)
           .limit(10)
           .to_pandas())

print(results.head(n=10))

|   label         | vector | distance |
|-----------------|-----------------------------------------------------------|-------------|
| whale           | [0.05180167, 0.008572296, -0.00027403078, -0.12351207, ...]| 447.551605  |
| dolphin         | [0.09493398, 0.02598409, 0.0057568997, -0.13548125, ...]| 451.570709  |
| aquarium_fish   | [-0.094619915, 0.13643932, 0.030785343, 0.12217164, ...]| 451.694672  |
| skunk           | [0.1975818, -0.04034014, 0.023241673, 0.03933424, ...]| 452.987640  |
| crab            | [0.05123004, 0.0696855, 0.016390173, -0.02554354, ...]| 454.392456  |
| chimpanzee      | [0.04187969, 0.0196794, -0.038968336, 0.10017315, ...]| 454.870697  |
| ray             | [0.10485967, 0.023477506, 0.06709562, -0.08323726, ...]| 454.880524  |
| sea             | [-0.08117988, 0.059666794, 0.09419422, -0.18542227, ...]| 454.975311  |
| shark           | [-0.01027703, -0.06132377, 0.060097754, -0.2388756, ...]| 455.291901  |
| keyboard        | [-0.18453166, 0.05200073, 0.07468183, -0.08227961, ...]| 455.424866  |

Here are the results everyone: all set and confirmed. Our initial accurate prediction is a whale, demonstrating the closest resemblance between the label and the image with minimal distance, just as we had hoped. What’s truly remarkable is that we achieved this without running a single epoch for a CNN model. That’s zero shot classification for you fellas. Here is your colab for your reference. See you in next one.

Train a CNN classification model with Lance Dataset

2024-06-26T00:00:00+00:00

In this previous post, I showed you how you can convert any Image Dataset to Lance format for faster retrieval and faster I/O operations. But can we use the same Lance formatted image dataset to train an image classification model? Well here it comes…

Lance Format: Saga for efficient image datasets

CNNs are widely used for the image related tasks in AI world. They’re great at figuring out what’s in a picture, spotting objects, and even breaking down images into meaningful parts. What makes them so useful is how they can learn important visual clues on their own, without needing humans intervention.

But when we’re dealing with massive image collections, just handling all that data can be a real headache. That’s where Lance file format comes in - it provides a clever new way to package up image data perfectly in the deep learning ecosystem for all our needs. The Lance format offers several key advantages that make it a powerful choice for machine learning applications, some of them are:

Lance uses a compressed columnar format, offering efficient storage, fast data loading, and quick random access, making it ideal for large-scale image datasets.
It supports diverse data types, including images, and text facilitating the processing of different modalities in machine learning pipelines.
Lance stores data on disk, ensuring persistence through system failures and enhancing privacy and security by allowing local storage and access.
It provides high-performance random access, up to 100 times faster than Parquet.
Lance enables vector search, finding nearest neighbors in under 1 millisecond, and integrates OLAP queries with vector search.
It features zero-copy, automatic versioning, which manages data versions automatically and reduces redundancy.

Integrating Lance and Convolutional Neural Networks

If you’ve been working with Convolutional Neural Networks (CNNs) for image classification, you know that data loading and preprocessing can be a real headache. But what if I told you there’s a way to make this process smoother and faster? Enter the Lance format.

In my previous post, I walked through the process of converting popular image datasets like cinic-10 and mini-imagenet to the Lance format. If you haven’t read that yet, I highly recommend you do so before continuing here. It’ll give you the foundation you need to fully appreciate what we’re about to dive into.

Now, let’s take the next step: using Lance-formatted data to train a CNN for image classification. We’ll use the cinic-10 dataset as our example, but the principles apply to other datasets as well.

Before we jump in, it’s important to understand how the Lance + PyTorch approach differs from the standard PyTorch method. Traditionally, PyTorch users rely on Torchvision’s ImageFolder to handle image and label loading. The Lance approach, however, requires us to create a custom dataset class. This class is designed to load binary-format images and their corresponding labels directly from the Lance dataset for creating the dataloaders.

You might be wondering, “Is it worth the effort to switch?” The answer is a resounding yes, especially if you’re dealing with large datasets or need faster training times. Lance’s secret weapon is its lightning-fast random access capability. This means that Lance dataloaders can feed data to your CNN much quicker than standard PyTorch dataloaders, potentially shaving hours off your training time.

In the following sections, we’ll dive into the details of implementing the Lance + PyTorch approach. By the end, you’ll have a powerful new tool in your deep learning toolkit that can significantly streamline your image classification workflows and reduce your model training time.

Load the Lance files to create the dataloaders

Lance-formatted image data is stored in binary format, which isn’t directly usable by Convolutional Neural Networks (CNNs). We need to convert this data into a format CNNs can process, such as PIL Image objects. Here’s the process we’ll follow:

Retrieve the binary image data: Extract the relevant image data from the Lance files.
Convert to PIL Image: Transform the binary data into a PIL Image object, creating a readable image format.
Handle grayscale images: Convert any grayscale images to RGB format for compatibility with CNNs that expect 3-channel color images.
Apply transformations: Use the provided transform function to apply necessary transformations like resizing or normalization.
Determine the labels: Look up the class index for each image’s label in the provided list of classes.
Return the data: Provide the transformed image and its corresponding label for CNN training.

To streamline this process, we’ll create a custom dataset class. This class will handle all these steps efficiently, preparing the Lance-formatted data for use with a CNN.

This custom dataset class manages all the necessary steps to prepare our Lance-formatted data for use with a CNN model. It essentially iterates over the dataset to retrieve the relevant images and labels. By using this class, we can easily integrate the Lance data into your PyTorch-based training pipeline.

# Define the image classes
classes = ('airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# transformation function 
transform_train = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((32, 32)),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

transform_val = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((32, 32)),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])

# Define the custom dataset class
class CustomImageDataset(data.Dataset):
    def __init__(self, classes, lance_dataset, transform=None):
        self.classes = classes
        self.ds = lance.dataset(lance_dataset)
        self.transform = transform

    def __len__(self):
        return self.ds.count_rows()

    def __getitem__(self, idx):
        raw_data = self.ds.take([idx], columns=['image', 'label']).to_pydict()
        img_data, label = raw_data['image'][0], raw_data['label'][0]

        img = Image.open(io.BytesIO(img_data))

        # Convert grayscale images to RGB
        if img.mode != 'RGB':
            img = img.convert('RGB')

        if self.transform:
            img = self.transform(img)

        label = self.classes.index(label)
        return img, label

Now that we have our custom Dataset class set up, we’re ready to proceed with training our model using the Lance dataset.

Using Lance dataset with CNNs: Putting It All Together

Now that we’ve created our custom dataset class, integrating Lance dataset into our CNN training process becomes straightforward. Here’s how it all comes together:

Import the custom dataset class into our CNN script.
Load the lance dataset and create lance dataloaders.
Use the lance dataloaders instead of the standard dataloaders to train our model.

From this point, the process follows a standard CNN training workflow. For our example, I’ve chosen to use ResNet-34 as our CNN architecture to enhance accuracy.

class Net(nn.Module):
    def __init__(self, num_classes):
        super(Net, self).__init__()
        self.resnet = models.resnet34(pretrained=True)
        num_ftrs = self.resnet.fc.in_features
        self.resnet.fc = nn.Linear(num_ftrs, num_classes)

    def forward(self, x):
        return self.resnet(x)

With this setup, we can train a CNN on our Lance dataset using just a single script.

One key advantage of using Lance-backed training is its performance. Compared to traditional methods, Lance-formatted data offers significant improvements in training speed. Here is the result when I compared the training time for 3 epochs with Lance vs Vanilla dataloaders

This shows an extensive improvement in training time for the Lance dataloaders as compared to the Vanilla ones.

Here is the complete notebook for the reference. For those wanting to explore further, there’s a repository showcasing various deep learning techniques that utilize Lance-formatted data. This resource can be valuable for expanding your understanding and application of Lance file format in different machine learning contexts.

Create a Movie Recommendation System using RAG and Genre Classification

2024-05-31T00:00:00+00:00

This article provides a comprehensive guide on creating a movie recommendation system by using vector similarity search and multi-label genre classification.

Here’s what we cover below:

Data ingestion and preprocessing techniques for movie metadata
Training a Doc2Vec model for Embeddings
Training a Neural network for genre classification task
Using Doc2Vec, LanceDB and the trained classifier get the relevant recommendations

Let’s get started!

Why use Embeddings for Recommendation System?

Scrolling through streaming platforms can be frustrating when the movie suggestions don’t match our interests. Building recommendation systems is a complex task, as there isn’t one metric that can measure the quality of recommendations. To improve this, we can combine embeddings and VectorDB for better recommendations.

These embeddings serve dual purposes: they can either be directly used as input to a classification model for genre classification or stored in a VectorDB for retrieval purposes. By storing embeddings in a VectorDB, efficient retrieval and query search for recommendations become possible at a later stage.

This architecture offers a holistic understanding of the underlying processes involved.

Data Ingestion and preprocessing techniques for movie metadata

Our initial task involves gathering and organizing information about movies. This includes gathering extensive details such as the movie’s type, plot summary, genres, audience ratings, and more.

Fortunately, we have access to a robust dataset on Kaggle containing information from various sources for approximately 45,000 movies. To follow along, please download the data from Kaggle and place it inside your working directory.

if you require additional data, you can supplement the dataset by extracting information from platforms like Rotten Tomatoes, IMDb, or even box-office records.

Our next step is to extract the core details from this dataset and generate a universal summary for each movie. Initially, I’ll combine the movie’s title, genres, and overviews into a single textual string. Then, this text will be tagged to create TaggedDocument instances, which will be utilized to train the Doc2Vec model later on.

Before moving forward, let’s install the relevant libraries to make our life easier.

pip install torch scikit-learn pylance lancedb nltk gensim scipy==1.12

Next, we’ll proceed with the ingestion and preprocessing of the data. To simplify the process, we’ll work with chunks of 1000 movies at a time. For clarity, we’ll only include movie indices with non-null values for genres, accurate titles, and complete overviews. This approach ensures that we’re working with high-quality, relevant data for our analysis.

import torch
import pandas as pd
import numpy as np
from gensim.models.doc2vec import Doc2Vec, TaggedDocument
from nltk.tokenize import word_tokenize
from tqdm import tqdm

import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from gensim.models.doc2vec import Doc2Vec, TaggedDocument
from nltk.tokenize import word_tokenize
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.model_selection import train_test_split
from tqdm import tqdm

import nltk
nltk.download('punkt')

# Read data from CSV file
movie_data = pd.read_csv('movies_metadata.csv', low_memory=False)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def preprocess_data(movie_data_chunk):
    tagged_docs = []
    valid_indices = []
    movie_info = []

    # Wrap your loop with tqdm
    for i, row in tqdm(movie_data_chunk.iterrows(), total=len(movie_data_chunk)):
        try:
            # Constructing movie text
            movies_text = ''
            genres = ', '.join([genre['name'] for genre in eval(row['genres'])])
            movies_text += "Overview: " + row['overview'] + '\n'
            movies_text += "Genres: " + genres + '\n'
            movies_text += "Title: " + row['title'] + '\n'
            tagged_docs.append(TaggedDocument(words=word_tokenize(movies_text.lower()), tags=[str(i)]))
            valid_indices.append(i)
            movie_info.append((row['title'], genres))
        except Exception as e:
            continue

    return tagged_docs, valid_indices, movie_info

Generating embeddings using Doc2Vec

Next, we’ll utilize the Doc2Vec model to generate embeddings for each movie based on the preprocessed text. We’ll allow the Doc2Vec model to train for several epochs to capture the essence of the various movies and their metadata in the multidimensional latent space. This process will help us represent each movie in a way that captures its unique characteristics and context.

def train_doc2vec_model(tagged_data, num_epochs=20):
    # Initialize Doc2Vec model
    doc2vec_model = Doc2Vec(vector_size=100, min_count=2, epochs=num_epochs)
    doc2vec_model.build_vocab(tqdm(tagged_data, desc="Building Vocabulary"))
    for epoch in range(num_epochs):
        doc2vec_model.train(tqdm(tagged_data, desc=f"Epoch {epoch+1}"), total_examples=doc2vec_model.corpus_count, epochs=doc2vec_model.epochs)
    
    return doc2vec_model

# Preprocess data and extract genres for the first 1000 movies
chunk_size = 1000
tagged_data = []
valid_indices = []
movie_info = []
for chunk_start in range(0, len(movie_data), chunk_size):
    movie_data_chunk = movie_data.iloc[chunk_start:chunk_start+chunk_size]
    chunk_tagged_data, chunk_valid_indices, chunk_movie_info = preprocess_data(movie_data_chunk)
    tagged_data.extend(chunk_tagged_data)
    valid_indices.extend(chunk_valid_indices)
    movie_info.extend(chunk_movie_info)

doc2vec_model = train_doc2vec_model(tagged_data)

The train_doc2vec_model function trains a Doc2Vec model on the tagged movie data, producing 100-dimensional embeddings for each movie. These embeddings act as input features for the neural network.

With our current training setup, we are sure that movies with identical genres and similar kinds of overviews will be positioned closer to each other in the latent space, reflecting their thematic and content similarities.

Extracting the unique genre labels

Next, our focus shifts to compiling the names of relevant movies along with their genres. Now won

To illustrate this, Let’s consider a movie with three genres: ‘Drama’, ‘Comedy’, and ‘Horror’. Using the MultiLabelBinarizer, we’ll represent these genres with lists of 0s and 1s. If a movie belongs to a particular genre, it will be assigned a 1; if it doesn’t, it will receive a 0. Now each row in our dataset will indicate which genres are associated with a specific movie. This approach simplifies the genre representation for easier analysis.

Let’s take the movie “Top Gun Maverick” as a reference. We’ll associate its genres using binary encoding. Suppose this movie is categorized only under ‘drama’, not ‘comedy’ or ‘horror’. When we apply the MultiLabelBinarizer, the representation would be: Drama: 1, Comedy: 0, Horror: 0. This signifies that “Top Gun Maverick” is classified as a drama but not as a comedy or horror. We’ll replicate this process for all the movies in our dataset to identify the unique genre labels present in our data.

Training a Neural Network for genre classification task

We’ll define a neural network consisting of four linear layers with ReLU activations. The final layer utilizes softmax activation to generate probability scores for various genres. If your objective is primarily classification within the genre spectrum, where you input a movie description to determine its relevant genres, you can establish a threshold value for the multi-label softmax output. This allows you to select the top ‘n’ genres with the highest probabilities.

Here’s the neural network class, hyperparameter settings, and the corresponding training loop for training our model.

# Extract genre labels for the valid indices
genres_list = []
for i in valid_indices:
    row = movie_data.loc[i]
    genres = [genre['name'] for genre in eval(row['genres'])]
    genres_list.append(genres)

mlb = MultiLabelBinarizer()
genre_labels = mlb.fit_transform(genres_list)

embeddings = []
for i in valid_indices:
    embeddings.append(doc2vec_model.dv[str(i)])

X_train, X_test, y_train, y_test = train_test_split(embeddings, genre_labels, test_size=0.2, random_state=42)

X_train_np = np.array(X_train, dtype=np.float32)
y_train_np = np.array(y_train, dtype=np.float32)
X_test_np = np.array(X_test, dtype=np.float32)
y_test_np = np.array(y_test, dtype=np.float32)

X_train_tensor = torch.tensor(X_train_np)
y_train_tensor = torch.tensor(y_train_np)
X_test_tensor = torch.tensor(X_test_np)
y_test_tensor = torch.tensor(y_test_np)

class GenreClassifier(nn.Module):
    def __init__(self, input_size, output_size):
        super(GenreClassifier, self).__init__()
        self.fc1 = nn.Linear(input_size, 512)
        self.bn1 = nn.BatchNorm1d(512)
        self.fc2 = nn.Linear(512, 256)
        self.bn2 = nn.BatchNorm1d(256)
        self.fc3 = nn.Linear(256, 128)
        self.bn3 = nn.BatchNorm1d(128)
        self.fc4 = nn.Linear(128, output_size)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(p=0.2)  # Adjust the dropout rate as needed

    def forward(self, x):
        x = self.fc1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc3(x)
        x = self.bn3(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc4(x)
        return x

# Move model to the selected device
model = GenreClassifier(input_size=100, output_size=len(mlb.classes_)).to(device)

# Define loss function and optimizer
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
epochs = 50
batch_size = 64

train_dataset = TensorDataset(X_train_tensor.to(device), y_train_tensor.to(device))
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)  # Move data to device
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * inputs.size(0)
    epoch_loss = running_loss / len(train_loader.dataset)
    print(f'Epoch [{epoch + 1}/{epochs}], Loss: {epoch_loss:.4f}')

That’s it! We’ve successfully trained a neural network for our genre classification task. Let’s test how our model is performing on the genre classification task.

from sklearn.metrics import f1_score

model.eval()
with torch.no_grad():
    X_test_tensor, y_test_tensor = X_test_tensor.to(device), y_test_tensor.to(device)  # Move test data to device
    outputs = model(X_test_tensor)
    test_loss = criterion(outputs, y_test_tensor)
    print(f'Test Loss: {test_loss.item():.4f}')


thresholds = [0.1] * len(mlb.classes_)
thresholds_tensor = torch.tensor(thresholds, device=device).unsqueeze(0)

# Convert the outputs to binary predictions using varying thresholds
predicted_labels = (outputs > thresholds_tensor).cpu().numpy()

# Convert binary predictions and actual labels to multi-label format
predicted_multilabels = mlb.inverse_transform(predicted_labels)
actual_multilabels = mlb.inverse_transform(y_test_np)

# Print the Predicted and Actual Labels for each movie
for i, (predicted, actual) in enumerate(zip(predicted_multilabels, actual_multilabels)):
    print(f'Movie {i+1}:')
    print(f'    Predicted Labels: {predicted}')
    print(f'    Actual Labels: {actual}')


# Compute F1-score
f1 = f1_score(y_test_np, predicted_labels, average='micro')
print(f'F1-score: {f1:.4f}')

# Saving the trained model
torch.save(model.state_dict(), 'trained_model.pth')

A movie recommendation system

To build a movie recommendation system, we’ll allow users to input a movie name, and based on that, we’ll return relevant recommendations. We’ll save our Doc2Vec embeddings in a vector database to achieve this. When a user inputs a query in the form of a new movie, we’ll first locate its embeddings in our vector database. Once we have this, we’ll find the ‘n’ number of movies whose embeddings are similar to our query movie. We can assess the similarity using various search algorithms like cosine similarity or finding the least Euclidean distance.

We’ll organize the data and the embeddings into a CSV file.

import lancedb
import numpy as np
import pandas as pd

data = []

for i in valid_indices:
    embedding = doc2vec_model.dv[str(i)]
    title, genres = movie_info[valid_indices.index(i)]
    data.append({"title": title, "genres": genres, "vector": embedding.tolist()})

db = lancedb.connect(".db")
tbl = db.create_table("doc2vec_embeddings", data, mode="Overwrite")
db["doc2vec_embeddings"].head()

Essentially, we establish a connection to LanceDB, set up our table, and add our movie data to it.

Each row in the table represents a single movie, with columns containing data like the title, genres, overview, and embeddings. For each movie title, we check if it exists in our dataset. If it does, we perform a cosine similarity search on all other movies and return the top 10 most relevant titles. This columnar format makes it easy to store and retrieve information for various tasks involving embeddings.

Using Doc2Vec Embeddings to get the relevant recommendations.

Our recommendation engine combines a neural network-based genre prediction model with a vector similarity search to provide relevant movie recommendations.

For a given query movie, first, we use our trained neural network to predict its genres. Based on these predicted genres, we filter our movie database to include only those movies that share at least one genre with the query movie, achieved by constructing an appropriate SQL filter.

We then perform a vector similarity search on this filtered subset to retrieve the most similar movies based on their vector representations. This approach ensures that the recommended movies are not only similar in terms of their vector characteristics but also share genre preferences with the query movie, resulting in more relevant and personalized recommendations.

# Function to get genres for a single movie query
def get_genres_for_query(model, query_embedding, mlb, thresholds, device):
    model.eval()
    with torch.no_grad():
        query_tensor = torch.tensor(query_embedding, dtype=torch.float32).unsqueeze(0).to(device)
        outputs = model(query_tensor)
        thresholds = [0.001] * len(mlb.classes_)
        thresold_tensor = torch.tensor(thresholds, device=device).unsqueeze(0)
        predicted_labels = (outputs >= thresold_tensor).cpu().numpy()
        predicted_multilabels = mlb.inverse_transform(predicted_labels)
        return predicted_multilabels


def movie_genre_prediction(movie_title):
    movie_index = movie_data.index[movie_data['title'] == movie_title].tolist()[0]
    query_embedding = doc2vec_model.dv[str(movie_index)]
    predicted_genres = get_genres_for_query(model, query_embedding, mlb, [0.1] * len(mlb.classes_), device=device)
    return predicted_genres

And now, after all the groundwork, we’ve arrived at the final piece of the puzzle. Let’s generate some relevant recommendations using embeddings and LanceDB.

def get_recommendations(title):
    pd_data = pd.DataFrame(data)
    title_vector = pd_data[pd_data["title"] == title]["vector"].values[0]
    predicted_genres = movie_genre_prediction(title)
    genres_movie = predicted_genres[0]  # Assuming predicted_genres is available

    genre_conditions = [f"genres LIKE '%{genre}%'" for genre in genres_movie]
    where_clause = " OR ".join(genre_conditions)

    result = (
        tbl.search(title_vector)
        .metric("cosine")
        .limit(10)
        .where(where_clause)
        .to_pandas()
    )
    return result[["title"]]

get_recommendations("Toy Story")

Some of the recommended movies are remarkably close matches. For example, when looking at “Toy Story,” which falls under “animation” and “family” movies, our recommendation system can find other movies in the similar genres.

That being said, here is the colab link for the complete reference.

Convert any Image dataset to Lance

2024-04-09T00:00:00+00:00

In our previous article, we explored the remarkable capabilities of the Lance format, a modern, columnar data storage solution designed to revolutionize the way we work with large image datasets in machine learning. For the same purpose, I have converted the cinic and mini-imagenet datasets to their lance versions. For this write-up, I will use the example of cinic dataset to explain how to convert any image dataset into the Lance format with a single script and unlocking the full potential of this powerful technology.

just in case, here are the cinic and mini-imagenet datasets in lance.

Processing Images

The process_images function is the heart of our data conversion process. It is responsible for iterating over the image files in the specified dataset, reading the data of each image, and converting it into a PyArrow RecordBatch object on the binary scale. This function also extracts additional metadata, such as the filename, category, and data type (e.g., train, test, or validation), and stores it alongside the image data.

def process_images(data_type):
    # Get the current directory path
    images_folder = os.path.join("cinic", data_type)

    # Define schema for RecordBatch
    schema = pa.schema([('image', pa.binary()), 
                        ('filename', pa.string()), 
                        ('category', pa.string()), 
                        ('data_type', pa.string())])

    # Iterate over the categories within each data type
    for category in os.listdir(images_folder):
        category_folder = os.path.join(images_folder, category)
        
        # Iterate over the images within each category
        for filename in tqdm(os.listdir(category_folder), desc=f"Processing {data_type} - {category}"):
            # Construct the full path to the image
            image_path = os.path.join(category_folder, filename)

            # Read and convert the image to a binary format
            with open(image_path, 'rb') as f:
                binary_data = f.read()

            image_array = pa.array([binary_data], type=pa.binary())
            filename_array = pa.array([filename], type=pa.string())
            category_array = pa.array([category], type=pa.string())
            data_type_array = pa.array([data_type], type=pa.string())

            # Yield RecordBatch for each image
            yield pa.RecordBatch.from_arrays(
                [image_array, filename_array, category_array, data_type_array],
                schema=schema
            )

By leveraging the PyArrow library, the process_images function ensures that the image data is represented in a format that is compatible with the Lance format. The use of RecordBatch objects allows for efficient data structuring and enables seamless integration with the subsequent steps of the conversion process.

One of the key features of this function is its ability to handle datasets with a hierarchical structure. It iterates over the categories within each data type, ensuring that the metadata associated with each image is accurately captured and preserved. This attention to detail is crucial, as it allows us to maintain the rich contextual information of us image dataset, which can be invaluable for tasks like classification, object detection, or semantic segmentation.

Writing to Lance

The write_to_lance function takes the data generated by the process_images function and writes it to Lance datasets, one for each data type (e.g., train, test, validation). This step is where the true power of the Lance format is unleashed.

The function first creates a PyArrow schema that defines the structure of the data to be stored in the Lance format. This schema includes the image data, as well as the associated metadata (filename, category, and data type). By specifying the schema upfront, the script ensures that the data is stored in a consistent and organized manner, making it easier to retrieve and work with in the future.

def write_to_lance():
    # Create an empty RecordBatchIterator
    schema = pa.schema([
        pa.field("image", pa.binary()),
        pa.field("filename", pa.string()),
        pa.field("category", pa.string()),
        pa.field("data_type", pa.string())
    ])

    # Specify the path where you want to save the Lance files
    images_folder = "cinic"
    
    for data_type in ['train', 'test', 'val']:
        lance_file_path = os.path.join(images_folder, f"cinic_{data_type}.lance")
        
        reader = pa.RecordBatchReader.from_batches(schema, process_images(data_type))
        lance.write_dataset(
            reader,
            lance_file_path,
            schema,
        )

Next, the function iterates through the different data types, creating a Lance dataset file for each one. The lance.write_dataset function is then used to write the RecordBatchReader, generated from the process_images function, to the respective Lance dataset files.

The benefits of this approach are numerous. By storing the data in the Lance format, you can take advantage of its columnar storage and compression techniques, resulting in significantly reduced storage requirements. Additionally, the optimized data layout and indexing capabilities of Lance enable lightning-fast data loading times, improving the overall performance and responsiveness of your machine learning pipelines.

Loading into Pandas

The final step in the process is to load the data from the Lance datasets into Pandas DataFrames, making the image data easily accessible for further processing and analysis in your machine learning workflows.

The loading_into_pandas function demonstrates this process. It first locates the Lance dataset files, created in the previous step, and creates a Lance dataset object for each data type. The function then iterates over the batches of data, converting them into Pandas DataFrames and concatenating them into a single DataFrame for each data type.

def loading_into_pandas():
    # Load Lance files from the same folder
    current_dir = os.getcwd()
    images_folder = os.path.join(current_dir, "cinic")
    
    data_frames = {}  # Dictionary to store DataFrames for each data type
    
    for data_type in ['test', 'train', 'val']:
        uri = os.path.join(images_folder, f"cinic_{data_type}.lance")

        ds = lance.dataset(uri)

        # Accumulate data from batches into a list
        data = []
        for batch in tqdm(ds.to_batches(columns=["image", "filename", "category", "data_type"], batch_size=10), desc=f"Loading {data_type} batches"):
            tbl = batch.to_pandas()
            data.append(tbl)

        # Concatenate all DataFrames into a single DataFrame
        df = pd.concat(data, ignore_index=True)
        
        # Store the DataFrame in the dictionary
        data_frames[data_type] = df
        
        print(f"Pandas DataFrame for {data_type} is ready")
        print("Total Rows: ", df.shape[0])
    
    return data_frames

This approach offers several advantages. By loading the data in batches, the function can efficiently handle large-scale image datasets without running into memory constraints. Additionally, the use of Pandas DataFrames provides a familiar and intuitive interface for working with the data, allowing you to leverage the rich ecosystem of Pandas-compatible libraries and tools for data manipulation, visualization, and analysis.

Moreover, the function stores the DataFrames in a list, indexed by the data type. This structure enables us to easily access the specific subsets of your dataset (e.g., train, test, validation) as needed, further streamlining your machine learning workflows. I mean it’s too smooth guys.

Putting It All Together

By running the provided script, you can convert your image datasets, whether they are industry-standard benchmarks or your own custom collections, into the powerful Lance format. This transformation unlocks a new level of efficiency and performance, empowering you to supercharge your machine learning projects. I have used the same script for the mini-imagenet too, make sure your data directory looks like this

here is the complete script for your reference..

import os
import pandas as pd
import pyarrow as pa
import lance
import time
from tqdm import tqdm

def process_images(data_type):
    # Get the current directory path
    images_folder = os.path.join("cinic", data_type)

    # Define schema for RecordBatch
    schema = pa.schema([('image', pa.binary()), 
                        ('filename', pa.string()), 
                        ('category', pa.string()), 
                        ('data_type', pa.string())])

    # Iterate over the categories within each data type
    for category in os.listdir(images_folder):
        category_folder = os.path.join(images_folder, category)
        
        # Iterate over the images within each category
        for filename in tqdm(os.listdir(category_folder), desc=f"Processing {data_type} - {category}"):
            # Construct the full path to the image
            image_path = os.path.join(category_folder, filename)

            # Read and convert the image to a binary format
            with open(image_path, 'rb') as f:
                binary_data = f.read()

            image_array = pa.array([binary_data], type=pa.binary())
            filename_array = pa.array([filename], type=pa.string())
            category_array = pa.array([category], type=pa.string())
            data_type_array = pa.array([data_type], type=pa.string())

            # Yield RecordBatch for each image
            yield pa.RecordBatch.from_arrays(
                [image_array, filename_array, category_array, data_type_array],
                schema=schema
            )

# Function to write PyArrow Table to Lance dataset
def write_to_lance():
    # Create an empty RecordBatchIterator
    schema = pa.schema([
        pa.field("image", pa.binary()),
        pa.field("filename", pa.string()),
        pa.field("category", pa.string()),
        pa.field("data_type", pa.string())
    ])

    # Specify the path where you want to save the Lance files
    images_folder = "cinic"
    
    for data_type in ['train', 'test', 'val']:
        lance_file_path = os.path.join(images_folder, f"cinic_{data_type}.lance")
        
        reader = pa.RecordBatchReader.from_batches(schema, process_images(data_type))
        lance.write_dataset(
            reader,
            lance_file_path,
            schema,
        )

def loading_into_pandas():
    # Load Lance files from the same folder
    current_dir = os.getcwd()
    print(current_dir)
    images_folder = os.path.join(current_dir, "cinic")
    
    data_frames = {}  # Dictionary to store DataFrames for each data type
    
    for data_type in ['test', 'train', 'val']:
        uri = os.path.join(images_folder, f"cinic_{data_type}.lance")

        ds = lance.dataset(uri)

        # Accumulate data from batches into a list
        data = []
        for batch in tqdm(ds.to_batches(columns=["image", "filename", "category", "data_type"], batch_size=10), desc=f"Loading {data_type} batches"):
            tbl = batch.to_pandas()
            data.append(tbl)

        # Concatenate all DataFrames into a single DataFrame
        df = pd.concat(data, ignore_index=True)
        
        # Store the DataFrame in the dictionary
        data_frames[data_type] = df
        
        print(f"Pandas DataFrame for {data_type} is ready")
        print("Total Rows: ", df.shape[0])
    
    return data_frames


if __name__ == "__main__":
    start = time.time()
    write_to_lance()
    data_frames = loading_into_pandas()
    end = time.time()
    print(f"Time(sec): {end - start}")

Take the different splits of the train, test and validation through different dataframes and utilize the information for your next image classifcation task

train = data_frames['train']
test = data_frames['test']
val = data_frames['val']

and this is how the training dataframe looks like

train.head()

image	filename	category	data_type
	image	filename	category	data_type
b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\...	n02130308_1836.png	cat	train
b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\...	cifar10-train-21103.png	cat	train
b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\...	cifar10-train-44957.png	cat	train
b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\...	n02129604_14997.png	cat	train
b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\...	n02123045_1463.png	cat	train

The benefits of this approach are numerous:

Storage Efficiency: The columnar storage and compression techniques employed by Lance result in significantly reduced storage requirements, making it an ideal choice for handling large-scale image datasets.
Fast Data Loading: The optimized data layout and indexing capabilities of Lance enable lightning-fast data loading times, improving the overall performance and responsiveness of your machine learning pipelines.
Random Access: The ability to selectively load specific data subsets from the Lance dataset allows for efficient data augmentation techniques and custom data loading strategies tailored to your unique requirements.
Unified Data Format: Lance can store diverse data types, such as images, text, and numerical data, in a single, streamlined format. This flexibility is invaluable in machine learning, where different modalities of data often need to be processed together.

By adopting the Lance format, we can literally elevate our machine learning workflow to new heights, unlocking unprecedented levels of efficiency, performance, and flexibility. Take the first step by running the provided script and converting your image datasets to the Lance format – the future of machine learning data management is awaiting for you, who knows if you find your second love with lance format.

Effortlessly Loading and Processing Images with Lance

2024-03-29T00:00:00+00:00

Working with large image datasets in machine learning can be challenging, often requiring significant computational resources and efficient data-handling techniques. While widely used for image storage, traditional file formats like JPEG or PNG are not optimized for efficient data loading and processing in Machine learning workflows. This is where the Lance format shines, offering a modern, columnar data storage solution designed specifically for machine learning applications.

The Lance format stores data in a compressed columnar format, enabling efficient storage, fast data loading, and fast random access to data subsets. Additionally, the Lance format is maintained on disk, which provides a couple of advantages: It will persist through a system failure and doesn’t rely on keeping everything in memory, which can run out. This also lends itself to enhanced data privacy and security, as the data doesn’t need to be transferred over a network.

One of the other key advantages of the Lance format is its ability to store diverse data types, such as images, text, and numerical data, in a unified format. Imagine having a data lake where each kind of data can be stored seamlessly without separating underlying data types. This flexibility is particularly valuable in machine learning pipelines, where different data types often need to be processed together. This unparalleled flexibility is a game-changer in machine learning pipelines, where different modalities of data often need to be processed together for tasks like multimodal learning, audio-visual analysis, or natural language processing with visual inputs.

With Lance, you can effortlessly consider all kinds of data, from images to videos and audio files to text data and numerical values, all within the same columnar storage format. This means you can have a single, streamlined data pipeline that can handle any combination of data types without the need for complex data transformations or conversions. Lance easily handles it without worrying about compatibility issues or dealing with separate storage formats for different data types. And the best part? You can store and retrieve all these diverse data types within the same column.

In contrast, while efficient for tabular data, traditional formats like Parquet may need to handle diverse data types better. By converting all data into a single, unified format using Lance, you can retrieve and process any type of data without dealing with multiple formats or complex data structures. In this article, I’ll walk through a Python code example that demonstrates how to convert a dataset of GTA5 images into the Lance format and subsequently load them into a Pandas DataFrame for further processing.

import os
import pandas as pd
import pyarrow as pa
import lance
import time
from tqdm import tqdm

We start by importing the necessary libraries, including os for directory handling, pandas for data manipulation, pyarrow for working with Arrow data formats, lance for interacting with the Lance format, and tqdm for displaying progress bars.

def process_images():
    current_dir = os.getcwd()
    images_folder = os.path.join(current_dir, "./image")

    # Define schema for RecordBatch
    schema = pa.schema([('image', pa.binary())])

    image_files = [filename for filename in os.listdir(images_folder)
                if filename.endswith((".png", ".jpg", ".jpeg"))]

    # Iterate over all images in the folder with tqdm
    for filename in tqdm(image_files, desc="Processing Images"):
        # Construct the full path to the image
        image_path = os.path.join(images_folder, filename)

        # Read and convert the image to a binary format
        with open(image_path, 'rb') as f:
            binary_data = f.read()

        image_array = pa.array([binary_data], type=pa.binary())

        # Yield RecordBatch for each image
        yield pa.RecordBatch.from_arrays([image_array], schema=schema)

The process_images function is responsible for iterating over all image files in a specified directory and converting them into PyArrow RecordBatch objects. It first defines the schema for the RecordBatch, specifying that each batch will contain a single binary column named ‘image’.

It then iterates over all image files in the directory, reads each image’s binary data, and yields a RecordBatch containing that image’s binary data.

def write_to_lance():
    schema = pa.schema([
        pa.field("image", pa.binary())
    ])

    reader = pa.RecordBatchReader.from_batches(schema, process_images())
    lance.write_dataset(
        reader,
        "image_dataset.lance",
        schema,
    )

The write_to_lance function creates a RecordBatchReader from the process_images generator and writes the resulting data to a Lance dataset named “image_dataset.lance”. This step converts the image data into the efficient, columnar Lance format, optimizing it for fast data loading and random access.

def loading_into_pandas():
    uri = "image_dataset.lance"
    ds = lance.dataset(uri)

    # Accumulate data from batches into a list
    data = []
    for batch in ds.to_batches(columns=["image"], batch_size=10):
        tbl = batch.to_pandas()
        data.append(tbl)

    # Concatenate all DataFrames into a single DataFrame
    df = pd.concat(data, ignore_index=True)
    print("Pandas DataFrame is ready")
    print("Total Rows: ", df.shape[0])

The loading_into_pandas function demonstrates how to load the image data from the Lance dataset into a Pandas DataFrame. It first creates a Lance dataset object from the “image_dataset.lance” file. Then, it iterates over batches of data, converting each batch into a Pandas DataFrame and appending it to a list. Finally, it concatenates all the DataFrames in the list into a single DataFrame, making the image data accessible for further processing or analysis.

if __name__ == "__main__":
    start = time.time()
    write_to_lance()
    loading_into_pandas()
    end = time.time()
    print(f"Time(sec): {end - start}")

The central part of the script calls the write_to_lance and loading_into_pandas functions, measuring the total execution time for the entire process. By leveraging the Lance format, this code demonstrates how to efficiently store and load large image datasets for machine learning applications. The columnar storage and compression techniques Lance uses result in reduced storage requirements and faster data loading times, making it an ideal choice for working with large-scale image data.

Moreover, the random access capabilities of Lance allow for selective loading of specific data subsets, enabling efficient data augmentation techniques and custom data loading strategies tailored to your machine learning workflow.

TLDR: Lance format provides a powerful and efficient solution for handling multimodal data in machine learning pipelines, streamlining data storage, loading, and processing tasks. By adopting Lance, we can improve our machine learning projects’ overall performance and resource efficiency while also benefiting from the ability to store diverse data types in a unified format and maintain data locality and privacy. Here is the whole script for your reference.

import os
import pandas as pd
import pyarrow as pa
import lance
import time
from tqdm import tqdm

def process_images():
    current_dir = os.getcwd()
    images_folder = os.path.join(current_dir, "./image")

    # Define schema for RecordBatch
    schema = pa.schema([('image', pa.binary())])

    image_files = [filename for filename in os.listdir(images_folder)
                if filename.endswith((".png", ".jpg", ".jpeg"))]

    # Iterate over all images in the folder with tqdm
    for filename in tqdm(image_files, desc="Processing Images"):
        # Construct the full path to the image
        image_path = os.path.join(images_folder, filename)

        # Read and convert the image to a binary format
        with open(image_path, 'rb') as f:
            binary_data = f.read()

        image_array = pa.array([binary_data], type=pa.binary())

        # Yield RecordBatch for each image
        yield pa.RecordBatch.from_arrays([image_array], schema=schema)

def write_to_lance():
    
    schema = pa.schema([
        pa.field("image", pa.binary())
    ])

    reader = pa.RecordBatchReader.from_batches(schema, process_images())
    lance.write_dataset(
        reader,
        "image_dataset.lance",
        schema,
    )

def loading_into_pandas():

    uri = "image_dataset.lance"
    ds = lance.dataset(uri)

    # Accumulate data from batches into a list
    data = []
    for batch in ds.to_batches(columns=["image"], batch_size=10):
        tbl = batch.to_pandas()
        data.append(tbl)

    # Concatenate all DataFrames into a single DataFrame
    df = pd.concat(data, ignore_index=True)
    print("Pandas DataFrame is ready")
    print("Total Rows: ", df.shape[0])


if __name__ == "__main__":
    start = time.time()
    write_to_lance()
    loading_into_pandas()
    end = time.time()
    print(f"Time(sec): {end - start}")

Imagine using Lance-formatted image data to make machine learning and deep learning projects faster. Something big is coming up, stay tuned.

Databases for ML Worflows

2024-03-15T00:00:00+00:00

In today’s world, when everyone’s curious about trying out generative AI tools and how they work, you’ve probably heard about embedded databases. Most of us tend to think about client-server based setups when databases come to mind. And honestly, that’s somewhat accurate.

However, client-server architectures aren’t really built to handle heavy analytical and ML workloads. Essentially, the processing tasks fall into two main categories: OLTP (online transactional processing) and OLAP (online analytical processing). So, when you’re changing your Instagram profile picture or uploading a photo on Facebook, you’re essentially involved in OLTP tasks, which focus on quick and easy processing. On the flip side, when we deal with OLAP, it’s all about handling complex computations such as retrieving queries from extensive datasets, combining tables, and aggregating data for big data purposes. Now, we need something that can handle our large ML workloads effectively and perform optimally across datasets ranging from small to large scales.

Columnar Oriented Datastores

Hard drives store data in terms of blocks, so whenever an operation is performed, the entire block containing the data is loaded into memory for reading by the OS. Now, Row-oriented databases aim to store whole rows of the database in the same block, whereas columnar databases store column entries in the same block.

This implies that when you need to perform column-oriented operations like updating columns, aggregations, or selecting a column entry, column-oriented databases outperform row-oriented ones in terms of speed. However, if you need to add a new data point entry with multiple columns, then row-oriented databases perform better.

Now, the point is, there’s something called Apache Arrow, which is a language-agnostic columnar memory format for flat and hierarchical data, organized for efficient analytic operations on modern hardware like CPUs and GPUs.

This is too technical. Let me break it down for you. Machine learning is all about feeding huge amounts of data into complex mathematical models to find patterns and make predictions, right? Now, Apache Arrow turbocharges this process by providing a standardized way to store and work with data that’s optimized for modern hardware like powerful GPUs. So Instead of dealing with clunky row-based formats, Arrow’s columnar layout lets you focus on the specific data features you need, drastically cutting down processing time. And since Arrow keeps data in memory, AKA RAM, rather than on sluggish disk storage like hard drives or SSDs, your models can crunch through datasets at blistering speeds. The end result? You can iterate faster, train better models, and stay ahead of the competition.

Still confused, right? I was too. Well, let’s take an example. If you’re building a model to predict housing prices based on factors like square footage, number of bedrooms, location, etc., Arrow’s columnar format would allow you to easily isolate and work with just the columns containing those specific features, ignoring any other irrelevant data columns. This focused, selective data access is more efficient than dealing with row-based formats where you’d have to sort through all the data indiscriminately.

Now that’s where the power of Arrow based columnar databases comes into play.

Lance Data Format

Building on the advantages of Apache Arrow’s columnar, in-memory approach for machine learning, there’s another game-changing data format that takes things to a whole new level – the Lance Data Format.

Designed from the ground up with modern ML workflows in mind, Lance is an absolute speed demon when it comes to querying and prepping data for training models. But it’s not just about raw speed – Lance has some seriously impressive versatility under the hood.

Unlike formats limited to tabular data, Lance can handle all kinds of data types like images, videos, 3D point clouds, audio, and more. It’s like a Swiss Army knife of data formats for ML. Btw, Don’t just take my word for it because I love LanceDB, instead – benchmarks have shown that Lance can provide random data access involving read and write operation a mind-boggling approximately 1000 times faster than Parquet, another popular columnar format. This blistering speed comes from unique storage memory layout used by Lance.

The other important thing LanceDB provides is the usage of Zero-copy versioning, essentially it means that when you create a new version of your data, LanceDB doesn’t have to make an entire copy – it just stores the changes efficiently. This saves a ton of time and storage space compared to traditional versioning methods. And optimized vector operations allow Lance to process data in bulk, taking full advantage of modern hardware like GPUs and vectorized CPUs. It’s all part of Lance’s cloud-native design.

In-process

Before understanding what Embedded Systems really do, First, we need to understand what a database management system (DBMS) is in Layman. Now in simple terms a DBMS is a software system that allows you to create, manage, and interact with databases (obviously duhh). I mean It provides a way to store, retrieve, and manipulate data in an organized and more efficient manner.

Now, an embedded database is a type of DBMS that is tightly integrated with the application layer. This means that the database that you are working with is not a separate process or service running on its own; instead, it runs within the same process as the application itself.

This tight integration makes it easier for the application to interact with the database, as there’s no need for inter-process communication or network communication unlike traditional client-server database systems where the database is a separate process or a service

This type of thing is called “in-process”. Remember this for the rest of your life. It might be the most important thing to remember when judging the other embedded databases out there.

In simple terms, the application can directly access and manipulate the database without going through additional layers or protocols. CUTTING THE BS.

On-disk storage

Remember I said earlier that Embedded systems have this ability to store the data in-memory and work at blistering speeds. While embedded databases can store data in memory, they also have this capability to store larger datasets on disk. This allows them to scale to large amounts of data (terabytes or billion dataset points, think of embeddings of billion tokens) while still providing relatively low query latencies and response times. So it’s like the best of the both worlds.

Serverless

Ok new term! Sometimes the terms “Embedded” and “Serverless” are sometimes used interchangeably in the database community, but they actually refer to different concepts. “Embedded” refers to the database being tightly integrated with the application as we seen earlier in case of Embedded databases while Serverless refers to the separation of storage and compute resources, and it’s often used in the context of microservices architectures.

To make it more concise, think of the serverless database as it’s composed of two different containers, the storage layer (where the data is stored) and the compute layer (where the data is processed). Now this separation allows the compute resources to be dynamically allocated and scaled up or down based on the workload.

And btw, when we are talking about the Serverless model, it is often associated with the cloud based services, where you don’t have to manage the underlying infra..

Scalability

Ok so you decided to work on the simple RAG system and wanted to test it out or just maybe play with it and came up with the various vector databases, you did your experiment, you are happy and done. Now when you got serious and came up with something and wanted to scale your RAG system for let’s say 1 billion embeddings, you open up your earlier setup, ingested more data, created more embeddings and when the time came, your traditional embedding database gave you nightmares in terms of the latency as well as stability.

Now, Think of an open-source embedding database designed to seamlessly handle a variety of setups and effortlessly scales up to billions of vectors locally., scales up to billions of embeddings, fetch the relevant embeddings with amazing searching capabilities and data never leaves your local machine, feels too good to be true right?! Well there is LanceDB again. I mean from the moment you dirty your hands for your next RAG system all upto the time you put something as big as production, LanceDB scales amazingly well…

Multimodality

The current embedded databases should not only support textual data but also be compatible with various file formats. It’s no longer just about text ingestion. Some time ago, I developed a GTA-5 multimodal RAG application that displays the GTA5 landscape upon entering a query. I highly recommend giving it a read to understand why the Multimodal RAG system is the new hotshot and why companies are eager to integrate it into their core. Honestly, I haven’t come across any embedded vector database other than LanceDB that can effortlessly ingest any kind of file format.

By the way, multimodality does make sense because LanceDB is built on top of the Lance format. As mentioned earlier, it supports diverse data types, including images, audio, and text, making it incredibly flexible for applications with various data formats.

Searching and Integrations

Ok, so we stored the Embeddings, we scaled our RAG too, now for a given query, we want to find the relevant embeddings, and that’s where LanceDB shines. Now Searching in LanceDB is as easy as it could be, you can just query your data in a number of ways - via SQL, full-text search , and vector search. As it supports the Hybrid search too which is one of the biggest favorable for the amazing search capabilities in LanceDB.

But it’s not just about searching guys, it integrates well enough with native Python, JavaScript/TypeScript, Pandas, Pydantic, that means you can easily integrate it with your favorite programming languages, in addition to that, it has direct integrations with cloud storage providers like AWS S3 and Azure Blob Storage. This means that we can directly query data stored on the cloud, without any added ETL steps.

Woah, I do love LanceDB

More or less, we looked at these things right :

Columnar Oriented Databases
Lance Format
In-process
On-disk storage
Serverless
Embedded systems.
Scalability
Multimodality
Searching and Integrations

Well all of them are bundled together with no cost upfront, ready to serve in one installation click and voila baby, you have your new best friend, maybe more than that, who knows? So what are you waiting for? Here is the reference, see you soon.

Create a Multimodal GTA-V RAG application

2024-03-03T00:00:00+00:00

Artificial Intelligence (AI) has been actively working with text for quite some time, but the world isn’t solely centered around words. If you take a moment to look around, you’ll find a mix of text, images, videos, audios, and their combinations.

Today we are going to work on Multimodality which is basically a concept that essentially empowers AI models with the capacity to perceive, listen, and comprehend data in diverse formats together with the text. Pretty much like how we do!

In an ideal situation, we should be able to mix different types of data together and show them to a generative AI model at the same time and iterate on it. It could be as simple as telling the AI model, “Hey, a few days ago, I sent you a picture of a brown, short dog. Can you find that picture for me?” and the model should then give us the details of that picture. Basically, we want the AI to understand things more like how we humans do, becoming really good at handling and responding to all kinds of information.

But the challenge here is to make a computer understand one data format with its related reference, and that could be a mix of text, audio, thermal imagery, and videos. Now to make this happen, we use something called Embeddings. It’s really a numeric vector which contains a bunch of numbers written together that might not mean much to us but are understood by machines very well.

Cat is equal to Cat

Let’s think of the text components for now, so we are currently aiming that our model should learn that words like “Dog” and “Cat” are closely linked to the word “Pet.” Now this understanding is easily achievable by using an embedding model which will convert these text words into their respective embeddings first and then the model is trained to follow a straightforward logic: if words are related, they are close together in the vector space, if not, they would be separated by the adequate distance.

But to help a model recognize that an image of a “Cat” and the word “Cat” are similar, we rely on Multimodal Embeddings. To simplify things a bit, imagine there is a magic box which is capable of handling various inputs – images, audios, text, and more.

Now, when we feed the box with an image of a “Cat” with the text “Cat,” it performs its magic and produces two numeric vectors. When these two vectors were given to a machine, it made machines think, “Hmm, based on these numeric values, it seems like both are connected to “Cat”. So that’s exactly what we were aiming for! Our goal was to help machines to recognize the close connection between an image of a “Cat” and the text “Cat”. However, to validate this concept, when we plot those two numeric vectors in a vector space, it turns out they are very close to each other. This outcome exactly mirrors what we observed earlier with the proximity of the two text words “Cat” and “Dog” in the vector space.

Ladies and gentlemen, that’s the essence of Multimodality. 👏

So we made our model to comprehend the association between “Cat” images and the word “Cat.” Well this is it, I mean if you are able to do this, you would have ingested the audio, images, videos as well as the word “Cat” and the model will understand how the cat is being portrayed across all kinds of file format..

RAG is here..

Well if you don’t know what RAG means, I would highly advise you to read this article here which I wrote some days back and loved by tons of people, not exaggerating it but yeah, it’s good to get the balls rolling..

So there are impressive models like DALLE-2 that provide text-to-image functionality. Essentially, you input text, and the model generates relevant images for you. But can we create a system similar to Multimodal RAG, where the model produces output images based on our own data? Alright, so the goal for today is to create an AI model that when asked something like, “How many girls were there in my party?” 💀 not only provides textual information but also includes a relevant image related to it. Think of it as an extension of a simple RAG system, but now incorporating images.

Before we dive in, remember that Multimodality isn’t limited to just text-to-image or image-to-text as it encompasses the freedom to input and output any type of data. However, for now, let’s concentrate on the interaction from image to text exclusively.

Contrastive learning

Now the question is, What exactly was that box doing? The magic it performs is known as Contrastive Learning. While the term might sound complex, it’s not that tricky. To simplify, consider a dataset with images, along with a caption describing what the image represents.

Alright, now what happens is: we give our text-image model with these Positive and Negative samples, where each sample consists of an image and a descriptive text. Positive samples are those where the image and text are correctly aligned – for instance, a picture of a cat matched with the text “this is an image of a cat.” Conversely, negative samples involve a mismatch, like presenting an image of a dog alongside the text “this is an image of a cat.”

Now we train our text-image model to recognize that positive samples offer accurate interpretations, while negative samples are misleading and should be disregarded during training. In formal terms this technique is called CLIP (Contrastive Language-Image Pre-training) introduced by OpenAI where authors trained an image-text model on something around 400 million image caption pairs taken from the internet and everytime model makes a mistake, the contrastive loss function increases and penalize it to make sure the model trains well. The same kind of principles are applied to the other modality combinations as well, so the voice of cat with the word cat is a positive sample for speech-text model, a video of cat with the descriptive text “this is a cat” is a positive sample for video-text model.

Show time

Well you don’t have to build that box from scratch because folks have already done it for us. There’s a Multimodal embedding model, like the “ViT-L/14” from OpenAI. This model can handle various data types, including text, images, videos, audios, and even thermal and gyroscope data. Now, onto the next question: how do we store those embeddings?

For that we’ll need a vector database that can efficiently fetch, query, and retrieve relevant embeddings for us, ideally one that supports multimodal data and doesn’t burn a hole in our wallets. That’s where LanceDB comes into play.

Vector database

When we talk about the vector database, there are ton of options available in the current market, but there is something about the LanceDB which makes it stands out as an optimal choice for a vector database, As far as I have used it, it address the limitations of traditional embedded databases in handling AI/ML workloads. When I say traditional, it typically means those database management tools which are not aligned with the usage of heavy computation that comes with the ML infra.

TLDR; LanceDB operates on a serverless architecture, meaning storage and compute are separated into two distinct units. This design makes it exceptionally fast for RAG use cases, ensuring fast fetching and retrieval. Additionally, it has some notable advantages – being open source, utilizing its Lance columnar data format built on top of Apache Arrow for high efficiency, persistent storage capabilities, and incorporating its own Disk Approximate Nearest Neighbor search. All these factors collectively make LanceDB an ideal solution for accessing and working with multimodal data. I love you LanceDB ❤️.

Data time

To add some excitement, I’ve crafted a GTA-V Image Captioning dataset, featuring thousands of images, each paired with a descriptive text illustrating the image’s content. Now, when we train our magic box, the expectation is clear – if I ask that box to provide me an image of “road with a stop sign,” it should deliver a GTA-V image of a road with a stop sign on it. Otherwise, what’s the point, right?

FAQ

We will be using “ViT-L/14” to convert our multimodal data into its respective embeddings.
LanceDB as our vector database to store the relevant embeddings.
GTA-V Image Captioning dataset for our magic box.

Environment Setup

I am using a MacBook Air M1, and it’s important to note that some kinds of dependencies and configurations may vary depending on the type of system that you are running, so it’s important to take that into account.

Here are the steps to install the relevant dependencies

# Create a virtual environment
python3 -m venv env

# Activate the virtual environment
source env/bin/activate

# Upgrade pip in the virtual environment
pip install --upgrade pip

# Install required dependencies
pip3 install lancedb clip torch datasets pillow 
pip3 install git+https://github.com/openai/CLIP.git

And don’t forget to get your access token from the hugging face to download the data.

Downloading the Data

Dataset can easily be fetched using the datasets library.

import clip
import torch
import os
from datasets import load_dataset

ds = load_dataset("vipulmaheshwari/GTA-Image-Captioning-Dataset")
device = torch.device("mps")
model, preprocess = clip.load("ViT-L-14", device=device)

Downloading the dataset may require some time, so please take a moment to relax while this process completes. Once the download is finished, you can visualize some sample points like this:

from textwrap import wrap
import matplotlib.pyplot as plt
import numpy as np

def plot_images(images, captions):
    plt.figure(figsize=(15, 7))
    for i in range(len(images)):
        ax = plt.subplot(1, len(images), i + 1)
        caption = captions[i]
        caption = "\n".join(wrap(caption, 12))
        plt.title(caption)
        plt.imshow(images[i])
        plt.axis("off")

# Assuming ds is a dictionary with "train" key containing a list of samples
sample_dataset = ds["train"]
random_indices = np.random.choice(len(sample_dataset), size=2, replace=False)
random_indices = [index.item() for index in random_indices]

# Get the random images and their captions
random_images = [np.array(sample_dataset[index]["image"]) for index in random_indices]
random_captions = [sample_dataset[index]["text"] for index in random_indices]

# Plot the random images with their captions
plot_images(random_images, random_captions)

# Show the plot
plt.show()

Storing the Embeddings

The dataset consists of two key features: the image and its corresponding descriptive text. Initially, our task is to create a LanceDB table to store the embeddings. This process is straightforward – all you need to do is define the relevant schema. In our case, the columns include “vector” for storing the multimodal embeddings, a “text” column for the descriptive text, and a “label” column for the corresponding IDs.

import pyarrow as pa
import lancedb
import tqdm

db = lancedb.connect('./data/tables')
schema = pa.schema(
  [
      pa.field("vector", pa.list_(pa.float32(), 512)),
      pa.field("text", pa.string()),
      pa.field("id", pa.int32())
  ])
tbl = db.create_table("gta_data", schema=schema, mode="overwrite")

Executing this will generate a table with the specified schema, and it’s ready to store the embeddings along with the relevant columns. It’s as straightforward as that – almost too easy!

Encode the Images

Now, we’ll simply take the images from the dataset, feed them into an encoding function that leverages our Multimodal Embedding model, and generate the corresponding embeddings. These embeddings will then be stored in the database.

def embed_image(img):
    processed_image = preprocess(img)
    unsqueezed_image = processed_image.unsqueeze(0).to(device)
    embeddings = model.encode_image(unsqueezed_image)
    
    # Detach, move to CPU, convert to numpy array, and extract the first element as a list
    result = embeddings.detach().cpu().numpy()[0].tolist()
    return result

So our embed_image function takes an input image, prepocesses it through our CLIP model preprocessor, encode the preprocessed image and returns a list representing the embeddings of that image. This returned embedding serves as a concise numerical representation, capturing all the key features and patterns within the image for downstream tasks or analysis. The next thing is to call this function for all the images and store the relevant embeddings in the database.

data = []
for i in range(len(ds["train"])):
    img = ds["train"][i]['image']
    text = ds["train"][i]['text']
    
    # Encode the image
    encoded_img = embed_image(img)
    data.append({"vector": encoded_img, "text": text, "id" : i})

Here, we’re just taking a list, adding the numeric embeddings, reference text and the current index id to it. All that’s left is to include this list in our LanceDB table. And voila, our datalake for the embeddings is set up and good to go!

tbl.add(data)
tbl.to_pandas()

Up until now, we’ve efficiently converted the images into their respective multimodal embeddings and stored them in the LanceDB table. Now the LanceDB tables offer a convenient feature: if there’s a need to add or remove images, it’s remarkably straightforward. Just encode the new image and add it, following to the same steps we followed for the previous images.

Query search

Our next move is to embed our text query using the same multimodal embedding model we used for our images. Remember that “box” I mentioned earlier? Essentially, we want this box to create embeddings for both our images and our texts which ensures that the representation of different types of data happens in the same way. Following this, we just need to initiate a search to find the nearest image embeddings that matches our text query.

def embed_txt(txt):
    tokenized_text = clip.tokenize([txt]).to(device)
    embeddings = model.encode_text(tokenized_text)
    
    # Detach, move to CPU, convert to numpy array, and extract the first element as a list
    result = embeddings.detach().cpu().numpy()[0].tolist()
    return result

res = tbl.search(embed_txt("a road with a stop")).limit(3).to_pandas()
res

0 | [0.064575195, .. ] | there is a stop sign...| 569 |	131.995728
1 | [-0.07989502, .. ] | there is a bus that... | 423 | 135.047852
2 | [0.06756592, .. ]  | amazing view of a ...	| 30  | 135.309937

Let’s slow down a bit and understand what just happened. Putting simply, the code snippet executes a search algorithm in its core to pinpoint the most relevant image embedding that aligns with our text query. The resulting output, as showcased above, gives us the embeddings which closely resembles our text query. In the result, the second column presents the embedding vector, while the third column contains the description of the image that closely matches our text query. Essentially, we’ve determined which image closely corresponds to our text query by examining the embeddings of both our text query and the image.

It’s similar to saying, If these numbers represent the word “Cat”, I spot an image with a similar set of numbers, so most likely it’s a match for an image of a “Cat”. 😺

If you are looking for the explanation of how the search happens, I will write a detailed explanation in the coming write ups because it’s so exciting to look under the hood and see how the searching happens. Essentially there is something called Approximate Nearest Neighbors (ANN) which is a technique used to efficiently find the closest points in high-dimensional spaces. ANN is extensively used in data mining, machine learning, computer vision and NLP use cases. So when we passed our embedded text query to the searching algorithm and asked it to give us the closest sample point in the vector space, it used a type of ANN algorithm to get it for us. Specifically LanceDB utilizes DANN (Deep Approximate Nearest Neighbor) for searching the relevant embeddings within its ecosystem..

In our results, we have five columns. The first is the index number, the second is the embedding vector, the third is the description of the image matching our text query, and the fourth is the label of the image. However, let’s focus on the last column – Distance. When I mentioned the ANN algorithm, it simply draws a line between the current data point (in our case, the embedding of our text query) and identifies which data point (image embedding) is closest to it. If you observe that the other data points in the results have a greater distance compared to the top one, it indicates they are a bit further away or more unrelated to our query. Just to make it clear, the calculation of distance is a part of the algorithm itself.

D-DAY

Now that we have all the necessary information, displaying the most relevant image for our query is straightforward. Simply take the relevant label of the top-matched embedding vector and showcase the corresponding image.

data_id = int(res['id'][0])
display(ds["train"][data_id]['image'])
print(ds["train"][data_id]['text'])

there is a truck driving down a street with a stop sign

What’s next?

To make things more interesting, I’m currently working on creating an extensive GTA-V captioning dataset. This dataset will include a larger number of images paired with their respective reference text, providing us with a richer set of queries to explore and experiment with.. Nevertheless, there’s always room for refining the model. We can explore creating a customized CLIP model, adjusting various parameters. Increasing the number of training epochs may afford the model more time to grasp the relevance between embeddings. Additionally, there’s an impressive multimodal embedding model developed by the Meta known as ImageBind. We can consider trying ImageBind as an alternative to our current multimodal embedding model and compare the outcomes. With numerous options available, the fundamental concept behind the Multimodal RAG workflow remains largely consistent.

Here’s how everything comes together in one frame and this is the Collab for your reference

Create LLM apps using RAG

2024-02-14T00:00:00+00:00

If you’re considering making a personalized bot for your documents or your website that responds to you, you’re in the right spot. I’m here to help you create a bot using Langchain and RAG strategies for this purpose.

Understanding the Limitations of ChatGPT and LLMs

ChatGPTs and other Large Language Models (LLMs) are extensively trained on text corpora to comprehend language semantics and coherence. Despite their impressive capabilities, these models have limitations that require careful consideration for particular use cases. One significant challenge is the potential for hallucinations, where the model might generate inaccurate or contextually irrelevant information.

Imagine requesting the model to enhance your company policies; in such scenarios, ChatGPTs and other Large Language Models might struggle to provide factual responses because they lack training on your company’s data. Instead, they may generate nonsensical or irrelevant responses, which can be unhelpful. So, how can we ensure that an LLM comprehends our specific data and generates responses accordingly? This is where techniques like Retrieval Augmentation Generation (RAG) come to the rescue.

What is RAG?

RAG or Retrieval Augmented Generation uses three main workflows to generate and give the better response

Information Retrieval: When a user asks a question, the AI system retrieves the relevant data from a well-maintained knowledge library or external sources like databases, articles, APIs, or document repositories. This is achieved by converting the query into a numerical format or vector that can be understood by machines.
LLM:The retrieved data is then presented to the LLM or Large Language Model, along with the user’s query. The LLM uses this new knowledge and its training data to generate the response.
Response: Finally, the LLM generates a response that is more accurate and relevant since it has been augmented with the retrieved information. I mean we gave LLM some additional information from our Knowledge library which allows LLMs to provide more contextually relevant and factual responses, solving the problem of models when they are just hallucinating or providing irrelevant answers.

Let’s take the example of company policies again. Suppose you have an HR bot that handles queries related to your Company policies. Now if someone asks anything specific to the policies, The bot can pull the most recent policy documents from the knowledge library, pass the relevant context to a well crafted prompt which is then passed further to the LLM for generating the response.

To make it easier, Imagine a LLM as your knowledgeable friend who seems to know everything, from Geography to Computer Science, from Politics to Philosophy. Now, picture yourself asking this friend a few questions:

“Who handles my laundry on weekends?”
“Who lives next door to me?”
“What brand of peanut butter do I prefer?”

Chances are, your friend wouldn’t be able to answer these questions. Most of the time, no. But let’s say this distant friend becomes closer to you over time, he comes to your place regularly, knows your parents very well, you both hangout pretty often, you go on outings, blah blah blah.. You got the point.

I mean he is gaining access to personal and insider information about you. Now, when you pose the same questions, he can somehow answer those questions with more relevance now because he is better suited with your personal insights.

Similarly, a LLM, when provided with additional information or access to your data, won’t guess or hallucinate. Instead, it can leverage that access data to provide more relevant and accurate answers.

To break it down, here are the exact steps to create any RAG application…

Extract the relevant information from your data sources.
Break the information into the small chunks.
Store the chunks as their embedddings into a vector database.
Create a prompt template which will be fed to the LLM with the query and the context.
Convert the query to it’s relevant embedding using same embedding model.
Fetch k number of relevant documents related to the query from the vector database.
Pass the relevant documents to the LLM and get the response.

FAQs

We will be using Langchain for this task, Basically it’s like a wrapper which lets you talk and manage your LLM operations better. Note that the Langchain is updating very fast and some functions and other classes might moved to the different modules. So if something doesn’t work, just check if you are importing the libraries from the right sources!
Along with it we will be using Hugging Face, an open-source library for building, training, and deploying state-of-the-art machine learning models, especially about NLP. To use the HuggingFace we need the access token, Get your access token here
For our models, we’ll need two key components: a LLM (Large Language Model) and an embedding model. While paid sources like OpenAI offer these, we’ll be utilizing open-source models to ensure accessibility for everyone.
Now we need a Vector Database to store our embeddings, For that task, we’ve got LanceDB – it’s like a super-smart data lake for handling lots of information. It’s a top-notch vector database, making it the go-to choice for dealing with complex data like vector embeddings.. And the best part? It won’t burn a dent in your pocket because it’s open source and free to use!!
To keep things simple, our data ingestion process will involve using a URL and some PDFs. While you can incorporate additional data sources if needed, we’ll concentrate solely on these two for now.

With Langchain for the interface, Hugging Face for fetching the models, along with open-source components, we’re all set to go! This way, we will save some bucks while still having everything we need. Let’s move to the next steps

Environment Setup

I am using a MacBook Air M1, and it’s important to note that certain dependencies and configurations may vary depending on the type of system you are using. Now open your favorite editor, create a python environment and install the relevant dependencies

# Create a virtual environment
python3 -m venv env

# Activate the virtual environment
source env/bin/activate

# Upgrade pip in the virtual environment
pip install --upgrade pip

# Install required dependencies
pip3 install lancedb langchain langchain_community prettytable sentence-transformers huggingface-hub bs4 pypdf pandas

# This is optional, I did it for removing a warning
pip3 uninstall urllib3
pip3 install 'urllib3<2.0'

Now create a .env file in the same directory to place your Hugging Face api credentials like this

HUGGINGFACEHUB_API_TOKEN = hf_KKNWfBqgwCUOHdHFrBwQ.....

Ensure the name “HUGGINGFACEHUB_API_TOKEN” remains unchanged, as it is crucial for authentication purposes.

If you prefer a straightforward approach without relying on external packages or file loading, you can directly configure the environment variable within your code like this..

HF_TOKEN = "hf_KKNWfBqgwCUOHdHFrBwQ....."
os.environ["HUGGINGFACEHUB_API_TOKEN"] = HF_TOKEN

Finally create a data folder in the project’s root directory, designated as the central repository for storing PDF documents. You can add some sample PDFs for testing purposes; for instance, I am using the Yolo V7 and Transformers paper for demonstration. It’s important to note that this designated folder will function as our primary source for data ingestion.

It seems like everything is in order, and we’re all set!

Step 1 : Extracting the relevant information

To get your RAG application running, the first thing we need to do is to extract the relevant information from the various data sources. It can be a website page, a PDF file, a notion link, a google doc whatever it is, it needs to be extracted from it’s original source first.

import os
from langchain_community.document_loaders import WebBaseLoader, PyPDFLoader, DirectoryLoader

# Put the token values inside the double quotes
HF_TOKEN = "hf_*******"
os.environ["HUGGINGFACEHUB_API_TOKEN"] = HF_TOKEN

# Loading the web url and data 
url_loader = WebBaseLoader("https://gameofthrones.fandom.com/wiki/Jon_Snow")
documents_loader = DirectoryLoader('data', glob="./*.pdf", loader_cls=PyPDFLoader)

# Creating the instances
url_docs = url_loader.load()
data_docs = documents_loader.load()

# Combining all the data that we ingested
docs = url_docs + data_docs

This will ingest all the data from the URL link and the PDFs.

Step 2 : Breaking the information into smaller chunks

We’ve all the necessary data for developing our RAG application. Now, it’s time to break down this information into smaller chunks. Later, we’ll utilize an embedding model to convert these chunks into their respective embeddings. But why is it important?

Think of it like this: If you’re tasked with digesting a 100-page book all at once and then asked a specific question about it, it would be challenging to retrieve the necessary information from the entire book to provide an answer. However, if you’re permitted to break the book into smaller, manageable chunks—let’s say 10 pages each—and each chunk is labeled with an index from 0 to 9, the process becomes much simpler. When the same question is posed after this breakdown, you can easily locate the relevant chunk based on its index and then extract the information needed to answer the question accurately.

Picture the book as your extracted information, with each 10-page segment representing a small chunk of data, and the index pages as the embedding. Essentially, we’ll apply an embedding model to these chunks to transform the information into their respective embeddings. While as humans, we may not directly comprehend or relate to these embeddings, they serve as numeric representations of the chunks to our application. This is how you can do this in Python

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(chunk_size = 1000, chunk_overlap = 50)
chunks = text_splitter.split_documents(docs)

Now the chunk_size parameter specifies the maximum number of characters that a chunk can contain, while the chunk_overlap parameter specifies the number of characters that should overlap between two adjacent chunks. With the chunk_overlap set to 50, the last 50 characters of the adjacent chunks will be shared between each other.

This approach helps to prevent important information from being split across two chunks, ensuring that each chunk contains sufficient contextual information of their neighbor chunks for the subsequent processing or analysis.

Shared information at the boundary of neighboring chunks enables a more seamless transition and understanding of the text’s content. The best strategy for choosing the chunk_size and chunk_overlap parameters largely depends on the nature of the documents and the purpose of the application.

Step 3 : Creating the embeddings and store them into a vectordatabase

There are two primary methods to generate embeddings for our text chunks. The first involves downloading a model, managing preprocessing, and conducting computations independently. Alternatively, we can leverage Hugging Face’s model hub, which offers a variety of pre-trained models for various NLP tasks, including embedding generation.

Opting for the latter approach allows us to utilize one of Hugging Face’s embedding models. With this method, we simply provide our text chunks to the chosen model, saving us from the resource-intensive computations on our local machines. 💀

Hugging Face’s model hub provides numerous options for embedding models, and you can explore the leaderboard to select the most suitable one for your requirements. For now, we’ll proceed with “sentence-transformers/all-MiniLM-L6-v2.” This model is pretty fast and highly efficient in our task!!

from langchain_community.embeddings import HuggingFaceEmbeddings

embedding_model_name = 'sentence-transformers/all-MiniLM-L6-v2'
embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name, model_kwargs={'device': 'cpu'})

Here’s a way to see the number of embeddings for each chunk

query = "Hello I want to see the length of the embeddings for this document."
len(embeddings.embed_documents([query])[0])

# 384

We have the embeddings for our chunks, now we need a vector database to store them.

When it comes to vector databases, there are plenty of options out there suiting various needs. Databases like Pinecone offer adequate performance and advanced features but come with a hefty price tag. On the other hand, open-source alternatives like FAISS or Chroma may lack some extras but are more than sufficient for those not who don’t require extensive scalability.

But wait, I am dropping a bomb here, I’ve recently come across LanceDB, So it’s an open-source vector database similar to FAISS and Chroma. What makes LanceDB stand out is not just its open-source nature but its unparalleled scalability. In fact, after a closer look, I realized that I haven’t done justice in highlighting the true value propositions of LanceDB earlier!!

Surprisingly, LanceDB is the most scalable vector database available, outperforming even the likes of Pinecone, Chroma, Qdrant, and others. Scaling up to a billion vectors locally on your laptop is a feat only achievable with LanceDB. I mean this capability is a game-changer, especially when you compare it to other vector databases struggling even with a hundred million vectors. What’s more mind blowing is that LanceDB manages to offer this unprecedented scalability at a fraction of cost, I mean they are offering the utilities and database tools at much cheaper rates than its closest counterparts.

So now, We’ll create an instance of LanceDB vector database by calling lancedb.connect("lance_database"). This line essentially sets up a connection to the LanceDB database named “lance_database.” Next, we create a table within the database named “rag_sample” using the create_table function. Now we initialzed this table with a single data entry which includes a numeric vector generated by the embed_query function. So text “Hello World” is first converted to it’s numeric representation (fancy name of embeddings) and then it’s mapped to id number 1. Like a key-value pair. Lastly, the mode=”overwrite” parameter ensures that if the table “rag_sample” already exists, it will be overwritten with the new data.

This happens with all the text chunks and it’s quite straightforward. This is how it looks in Python..

import lancedb
from langchain_community.vectorstores import LanceDB

db = lancedb.connect("lance_database")
table = db.create_table(
    "rag_sample",
    data=[
        {
            "vector": embeddings.embed_query("Hello World"),
            "text": "Hello World",
            "id": "1",
        }
    ],
    mode="overwrite",
)

docsearch = LanceDB.from_documents(chunks, embeddings, connection=table)

NO ROCKET SCIENCE HA!

Step 4 : Create a prompt template which will be fed to the LLM

Ok now comes the prompt template. So when you write a question to the ChatGPT and it answers that question, you are basically providing a prompt to the model so that it can understand what the question is. When companies train the models, they decide what kind of prompt they are going to use for invoking the model and ask the question. For example, if you are working with “Mistral 7B instruct” and you want the optimal results it’s recommended to use the following chat template:

<s>[INST] Instruction [/INST] Model answer</s>[INST] Follow-up instruction [/INST]

Note that <s> and </s> are special tokens to represent beginning of string (BOS) and end of string (EOS) while [INST] and [/INST] are regular strings. It’s just that the Mistral 7B instruct is made in such a way that the model looks for those special tokens to understand the question better. Different types of LLMs have different kinds of instructed prompts.

Now for our case we are going to use huggingfaceh4/zephyr-7b-alpha which is a text generation model. Just to make it clear, Zephyr-7B-α has not been aligned or formated to human preferences with techniques like RLHF (Reinforcement Learning with Human Feedback) or deployed with in-the-loop filtering of responses like ChatGPT, so the model can produce problematic outputs (especially when prompted to do so).

Instead of writing a Prompt of our own, I will use ChatPromptTemplate class which creates a prompt template for the chat models. In layman terms, instead of writing a specified prompt I am letting ChatPromptTemplate to do it for me. Here is an example prompt template that is being generated from the manual messsages.

from langchain_core.prompts import ChatPromptTemplate

chat_template = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a helpful AI bot. Your name is {name}."),
        ("human", "Hello, how are you doing?"),
        ("ai", "I'm doing well, thanks!"),
        ("human", "{user_input}"),
    ]
)

messages = chat_template.format_messages(name="Bob", user_input="What is your name?")

If you don’t want to write the manual instructions, you can just use the from_template function to generate a more generic prompt template which I used for this project. Here it is..

from langchain_core.prompts import ChatPromptTemplate

template = """
{query}
"""

prompt = ChatPromptTemplate.from_template(template)

Our prompt is set! We’ve crafted a single message, assuming it’s from a human xD . If you’re not using the from_messages function, the ChatPromptTemplate will ensure your prompt works seamlessly with the language model by reserving some additional system messages. There’s always room for improvement with more generic prompts to achieve better results. For now this setup should work..

Step 5 : Convert the query to it’s relevant embedding using same embedding model.

Now, let’s talk about the query or question we want to ask our RAG application. We can’t just pass the query to our model and expect information in return. Instead, we need to pass the query through the same embedding model used for the chunks earlier. Why is this important? Well, by embedding queries, we allow models to compare them efficiently with previously processed chunks of text. This enables tasks like finding similar documents or generating relevant responses.

To understand it better, Imagine you and your friend speak different languages, like English and Hindi, and you need to understand each other’s writings. If your friend hands you a page in Hindi, you won’t understand it directly. So, your friend translates it first, turning the Hindi into English for you. So now if your friend asks you a question in Hindi, you can easily translate that question into English first and look up for the relevant answers in that translated English Text..

Similarly, we initially transformed textual information into their corresponding embeddings. Now, when you pose a query, it undergoes a similar kind of conversion into the numeric form using the same embedding model applied previously to process our textual chunks. This consistent approach allows for efficient retrieval of relevant responses.

Step 6 : Fetch K number of documents.

Now, let’s talk about the retriever. Its job is to dive into the vector database and perform a search to find relevant documents. It returns a set number, let’s call it “k”, of these documents, which are ranked based on their contextual relevance to the query or question you asked. You can set “k” as a parameter, indicating how many relevant documents you want - whether it’s 2, 5, or 10. Generally, if you have a smaller amount of data, it’s best to stick with a lower “k”, around 2. For longer documents or larger datasets, a “k” between 10 and 20 is recommended.

Different search techniques can be employed to fetch relevant documents more effectively and quickly from a vector database. The choice depends on various factors such as your specific use case, the amount of data you have, what kind of vector database you are using and the context of your problem.

retriever = docsearch.as_retriever(search_kwargs={"k": 3})
docs = retriever.get_relevant_documents("what did you know about Yolo V7?")
print(docs)

When you run this code, the retriever will fetch 3 most most relevant documents from the vector database. All these documents will be the contexts for our LLM model to generate the response for our query.

Step 7 : Pass the relevant documents to the LLM and get the response.

So far, we’ve asked our retriever to fetch a set number of relevant documents from the database. Now, we need a language model (LLM) to generate a relevant response based on that context. To ensure robustness, let’s remember that at the beginning of this blog, I mentioned that LLMs like ChatGPT can sometimes generate irrelevant responses, especially when asked about specific use cases or contexts. However, this time, we’re providing the context from our own data to the LLM as a reference. So, it will consider this reference along with its general capabilities to answer the question. That’s the whole idea behind using RAG!

Now, let’s dive into implementing the language model (LLM) aspect of our RAG setup. We’ll be using zephyr model architecture from the Hugging Face Hub. Here’s how we do it in Python:

from langchain_community.llms import HuggingFaceHub

# Model architecture
llm_repo_id = "huggingfaceh4/zephyr-7b-alpha"
model_kwargs = {"temperature": 0.5, "max_length": 4096, "max_new_tokens": 2048}
model = HuggingFaceHub(repo_id=llm_repo_id, model_kwargs=model_kwargs)

In this code snippet, we’re instantiating our language model using the Hugging Face Hub. Specifically, we’re selecting the zephyr 7 billion model which is placed in this repository ID “huggingfaceh4/zephyr-7b-alpha”. Choosing this model isn’t arbitrary; as I said before, it’s based on the model’s suitability for our specific task and requirements. As we are already implementing only Open Source components, Zephyr 7 billion works good enough to generate a useful response with minimal overhead and low latency.

This model comes with some additional parameters to fine-tune its behavior. We’ve set the temperature to 0.5, which controls the randomness of the generated text. As a lower temperature tends to result in more conservative and predictable outputs and when the temperature is set to max which is 1, the model tries to be as creative as it could, so based on what type of output you want for your use case, you can tweak this parameter. For the sake of the simplicity and demonstration purposes, I set it to 0.5 to make sure we get decent results. Next is max_length parameter which defines the maximum length of the generated text and it includes the size of your prompt as well as the response.

max_new_tokens sets the threshold on the maximum number of new tokens that can be generated. As a general rule of thumb, the max_new_tokens should always be less than or equal to the max_length parameter. Why? Think about it..

Step 8 : Create a chain for invoking the LLM.

We have everything we want for our RAG application. The last thing we need to do is to create a chain for invoking the LLM on our query to generate the response. There are different types of chains for the different types of use cases, if you like your LLM to remember the context of the chat over the time like the ChatGPT , you would need a memory instance which can be shared among multiple conversation pieces, for such cases, there are conversational chains available.

For now we just need a chain which can combine our retrieved contexts and pass it with the query to the LLM to generate the response.

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

rag_chain = (
    {"context": retriever,  "query": RunnablePassthrough()}
    | prompt
    | model
    | StrOutputParser()
)

response = rag_chain.invoke("Who killed Jon Snow?")

We have our Prompt, model, context and the query! All of them are combined into a single chain. It’s pretty much what all the chains does! Now before running the final code, I want to give a quick check on these two helper functions: RunnablePassthrough() and StrOutputParser().

The RunnablePassthrough class in LangChain serves to pass inputs unchanged or with additional keys. In our chain, a prompt expects input in the form of a map with keys “context” and “question.” However, user input only includes the “question.” or the “query”. Here, RunnablePassthrough is utilized to pass the user’s question under the “question” key while retrieving the context using a retriever. It just ensures that the input to the prompt conforms to the expected format.

Secondally, StrOutputParser is typically employed in RAG chains to parse the output of the model into a human-readable string. In the layman terms, It is responsible for transforming the model’s output into a more coherent and grammatically correct sentence, which is generally better readable by Humans! That’s it!

D-Day

To make sure we get the entire idea even if the response gets cut off, I’ve implemented a function called get_complete_sentence(). Basically this function helps extract the last complete sentence from the text. So, even if the response hits the maximum token limit that we set upon and it gets truncated midway, we will still get a coherent understanding of the message.

For practical testing, I suggest storing some low sized PDFs in the data folder of your project. You can choose PDFs related to various topics or domains that you want the chatbot to interact with. Additionally, providing a URL as a reference for the chatbot can be helpful for testing. For example, you could use a Wikipedia page, a research paper, or any other online document relevant to your testing goals. During my testing, I used a URL containing information about Jon Snow from Game of Thrones, and PDFs of Transformers paper, and the YOLO V7 paper to evaluate the bot’s performance. Let’s see how our bot performs in varied content.

import os
import time
import lancedb
from langchain_community.vectorstores import LanceDB

from langchain_community.llms import HuggingFaceHub
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import LanceDB
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import WebBaseLoader, PyPDFLoader, DirectoryLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from prettytable import PrettyTable

HF_TOKEN = "hf*********"
os.environ["HUGGINGFACEHUB_API_TOKEN"] = HF_TOKEN

# Loading the web URL and breaking down the information into chunks
start_time = time.time()

loader = WebBaseLoader("https://gameofthrones.fandom.com/wiki/Jon_Snow")
documents_loader = DirectoryLoader('data', glob="./*.pdf", loader_cls=PyPDFLoader)

# URL loader
url_docs = loader.load()

# Document loader
data_docs = documents_loader.load()

# Combining all the information into a single variable
docs = url_docs + data_docs

# Specify chunk size and overlap
chunk_size = 256
chunk_overlap = 20
text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
chunks = text_splitter.split_documents(docs)

# Specify Embedding Model
embedding_model_name = 'sentence-transformers/all-MiniLM-L6-v2'
embeddings = HuggingFaceEmbeddings(model_name=embedding_model_name, model_kwargs={'device': 'cpu'})

# Specify Vector Database
vectorstore_start_time = time.time()
database_name = "LanceDB"
db = lancedb.connect("src/lance_database")
table = db.create_table(
    "rag_sample",
    data=[
        {
            "vector": embeddings.embed_query("Hello World"),
            "text": "Hello World",
            "id": "1",
        }
    ],
    mode="overwrite",
)
docsearch = LanceDB.from_documents(chunks, embeddings, connection=table)
vectorstore_end_time = time.time()

# Specify Retrieval Information
search_kwargs = {"k": 3}
retriever = docsearch.as_retriever(search_kwargs = {"k": 3})

# Specify Model Architecture
llm_repo_id = "huggingfaceh4/zephyr-7b-alpha"
model_kwargs = {"temperature": 0.5, "max_length": 4096, "max_new_tokens": 2048}
model = HuggingFaceHub(repo_id=llm_repo_id, model_kwargs=model_kwargs)

template = """
{query}
"""

prompt = ChatPromptTemplate.from_template(template)

rag_chain_start_time = time.time()
rag_chain = (
    {"context": retriever, "query": RunnablePassthrough()}
    | prompt
    | model
    | StrOutputParser()
)
rag_chain_end_time = time.time()

def get_complete_sentence(response):
    last_period_index = response.rfind('.')
    if last_period_index != -1:
        return response[:last_period_index + 1]
    else:
        return response

# Invoke the RAG chain and retrieve the response
rag_invoke_start_time = time.time()
response = rag_chain.invoke("Who killed Jon Snow?")
rag_invoke_end_time = time.time()

# Get the complete sentence
complete_sentence_start_time = time.time()
complete_sentence = get_complete_sentence(response)
complete_sentence_end_time = time.time()

# Create a table
table = PrettyTable()
table.field_names = ["Task", "Time Taken (Seconds)"]

# Add rows to the table
table.add_row(["Vectorstore Creation", round(vectorstore_end_time - vectorstore_start_time, 2)])
table.add_row(["RAG Chain Setup", round(rag_chain_end_time - rag_chain_start_time, 2)])
table.add_row(["RAG Chain Invocation", round(rag_invoke_end_time - rag_invoke_start_time, 2)])
table.add_row(["Complete Sentence Extraction", round(complete_sentence_end_time - complete_sentence_start_time, 2)])

# Additional information in the table
table.add_row(["Embedding Model", embedding_model_name])
table.add_row(["LLM (Language Model) Repo ID", llm_repo_id])
table.add_row(["Vector Database", database_name])
table.add_row(["Temperature", model_kwargs["temperature"]])
table.add_row(["Max Length Tokens", model_kwargs["max_length"]])
table.add_row(["Max New Tokens", model_kwargs["max_new_tokens"]])
table.add_row(["Chunk Size", chunk_size])
table.add_row(["Chunk Overlap", chunk_overlap])
table.add_row(["Number of Documents", len(docs)])


print("\nComplete Sentence:")
print(complete_sentence)

# Print the table
print("\nExecution Timings:")
print(table)

To enhance readability and present the execution information in a structured tabular format, I have used PrettyTable library. You can add it to your virtual environment by using the command pip3 install prettytable.

So this is the response I received in less than < 1 minute, which is quite considerable for the starters. The time it takes can vary depending on your system’s configuration, but I believe you’ll get decent results in just a few minutes. So, please be patient if it’s taking a bit longer.

Human: 
Question : Who killed Jon Snow?

Answer: 
In the TV series Game of Thrones, Jon Snow was stabbed by his 
fellow Night's Watch members in season 5, episode 9, 
"The Dance of Dragons." However, he was later resurrected by Melisandre 
in season 6, episode 3, "Oathbreaker." So, technically, 
no one killed Jon Snow in the show.


Execution Timings:
+------------------------------+----------------------------------------+
|             Task             |          Time Taken (Seconds)          |
+------------------------------+----------------------------------------+
|     Vectorstore Creation     |                 16.21                  |
|       RAG Chain Setup        |                  0.03                  |
|     RAG Chain Invocation     |                  2.06                  |
| Complete Sentence Extraction |                  0.0                   |
|       Embedding Model        | sentence-transformers/all-MiniLM-L6-v2 |
| LLM (Language Model) Repo ID |     huggingfaceh4/zephyr-7b-alpha      |
|       Vector Database        |                LanceDB                 |
|         Temperature          |                  0.5                   |
|      Max Length Tokens       |                  4096                  |
|        Max New Tokens        |                  2048                  |
|          Chunk Size          |                  256                   |
|        Chunk Overlap         |                   20                   |
|     Number of Documents      |                   39                   |
+------------------------------+----------------------------------------+

Have fun experimenting with various data sources! You can try changing the website addresses, adding new PDF files or maybe change the template a bit. LLMs are fun, you never know what you get!

What’s next?

There are plenty of things we can adjust here. We could switch to a more effective embedding model for better indexing, try different searching techniques for the retriever, add a reranker to improve the ranking of documents, or use a more advanced LLM with a larger context window and faster response times. Essentially, every RAG application is just an enhanced version based on these factors. However, the fundamental concept of how a RAG application works always remains the same.

Here is the collab link for the reference..

Jyosa Labs Blogs

FP8, KV Cache Quantization, and the "!!!!" Bug

part 1: floating point fundamentals

the mantissa (where your actual digits live)

the implicit leading 1 (free bit hack)

subnormal numbers (filling the gap near zero)

FP8 E4M3: the format that started the whole mess

why is the max value exactly 448?

FP8 E4M3 has NO infinity!

E4M3 vs E5M2: two flavors of FP8

part 2: the vLLM KV cache bug

the problem

vLLM’s scale calculation: the root cause

why one-time scale calibration is dangerous

the clipping: where things start going wrong

the identical values problem

layernorm variance collapse: the NaN factory

NaN propagation: one bad apple spoils the batch

why ALL future requests fail too

the complete causal chain

the fix

takeaways

Attention is all you need guys

here we go!

the sequence problem or why rnns kinda suck

well what really is a transformer?!

attention baby!

the core idea (why attention is actually genius)

positional embeddings : gives you the wings

how it actually works (the full picture)

multi-head attention (because one perspective isn’t enough)

feed-forward network (the other half of the layer)

the autoregressive thing (why decoders are slightly different)

a quick note on kv-cache (the inference trick)

Create a custom retriever

Why Custom Retrievers Matter

Implementation Breakdown

Part 1: Core Dependencies and Imports

Part 2: Understanding LlamaIndex Custom Retrievers

BaseRetriever Abstraction

But Why This Matters?

Plugging Superlinked into LlamaIndex

Part 3: Superlinked Schema Definition and Setup

Part 4: Vector Space Configuration

Part 5: Data Pipeline and Executor Setup

Part 6: The Retrieval Engine

Part 7: Result Processing and Node Creation

Show Time: Executing the pipeline

Real Estate Natural Language Querying Agent

The Problem with Current Real Estate Search

What’s Natural Language Querying All About?

What it took to made it..

Data Requirements

How I Built This Thing

Step 1: Vector Spaces for Property Magic

Step 2: Weighted Query Construction

Why Agents + NLQ for Real Estate Search?

Step 3: The Agentic Layer - Where the Magic Happens

PropertyRetrievalTool

PropertyRecommendationTool

QueryRefinementTool

NarrativeInsightTool

MultiStepQueryTool

Step 4: Intent Classification – Figuring Out What the User Really Wants

Example 1: Basic Search

Example 2: Personalised Recommendations

Example 3: Query Refinement

Example 4: Market Insight Generation

Example 5: Comparative Search

An agent to read research papers better

What’s the Plan?

Step 1 : Setting up the toolbox

Step 2 : Data needs

Step 3 : Defining the Superlinked Schema

Step 4 : Building the index

Step 5 : Crafting the query

Step 6 : Building tools

Step 7 : Building the Kernel Agent

Creating a Fintech AI agent without any framework

What’s the deal here?