Part 4: Improving Retrieval Quality with Token-Aware Chun...

Making RAG Smarter with Token-Aware Chunking, HyDE, and Context-Aware Search

In Part 3, we improved chunking and optimized context. The system was faster and cleaner… but still not always correct.

What broke after Part 3?

By this point, the system looked solid:

Smarter chunking
Context compression
FAISS + re-ranking
Streaming responses

But when I started using it more realistically, a few problems showed up:

1. Token limits were still hurting quality

Even with better chunks, we were still not controlling how much context we send to the model.

2. Vague queries failed badly

Questions like:

“Explain this”
“What does it mean?”

…would often retrieve irrelevant chunks.

3. Follow-up questions felt disconnected

The system didn’t “remember” what we were talking about.

At this point, it stopped feeling like a “retrieval problem”
…and more like a context understanding problem.

So in Part 4, I focused on making RAG smarter.

What we’re building in this part

Token-aware chunking (based on actual LLM limits)
HyDE (Hypothetical Document Embeddings)
Early version of context-aware retrieval

Updated Pipeline

Before jumping into code, here’s how the pipeline evolved:

Before (Part 3):

Query → Embedding → FAISS → Re-rank → Context → LLM

Now (Part 4):

Query → (HyDE) → Better Query → Embedding  
     → FAISS → Re-rank → Token-aware Context  
     → LLM (with constraints)

Problem 1: Token limits are real (and we were ignoring them)

Until now, chunking was based on:

characters
sentences
separators

But LLMs don’t think in characters.

They think in tokens.

Why this matters

You might send:

2 chunks → fine
5 chunks → maybe fine
10 chunks → silently truncated or degraded

And the worst part?

You won’t even know it’s happening.

Solution: Token-aware chunking

Instead of guessing chunk sizes, we measure them using a tokenizer.

import tiktoken

tokenizer = tiktoken.get_encoding("cl100k_base")

def get_token_length(text):
    return len(tokenizer.encode(text))

Now chunking becomes token-driven instead of size-driven.

Token-based chunking strategy

Key idea:

Build chunks until a token limit
If exceeded → split intelligently
Maintain overlap using tokens, not characters

MAX_TOKENS = 250
OVERLAP_TOKENS = 50

Smarter chunk building

Instead of blindly splitting:

Prefer paragraphs
Then sentences
Then fallback

def generate_chunks_recursive_tokens(text, page_num):
    paragraphs = text.split("\n\n")
    current_chunk = []
    current_tokens = 0

    for paragraph in paragraphs:
        paragraph_tokens = get_token_length(paragraph)

        if current_tokens + paragraph_tokens > MAX_TOKENS:
            chunks.append({
                "text": "\n\n".join(current_chunk),
                "page": page_num
            })

            current_chunk, current_tokens = _get_overlap(current_chunk)

        current_chunk.append(paragraph)
        current_tokens += paragraph_tokens

Why this works better

Matches actual LLM limits
Avoids hidden truncation
Improves context density
Makes responses more reliable

Problem 2: RAG fails on vague queries

This was the bigger issue.

Even with good chunking, queries like:

“Explain this concept”

…don’t contain enough semantic signal.

So FAISS retrieves something… but often not the right thing.

Solution: HyDE (Hypothetical Document Embeddings)

This is one of the most interesting tricks in RAG.

Instead of embedding the raw query…

👉 We first generate a hypothetical answer
👉 Then embed that

Why this works

A vague query becomes a rich semantic representation.

Example:

User query:

Explain this

HyDE generates:

This concept refers to a method where...

Now embedding this gives:

More keywords
Better semantic alignment
Stronger retrieval

Implementation

Step 1: Generate hypothetical answer

def generate_hypothetical_answer(query, chat_history):
    recent_user = [m["content"] for m in chat_history[-4:] if m["role"] == "user"]
    history_text = " | ".join(recent_user[-2:])

    prompt = (
        f"Write a 2-sentence technical summary answering: {query}\n"
        f"Recent user context: {history_text}"
    )

    response = ollama.generate(
        model=HYDE_MODEL,
        prompt=prompt,
        stream=False
    )

    return response['response']

Step 2: Augment the query

hypothetical_answer = generate_hypothetical_answer(query, chat_history)
search_query = f"{query} {hypothetical_answer}"

Step 3: Embed the enriched query

response = ollama.embed(model=EMBED_MODEL, input=search_query)

Impact

Better retrieval for vague queries
Improved relevance
More stable responses

At the cost of:

~1–2 seconds extra latency

Worth it? In most cases — yes.

Early Step: Context-aware retrieval

Another issue we started addressing:

Follow-up questions were treated as completely new queries.

Example:

User: Explain attention mechanism
User: How does it work?

Second query loses context.

What we added

A simple but effective improvement:

Detect vague follow-ups
Inject previous page context

if not target_page and vague_followups.match(query.strip()):
    last_pages = get_last_referenced_pages(chat_history)

    if last_pages:
        query = f"{query} page {last_pages[0]}"

Result

Follow-ups become meaningful
Retrieval stays anchored
Conversation feels connected

Putting it all together

Now the system:

Understands token limits
Improves weak queries
Handles basic conversation flow

From learning project → real system

This is where things started to change.

Earlier:

It worked
It demonstrated RAG

Now:

It behaves more like a real assistant
Handles imperfect queries
Works under constraints

What’s next?

We’ve improved:

Data representation (chunks)
Query understanding (HyDE)

But one big gap still remains:

We still don’t know why RAG fails when it fails.

In Part 5, we’ll go deeper into:

Debugging the RAG pipeline
Visualizing FAISS vs re-ranking
Understanding retrieval quality
Making the system more transparent

Code

Full implementation available here:

👉 https://github.com/SharathKurup/chatPDF/blob/token_aware_rag/

Final thoughts

At a high level, RAG seems simple:

Retrieve → augment → generate

But in practice, most of the work is here:

How you represent data
How you interpret queries
How you control context

This part was about tightening those pieces.

And it makes a noticeable difference.

If you’ve been building with RAG, I’d recommend trying:

Token-aware chunking
HyDE

They’re relatively small changes — but high impact.

Let me know what you think or what you’d improve.

推荐订阅源

DEV Community