A Coding Tutorial of Model Context Protocol Focusing on Semantic Chunking, Dynamic Token Management, and Context Relevance Scoring for Efficient LLM Interactions

Managing context effectively is a critical challenge when working with large language models, especially in environments like Google Colab, where resource constraints and long documents can quickly exceed available token windows. In this tutorial, we guide you through a practical implementation of the Model Context Protocol (MCP) by building a ModelContextManager that automatically chunks incoming […] The post A Coding Tutorial of Model Context Protocol Focusing on Semantic Chunking, Dynamic Token Management, and Context Relevance Scoring for Efficient LLM Interactions appeared first on MarkTechPost.

Apr 28, 2025 - 08:03
 0
A Coding Tutorial of Model Context Protocol Focusing on Semantic Chunking, Dynamic Token Management, and Context Relevance Scoring for Efficient LLM Interactions

Managing context effectively is a critical challenge when working with large language models, especially in environments like Google Colab, where resource constraints and long documents can quickly exceed available token windows. In this tutorial, we guide you through a practical implementation of the Model Context Protocol (MCP) by building a ModelContextManager that automatically chunks incoming text, generates semantic embeddings using Sentence-Transformers, and scores each chunk based on recency, importance, and relevance. You’ll learn how to integrate this manager with a Hugging Face sequence-to-sequence model, demonstrated here with FLAN-T5, to add, optimize, and retrieve only the most pertinent pieces of context. Along the way, we’ll cover token counting with a GPT-2 tokenizer, context-window optimization strategies, and interactive sessions that let you query and visualize your dynamic context in real time.

import torch
import numpy as np
from typing import List, Dict, Any, Optional, Union, Tuple
from dataclasses import dataclass
import time
import gc
from tqdm.notebook import tqdm

We import essential libraries for building a dynamic context manager: torch and numpy handle tensor and numerical operations, while typing and dataclasses provide structured type annotations and data containers. Utility modules, such as time and gc, support timestamping and memory cleanup, as well as tqdm.notebook offers interactive progress bars for chunk processing in Colab.

@dataclass
class ContextChunk:
    """A chunk of text with metadata for the Model Context Protocol."""
    text: str
    embedding: Optional[torch.Tensor] = None
    importance: float = 1.0
    timestamp: float = 0.0
    metadata: Dict[str, Any] = None
   
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}
        if self.timestamp == 0.0:
            self.timestamp = time.time()

The ContextChunk dataclass encapsulates a single segment of text along with its embedding, a user-assigned importance score, a timestamp, and arbitrary metadata. Its __post_init__ method ensures that each chunk is stamped with the current time upon creation and that metadata defaults to an empty dictionary if none is provided.

class ModelContextManager:
    """
    Manager for implementing Model Context Protocol in LLMs on Google Colab.
    Handles context window optimization, token management, and relevance scoring.
    """
   
    def __init__(
        self,
        max_context_length: int = 8192,
        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
        relevance_threshold: float = 0.7,
        recency_weight: float = 0.3,
        importance_weight: float = 0.3,
        semantic_weight: float = 0.4,
        device: str = "cuda" if torch.cuda.is_available() else "cpu"
    ):
        """
        Initialize the Model Context Manager.
       
        Args:
            max_context_length: Maximum number of tokens in context window
            embedding_model: Model to use for text embeddings
            relevance_threshold: Threshold for chunk relevance to be included
            recency_weight: Weight for recency in relevance calculation
            importance_weight: Weight for importance in relevance calculation
            semantic_weight: Weight for semantic similarity in relevance calculation
            device: Device to run computations on
        """
        self.max_context_length = max_context_length
        self.device = device
        self.chunks = []
        self.current_token_count = 0
        self.relevance_threshold = relevance_threshold
       
        self.recency_weight = recency_weight
        self.importance_weight = importance_weight
        self.semantic_weight = semantic_weight
       
        try:
            from sentence_transformers import SentenceTransformer
            print(f"Loading embedding model {embedding_model}...")
            self.embedding_model = SentenceTransformer(embedding_model).to(self.device)
            print(f"Embedding model loaded successfully on {self.device}")
        except ImportError:
            print("Installing sentence-transformers...")
            import subprocess
            subprocess.run(["pip", "install", "sentence-transformers"])
            from sentence_transformers import SentenceTransformer
            self.embedding_model = SentenceTransformer(embedding_model).to(self.device)
            print(f"Embedding model loaded successfully on {self.device}")
           
        try:
            from transformers import GPT2Tokenizer
            self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
        except ImportError:
            print("Installing transformers...")
            import subprocess
            subprocess.run(["pip", "install", "transformers"])
            from transformers import GPT2Tokenizer
            self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
   
    def add_chunk(self, text: str, importance: float = 1.0, metadata: Dict[str, Any] = None) -> None:
        """
        Add a new chunk of text to the context manager.
       
        Args:
            text: The text content to add
            importance: Importance score (0-1)
            metadata: Additional metadata for the chunk
        """
        with torch.no_grad():
            embedding = self.embedding_model.encode(text, convert_to_tensor=True)
       
        chunk = ContextChunk(
            text=text,
            embedding=embedding,
            importance=importance,
            timestamp=time.time(),
            metadata=metadata or {}
        )
       
        self.chunks.append(chunk)
        self.current_token_count += len(self.tokenizer.encode(text))
       
        if self.current_token_count > self.max_context_length:
            self.optimize_context()
   
    def optimize_context(self) -> None:
        """Optimize context by removing less relevant chunks to fit within token limit."""
        if not self.chunks:
            return
           
        print("Optimizing context window...")
       
        scores = self.score_chunks()
       
        sorted_indices = np.argsort(scores)[::-1]
       
        new_chunks = []
        new_token_count = 0
       
        for idx in sorted_indices:
            chunk = self.chunks[idx]
            chunk_tokens = len(self.tokenizer.encode(chunk.text))
           
            if new_token_count + chunk_tokens <= self.max_context_length:
                new_chunks.append(chunk)
                new_token_count += chunk_tokens
            else:
                if scores[idx] > self.relevance_threshold * 1.5:
                    for i, included_chunk in enumerate(new_chunks):
                        included_idx = sorted_indices[i]
                        if scores[included_idx] < self.relevance_threshold:
                            included_tokens = len(self.tokenizer.encode(included_chunk.text))
                            if new_token_count - included_tokens + chunk_tokens <= self.max_context_length:
                                new_chunks.remove(included_chunk)
                                new_token_count -= included_tokens
                                new_chunks.append(chunk)
                                new_token_count += chunk_tokens
                                break
       
        removed_count = len(self.chunks) - len(new_chunks)
        self.chunks = new_chunks
        self.current_token_count = new_token_count
       
        print(f"Context optimized: Removed {removed_count} chunks, {len(new_chunks)} remaining, using {new_token_count}/{self.max_context_length} tokens")
       
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
   
    def score_chunks(self, query: str = None) -> np.ndarray:
        """
        Score chunks based on recency, importance, and semantic relevance.
       
        Args:
            query: Optional query to calculate semantic relevance against
           
        Returns:
            Array of scores for each chunk
        """
        if not self.chunks:
            return np.array([])
           
        current_time = time.time()
        max_age = max(current_time - chunk.timestamp for chunk in self.chunks) or 1.0
        recency_scores = np.array([
            1.0 - ((current_time - chunk.timestamp) / max_age)
            for chunk in self.chunks
        ])
       
        importance_scores = np.array([chunk.importance for chunk in self.chunks])
       
        if query is not None:
            query_embedding = self.embedding_model.encode(query, convert_to_tensor=True)
            similarity_scores = np.array([
                torch.cosine_similarity(chunk.embedding, query_embedding, dim=0).item()
                for chunk in self.chunks
            ])
           
            similarity_scores = (similarity_scores - similarity_scores.min()) / (similarity_scores.max() - similarity_scores.min() + 1e-8)
        else:
            similarity_scores = np.ones(len(self.chunks))
       
        final_scores = (
            self.recency_weight * recency_scores +
            self.importance_weight * importance_scores +
            self.semantic_weight * similarity_scores
        )
       
        return final_scores
   
    def retrieve_context(self, query: str = None, k: int = None) -> str:
        """
        Retrieve the most relevant context for a given query.
       
        Args:
            query: The query to retrieve context for
            k: The maximum number of chunks to return (None = all relevant chunks)
           
        Returns:
            String containing the combined relevant context
        """
        if not self.chunks:
            return ""
           
        scores = self.score_chunks(query)
       
        relevant_indices = np.where(scores >= self.relevance_threshold)[0]
       
        relevant_indices = relevant_indices[np.argsort(scores[relevant_indices])[::-1]]
       
        if k is not None:
            relevant_indices = relevant_indices[:k]
           
        relevant_texts = [self.chunks[i].text for i in relevant_indices]
        return "\n\n".join(relevant_texts)
   
    def get_stats(self) -> Dict[str, Any]:
        """Get statistics about the current context state."""
        return {
            "chunk_count": len(self.chunks),
            "token_count": self.current_token_count,
            "max_tokens": self.max_context_length,
            "usage_percentage": self.current_token_count / self.max_context_length * 100 if self.max_context_length else 0,
            "avg_chunk_size": self.current_token_count / len(self.chunks) if self.chunks else 0,
            "oldest_chunk_age": time.time() - min(chunk.timestamp for chunk in self.chunks) if self.chunks else 0,
        }


    def visualize_context(self):
        """Visualize the current context window distribution."""
        try:
            import matplotlib.pyplot as plt
            import pandas as pd
           
            if not self.chunks:
                print("No chunks to visualize")
                return
           
            scores = self.score_chunks()
            chunk_sizes = [len(self.tokenizer.encode(chunk.text)) for chunk in self.chunks]
            timestamps = [chunk.timestamp for chunk in self.chunks]
            relative_times = [time.time() - ts for ts in timestamps]
            importance = [chunk.importance for chunk in self.chunks]
           
            df = pd.DataFrame({
                'Size (tokens)': chunk_sizes,
                'Age (seconds)': relative_times,
                'Importance': importance,
                'Score': scores
            })
           
            fig, axs = plt.subplots(2, 2, figsize=(14, 10))
           
            axs[0, 0].bar(range(len(chunk_sizes)), chunk_sizes)
            axs[0, 0].set_title('Token Distribution by Chunk')
            axs[0, 0].set_ylabel('Tokens')
            axs[0, 0].set_xlabel('Chunk Index')
           
            axs[0, 1].scatter(chunk_sizes, scores)
            axs[0, 1].set_title('Score vs Chunk Size')
            axs[0, 1].set_xlabel('Tokens')
            axs[0, 1].set_ylabel('Score')
           
            axs[1, 0].scatter(relative_times, scores)
            axs[1, 0].set_title('Score vs Chunk Age')
            axs[1, 0].set_xlabel('Age (seconds)')
            axs[1, 0].set_ylabel('Score')
           
            axs[1, 1].scatter(importance, scores)
            axs[1, 1].set_title('Score vs Importance')
            axs[1, 1].set_xlabel('Importance')
            axs[1, 1].set_ylabel('Score')
           
            plt.tight_layout()
            plt.show()
           
        except ImportError:
            print("Please install matplotlib and pandas for visualization")
            print('!pip install matplotlib pandas')

The ModelContextManager class orchestrates the end-to-end handling of context for LLMs by chunking input text, generating embeddings, and tracking token usage against a configurable limit. It implements relevance scoring (combining recency, importance, and semantic similarity), automatic context pruning, retrieval of the most pertinent chunks, and convenient utilities for monitoring and visualizing context statistics.

class MCPColabDemo:
    """Demonstration of Model Context Protocol in Google Colab with a Language Model."""
   
    def __init__(
        self,
        model_name: str = "google/flan-t5-base",
        max_context_length: int = 2048,
        device: str = "cuda" if torch.cuda.is_available() else "cpu"
    ):
        """
        Initialize the MCP Colab demo with a specified model.
       
        Args:
            model_name: Hugging Face model name
            max_context_length: Maximum context length for the MCP manager
            device: Device to run the model on
        """
        self.device = device
        self.context_manager = ModelContextManager(
            max_context_length=max_context_length,
            device=device
        )
       
        try:
            from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
            print(f"Loading model {model_name}...")
            self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            print(f"Model loaded successfully on {device}")
        except ImportError:
            print("Installing transformers...")
            import subprocess
            subprocess.run(["pip", "install", "transformers"])
            from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
            self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            print(f"Model loaded successfully on {device}")
   
    def add_document(self, text: str, chunk_size: int = 512, overlap: int = 50) -> None:
        """
        Add a document to the context by chunking it appropriately.
       
        Args:
            text: Document text
            chunk_size: Size of each chunk in characters
            overlap: Overlap between chunks in characters
        """
        chunks = []
        for i in range(0, len(text), chunk_size - overlap):
            chunk = text[i:i + chunk_size]
            if len(chunk) > 20:  
                chunks.append(chunk)
       
        print(f"Adding {len(chunks)} chunks to context...")
        for i, chunk in enumerate(tqdm(chunks)):
            pos = i / len(chunks)
            importance = 1.0 - 0.5 * min(pos, 1 - pos)
           
            self.context_manager.add_chunk(
                text=chunk,
                importance=importance,
                metadata={"source": "document", "position": i, "total_chunks": len(chunks)}
            )
   
    def process_query(self, query: str, max_new_tokens: int = 256) -> str:
        """
        Process a query using the context manager and model.
       
        Args:
            query: The query to process
            max_new_tokens: Maximum number of tokens in response
           
        Returns:
            Model response
        """
        self.context_manager.add_chunk(query, importance=1.0, metadata={"type": "query"})
       
        relevant_context = self.context_manager.retrieve_context(query=query)
       
        prompt = f"Context: {relevant_context}\n\nQuestion: {query}\n\nAnswer:"
       
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
       
        print("Generating response...")
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                temperature=0.7,
                top_p=0.9,
            )
       
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
       
        self.context_manager.add_chunk(
            response,
            importance=0.9,
            metadata={"type": "response", "query": query}
        )
       
        return response
   
    def interactive_session(self):
        """Run an interactive session in the notebook."""
        from IPython.display import clear_output
       
        print("Starting interactive MCP session. Type 'exit' to end.")
        conversation_history = []
       
        while True:
            query = input("\nYour query: ")
           
            if query.lower() == 'exit':
                break
               
            if query.lower() == 'stats':
                print("\nContext Statistics:")
                stats = self.context_manager.get_stats()
                for key, value in stats.items():
                    print(f"{key}: {value}")
                self.context_manager.visualize_context()
                continue
               
            if query.lower() == 'clear':
                self.context_manager.chunks = []
                self.context_manager.current_token_count = 0
                conversation_history = []
                clear_output(wait=True)
                print("Context cleared!")
                continue
           
            response = self.process_query(query)
            conversation_history.append((query, response))
           
            print("\nResponse:")
            print(response)
            print("\n" + "-"*50)
           
            stats = self.context_manager.get_stats()
            print(f"Context usage: {stats['token_count']}/{stats['max_tokens']} tokens ({stats['usage_percentage']:.1f}%)")

The MCPColabDemo class ties the context manager to a seq2seq LLM, loading FLAN-T5 (or any specified Hugging Face model) on the chosen device, and provides utility methods for chunking and ingesting entire documents, processing user queries by prepending only the most relevant context, and running an interactive Colab session complete with real-time stats, visualizations, and commands for clearing or inspecting the evolving context window.

def run_mcp_demo():
    """Run a simple demo of the Model Context Protocol."""
    print("Running Model Context Protocol Demo...")
   
    context_manager = ModelContextManager(max_context_length=4096)
   
    print("Adding sample chunks...")
   
    context_manager.add_chunk(
        "The Model Context Protocol (MCP) is a framework for managing context "
        "windows in large language models. It helps optimize token usage and improve relevance.",
        importance=1.0
    )
   
    context_manager.add_chunk(
        "Context management involves techniques like sliding windows, chunking, "
        "and relevance filtering to handle large documents efficiently.",
        importance=0.8
    )
   
    for i in range(10):
        context_manager.add_chunk(
            f"This is test chunk {i} with some filler content to simulate a larger context "
            f"window that needs optimization. This helps demonstrate the MCP functionality "
            f"for context window management in language models on Google Colab.",
            importance=0.5 - (i * 0.02)  
        )
   
    stats = context_manager.get_stats()
    print("\nInitial Statistics:")
    for key, value in stats.items():
        print(f"{key}: {value}")
       
    query = "How does the Model Context Protocol work?"
    print(f"\nRetrieving context for: '{query}'")
    context = context_manager.retrieve_context(query)
    print(f"\nRelevant context:\n{context}")
   
    print("\nVisualizing context:")
    context_manager.visualize_context()
   
    print("\nDemo complete!")

The run_mcp_demo function ties everything together in a single script: it instantiates the ModelContextManager, adds a series of sample chunks with varying importance, prints out initial statistics, retrieves and displays the most relevant context for a test query, and finally visualizes the context window, providing a complete, end-to-end demonstration of the Model Context Protocol in action.

if __name__ == "__main__":
    run_mcp_demo()

Finally, this standard Python entry-point guard ensures that the run_mcp_demo() function executes only when the script is run directly (rather than imported as a module), triggering the end-to-end demonstration of the Model Context Protocol workflow.

In conclusion, we will have a fully functional MCP system that not only curbs runaway token usage but also prioritizes context fragments that truly matter for your queries. The ModelContextManager equips you with tools to balance semantic relevance, temporal freshness, and user-assigned importance. At the same time, the accompanying MCPColabDemo class provides an accessible framework for real-time experimentation and visualization. Armed with these patterns, you can extend the core principles by adjusting relevance thresholds, experimenting with various embedding models, or integrating with alternative LLM backends to tailor your domain-specific workflows. Ultimately, this approach enables you to create concise yet highly relevant prompts, resulting in more accurate and efficient responses from your language models.


Here is the Colab Notebook. Also, don’t forget to follow us on Twitter and join our Telegram Channel and LinkedIn Group. Don’t Forget to join our 90k+ ML SubReddit.

                        </div>
                                            <div class= Read More