Showing posts with label large-language-models. Show all posts
Showing posts with label large-language-models. Show all posts

Sunday, July 19, 2026

Book Review: Domain Specific Small Language Models

Artificial Intelligence (AI) powered applications are changing the way we consume and use information. Retrieval Augmented Generation (RAG) systems are becoming ubiquitous with AI generated answers to our questions replacing Google / Bing searches returning a list of links to matching documents. Programmers and non-programmers alike are using Agents backed by Large Language Models (LLM) to write code, create schedules, and a host of other things that we had to do by hand just a few years ago.

Although AI is now everywhere, the intelligence powering them come from large general-purpose LLMs that are run by only a few large companies like OpenAI and Anthropic. As AI gets more integrated into our lives, it will become harder for us to do without it. There is a risk of price gouging (aka a more realistic pricing model) as that happens, as Edward Tufte has noted.

Even with current pricing, LLM pipelines work out to be significantly more expensive and slower than their (albeit less powerful) pre-AI incarnations. A good place to look to lower costs and latency may be to consider whether all the functions that are using LLMs in a pipeline really need the power of these models, or if they can be achieved with (smaller and potentially faster) Domain Specific Small Language Models (SLM). I wanted to educate myself about this, so I picked up Domain Specific Small Language Models to learn more about SLMs. Here is my review.

The book is organized in three parts. Part 1 is basic background material, Part 2 covers the information to get you started with fine-tuning and deploying domain-specific SLMs in your pipelines, Part 3 covers some case studies, and Part 4 covers advanced topics mostly covering late-breaking breakthroughs and tooling popular in real-world pipelines.

Part 1 covers a brief history of LLMs, the Transformer architecture that they use, the contributions of the open source community which led to the development of SLMs, the case for SLMs, etc. If you are looking to learn about SLMs, it is very likely that you have a Machine Learning (ML) background, so you may be familiar with lot of this material already. However, it can be a useful refresher that can round out possible gaps in your knowledge.

Part 2 covers fine-tuning SLMs to adapt them to specific domains. The focus is on fine-tuning generative models (either encoder-decoder or decoder-only Transformer models), which tend to be larger and therefore require specialized training strategies such as LoRA (Low Rank Adaptation). It also discusses strategies to deploy SLMs at scale for inference. Model quantization is a big part of the deployment strategy for inference, and that gets its own chapter. In addition, there is a chapter on ONNX (Open Neural Network eXchange), which is an open format for models such that they can portably run on any ONNX supported runtime (spanning programming language environments).

Part 3 covers some real-world use cases of SLMs, such as Code Generation using fine-tuned SLMs such as StarCoder and CodeGen, and generating chemical and protein structures using ProtGPT2, AntibodyGPT, and CrystaLLM.

Part 4 covers additional advanced quantization techniques (FlexGen, SmoothQuant, BitNet), profiling using the ONNX Runtime (ORT) platform, and deployment platforms such as vLLM and Ollama. It also covers the integration of SLMs with RAG and Agentic pipelines, including a discussion of components to support Agentic pipelines such as GraphRAG, Tools and Memory. Finally, it covers test time compute, which is typically used for inference by reasoning engines, as well as adapting SLMs for reasoning using GRPO (Group Relative Policy Optimization).

There are code examples throughout the book, mostly based on the HuggingFace API. This makes sense, since much of SLM development happens using their API and is hosted on their repositories.

Overall, I think the book can jumpstart your journey into SLMs, but the field is moving too fast for one book to cover. What you will get out of the book are solid foundations that will help you navigate more recent developments in the field and decide if it is something that makes sense to explore for your use case.

Sunday, June 15, 2025

Book Review: Essential Graph RAG

Coming from a background of Knowledge Graph (KG) backed Medical Search, I don't need to be convinced about the importance of manually curated structured knowledge on the quality of search results. Traditional search is being rapidly replaced with Generative AI using a technique called Retrieval Augmented Generation (RAG), where the pipeline produces an answer summarizing the search results retrieved instead of the ten blue links that the searcher had to parse and retrieve an answer from earlier. In any case, I had been experimenting with Using KGs to enhance RAG to support this intuition, and when Microsoft announced their work on GtaphRAG, it felt good to be vindicated. So when Manning reached out to me to ask if I would be interested in reviewing the book Essential GraphRAG by Tomaž Bratanič and Oskar Hane, I jumped at the chance.

Both authors are from Neo4j, so it is not surprising that the search component is also Neo4j, even for vector search, and hybrid search is really vector + graph search (rather than the more common vector + lexical search). However, most people nowadays would prefer a multi-backend search that would include graph search as well as vector and lexical search, so the examples can help you learn (a) how to use Neo4j for vector search and (b) how to implement graph search with Neo4j. Since Neo4j is a leading graph database provider, this is useful information to know if you decide to incorporate graph search into your repertoire of tools, as you very likely are if you are reading this book.

The book is available under the Manning Early Access Program (MEAP) and is expected to be published in August 2025. It is currently organized into 8 chapters as follows:

Improving LLM accuracy -- here the authors introduce what LLMs are, what they are capable of as well as their limitations when used for question answering, i.e. not knowing about recent events post its training date, its tendency to hallucinate when it cannot answwe a question from the knowledge it was trained on, and its inability to know of company confidential or otherwise private information, since it is trained on public data only. They cover solutions to mitigate this, i.e. finetuning and RAG, and why RAG is a better alternaive in most cass. Finally they cover why KGs are the best general purpose datastore for RAG pipelines.

Vector Similarity Search and Hybrid Search -- here the authors cover the fundamentals of vector search, such as vector similarity functions, embedding models used to support vector search, and the reasoning behind chunking. They describe what a typical RAG pipeline looks like, although as mentioned earlier, they showcase Neo4j's vector search capabilities instead of relying on more popular vecror search alternatives. I thought it was good information though, since I wasn't aware that Neo4j supported vector search. They also cover hybrid search, in this case vector + graph search (this is a book about GraphRAG after all). Although I can definitely see Graph Search as one of the components of a hybrid search pipeline.

Advanced Vector Retrieval Strategies -- in this chapter, the authors introduce some interesting techniques to make your Graph Search produce more relevant context for your GraphRAG pipeline. Techniques on the query side include Step Back Prompting (SBP) to look for more generic concepts then drill down using Graph Search to improve recall, and the Parent Document Retriever pattern of retrieving parent documents of the chunks that matched, rather than the chunks themselves. On the indexing side, they talk about creating additional synthetic chunks that summarize actual chunks and can be queried as well as the chunks, and representing document chunks as pre-generated questions the chunk can answer instead of its text content.

Text2Cypher -- in this chapter, the authors show how an LLM can be prompted using Few Shot Learning (FSL) to generate Cypher queries from natural language. Users would type in a query using natural language, knowing nothing about the schema structure of the underlying Graph Database. The LLM, through detailed prompts and examples, would translate the natural language query to Cypher query. The authors also reference pre-trained models from Neo4j that have been fine-tuned to do this. While these models are generally not as effective as the one built from LLMs through prompting, they are more efficient on large volumes of data.

Agentic RAG -- Agentic RAG allows autonomous / semi-autonomous LLM backed software components, called Agents, to modify and enhance the standard control flow for RAG. One change could be for an Agent (the Router) to determine query intent and call on one or more retrieveers from the available pool of retrievers, or another (the Critic) to determine if the answer generated so far is adequate given the user's query, and if not, to rerun the pipeline with a modified query until the query is fully answered. The authors go on to describe a system (with code) consisting of a Router and Critic and several Retrieval Agents.

Constructing Knowledge Graph with LLM -- this chapter focuses on the index creation. Search is traditionally done on unstructured data such as text documents. This chapter describes using the LLM to extract entities of known types (PERSON, ORGANIZATION, LOCATION, etc), followed by a manual / semi-manual Graph Modeling step to set up relations between these extracted entities and build a schema. It then talks a little about convert specific query types into structured Cypher queries that leverage this schema.

Microsoft GraphRAG Implementation -- this chapter deals specifically with Microsoft's GraphRAG implementation. While most people think of GraphRAG as any infrastructure that supports incorporating Graph Search into a RAG pipeline, Microsoft specifies it as a multi-step recipe to build your KG from your data sources and use results from your KG to support a RAG pipeline. The steps involved are structured extraction and community detection, followed by summarization of community chunks into synthetic nodes. To some extent this is similar to Chonkie's Semantic Double Pass Merging (SDPM) chunker, except that the size of the skip window is unbounded. These synthetic chunks can be useful to answer global questions that span multiple ideas across the corpus. However, as the authors show, this approach can be effective for local queries as well.

RAG Application Evaluation -- because of the stochastic nature of LLMs, evaluating RAG pipelines in general present some unique challenges. Here these challenges are investigated with particular reference to GraphRAG systems, i.e. where the retrieval context is provided by Knowledge Graphs. The authors describe some metrics fro the RAGAS library, where LLMs are used to generate these metrics from outputs at different stages of the RAG pipeline. It also discusses ideas for setting up an evaluation dataset. The metrics covered in the example sare RAGAS context recall, faithfulness and answwr correctness.

Overall, the book takes a very practical, hands-on approach to the subject. It is filled with code examples and practical advice for leveraging KGs in RAG, and using Large Language Models (LLM) to build KGs, as well as evaluating such pipelines. If you were thinking of incorporating Graph Search into your search pipeline, be it traditional, hybrid, RAG or agentic, you will find the information in the book useful and beneficial.

Monday, July 29, 2024

Experiments with Prompt Compression

I recently came across Prompt Compression (in the context of Prompt Engineering on Large Language Models) on this short course on Prompt Compression and Query Optimization from DeepLearning.AI. Essentially it involves compressing the prompt text using a trained model to drop non-essential tokens. The resulting prompt is shorter (and in cases of the original context being longer than the LLM's context limit, not truncated) but retains the original semantic meaning. Because it is short, the LLM can process it faster and cheaper, and in some cases get around the Lost In the Middle problems observed with long contexts.

The course demonstrated Prompt Compression using the LLMLingua library (paper) from Microsoft. I had heard about LLMLingua previously from my ex-colleague Raahul Dutta, who blogged about it on his Edition 26: LLMLingua - A Zip Technique for Prompt post, but at the time I thought maybe it was more in the realm of research. Seeing it mentioned in the DeepLearning.AI course made it feel more mainstream, so I tried it out a single query from my domain using their Quick Start example, compressing the prompt with the small llmlingua-2-bert-base-multilingual-cased-meetingbank model, and using Anthropic's Claude-v2 on AWS Bedrock as the LLM.

Compressing the prompt for the single query gave me a better answer than without compression, at least going by inspecting the answer produced by the LLM before and after compression. Encouraged by these results, I decided to evaluate the technique using a set of around 50 queries I had lying around (along with a vector search index) from a previous project. This post describes the evaluation process and the results I obtained from it.

My baseline was a naive RAG pipeline, with the context retrieved by vector matching the query against the corpus, and then incorporated into a prompt that looks like this. The index is an OpenSearch index containing vectors of document chunks, vectorization was done using the all-MiniLM-L6-v2 pre-trained SentenceTransformers encoder, and the LLM is Claude-2 (on AWS Bedrock as mentioned previously).

1
2
3
4
5
6
7
8
9
Human: You are a medical expert tasked with answering questions
expressed as short phrases. Given the following CONTEXT, answer the QUESTION.

CONTEXT:
{context}

QUESTION: {question}

Assistant:

While the structure of the prompt is pretty standard, LLMLingua explicitly requires the prompt to be composed of an instruction (the System prompt beginning with Human:), the demonstration (the {context}) and the question (the actual quary to the RAG pipeline). The LLMLingua Compressor's compress function expects these to be passed separately as parameters. Presumably, it compresses the demonstration with respect to the instruction and the question, i.e. context tokens that are non-essential given the instruction and question are dropped during the compression process.

The baseline for the experiment uses the context as retrieved from the vector store without compression, and we evaluate the effects of prompt compression using the two models listed in LLMLingua's Quick Start -- llmlingua-2-bert-base-multilingual-cased-meetingbank (small model) and llmlingua-2-bert-base-multilingual-cased-meetingbank (large model). The three pipelines -- baseline, compression using small model, and compression using large model -- are run against my 50 query dataset. The examples imply that the compressed prompt can be provided as-is to the LLM, but I found that (at least with the small model), the resulting compressed prompt generates answers that does not always capture all of the question's nuance. So I ended up substituting only the {context} part of the prompt with the generated compressed prompt in my experiments.

Our evaluation metric is Answer Relevance as defined by the RAGAS project. It is a measure of how relevant the generated answer is given the question. To calculate this, we prompt the LLM to generate a number of (in our case, upto 10) questions from the generated answer. We then compute the cosine similarity of the vector of each generated question with the vector of the actual question. The average of these cosine similarities is the Answer Relevance. Question Generation from the answer is done by prompting Claude-2 and vectorization of the original and generated questions are done using the same SentenceTransformer encoder we used for retrieval.

Contrary to what I saw in my first example, the results were mixed when run against the 50 queries. Prompt Compression does result in faster response times, but it degraded the Answer Relevance scores more times than improve it. This is true for both the small and large compression models. Here are plots of the difference of the Answer Relevance score for the compressed prompt against the baseline uncompressed prompt for each compression model. The vertical red line separates the cases where compression is hurting answer relevance (left side) versus improving answer relevance (right side). In general, it seems like compression helps when the input prompt is longer, which intuitively makes sense. But there doesn't seem to be a simple way to know up front if prompt compression is going to help or hurt.

I used the following parameters to instantiate LLMLingua's PromptCompressor object and to call its compress_prompt function. These are the same parameters that were shown in the Quick Start. It is possible I may have gotten different / better results if I had experimented a bit with the parameters.

1
2
3
4
5
6
7
8
9
from llmlingua import PromptCompressor

compressor = PromptCompressor(model_name=model_name, use_llmlingua2=True)

compressed = compressor.compress_prompt(contexts, instruction=instruction, question=query,
    target_token=500, condition_compare=True, condition_in_question="after", 
    rank_method="longllmlingua", use_sentence_level_filter=False, context_budget="+100",
    dynamic_context_compression_ratio=0.4, reorder_context="sort")
compressed_context = compressed["compressed_prompt"]

A few observations about the compressed context. The number of context documents changes before and after compression. In my case, all input contexts had 10 chunks, and the output would vary between 3-5 chunks, which probably leads to the elimination of Lost in the Middle side-effects as claimed in LLMLingua's documentation. Also, the resulting context chunks are shorter and seems to be a string of keywords rather than coherent sentences, basically unintelligible to human readers, but intelligible to the LLM.

Overall, Prompt Compression seems like an interesting and very powerful technique which can result in savings in time and money if used judiciously. Their paper shows very impressive results on some standard benchmark datasets with supervised learning style metrics using a variety of compression ratios. I used Answer Relevance because it can be computed without needing domain experts to grade additional answers. But it is likely that I am missing some important optimization, so I am curious if any of you have tried it, and if your results are different from mine. If so, would appreciate any pointers to things you think I might be missing.

Sunday, June 30, 2024

Table Extraction from PDFs using Multimodal (Vision) LLMs

Couple of weeks ago a colleague and I participated in an internal hackathon where the task was to come up with an interesting use case using the recent multi-modal Large Language Models (LLMs). Multi-modal LLMs take not only text inputs via their prompt like earlier LLMs, but can also accept non-text modalities such as images and audio. Some examples of multi-modal LLMs are GPT-4o from OpenAI, Gemini 1.5 from Google, and Claude-3.5 Sonnet from Anthropic. The hackathon provided access to GPT-4o through Azure, Microsoft's Cloud Computing Platform. We did not win, there were other entries that were better than ours both in terms of the originality of their idea as well as quality of their implementations. However, we learned some cool new things during the hackathon, and figured that these might be of general interest to others as well, hence this post.

Our idea was to use GPT-4o to extract and codify tables found in academic papers as semi-structured data (i.e. JSON). We could then either query the JSON data for searching within tables, or convert it to Markdown for downstream LLMs to query them easily via their text interface. We had originally intended to extend the idea to figures and charts, but we could not get that pipeline working end to end.

Here is what our pipeline looked like.

  1. Academic papers are usually available as PDFs. We use the PyMuPDF library to split the PDF file into a set of image files, where each image file corresponds to a page in the paper.
  2. We then send each page image through the Table Transformer, which returns bounding box information for each table it detects in the page, as well as a confidence score. The Table Transformer model we used was microsoft/table-transformer-detection.
  3. We crop out each table from the pages using the bounding box information, and then send each table to GPT-4o as part of a prompt asking to convert it to a JSON structure. GPT-4o responds with a JSON structure representing the table.

This pipeline was based on my colleague's idea. I like how it progressively simplifies the task by splitting each page of the incoming PDF into its own image, then uses a pre-trained Table Transformer to crop out the tables from them, and only then passes the table to GPT-4o to convert to JSON. That table image is passed into the prompt as a "data URL" which is just the base-64 encoding of the image formatted as "data:{mime_type};base64,{base64_encoded_data}. The Table Transformer, while not perfect, proved remarkably successful at identifying tables in the text. I say remarkable because we used a pre-trained model, but perhaps it is not that remarkable once you consider that it probably trained on tables in academic papers as well.

Our prompt for GPT-4o looked something like this:

System: You are an AI model that specializes in detecting the tables and extracting, interpreting table content from images. Follow below instruction step by step:
1. Recognize whether the given image is table or not, if it's not a table print "None". if it's a table go to next step.
2. accurately convert the table's content into a structured structured Json format

general instruction:
1. do not output anything extra. 
2. a table must contains rows and columns

User: Given the image, detect whether it's a table or not, if it's a table then convert it to Json format
{image_data_url}

For the figure pipeline, I tried to use an OWL-VIT (Vision Transformer for Open World Localization) model in place of the Table Transformer. But it was not as successful at detecting figures in the text, probably since SAM seems to be fine-tuned to detect objects in natural images. Unfortunately we couldn't find a pre-trained model that would work for this particular case. Another issue was converting the fgure into a semi-structured JSON representation, we ended up asking GPT-4o to describe the image as text.

One suggestion by some of my TWIML non-work colleagues was to ask GPT-4o to return the bounding boxes for the images it finds in it, and then use that to extract the figures to send to GPT-4o for describing. It didn't work unfortunately, but was definitely worth trying. As LLMs get more and more capable, I think it makes sense to rethink our pipelines to delegate more and more work to the LLM. Or at least verify that it can't do something before moving on to older (and harder to implement) solutions.

Saturday, May 18, 2024

Finetuning RAGAS Metrics using DSPy

Last month, I decided to sign-up for the Google AI Hackathon, where Google provided access to their Gemini Large Language Model (LLM) and tasked participants with building a creative application on top of it. I have worked with Anthropic's Claude and OpenAI's GPT-3 at work previously, and I was curious to see how Gemini stacked up against them. I was joined in that effort by David Campbell and Mayank Bhaskar, my non-work colleagues from the TWIML (This Week In Machine Learning) Slack. Winners for the Google AI Hackathon were declared last Thursday, and whilte our project sadly did not win anything, the gallery provides examples of some very cool applications of LLMs (and Gemini in particular) for both business and personal tasks.

Our project was to automate the evaluation of RAG (Retrieval Augmented Generation) pipelines using LLMs. I have written previously about the potential of LLMs to evaluate search pipelines, but the scope of this effort is broader in that it attempts to evaluate all aspects of the RAG pipeline, not just search. We were inspired by the RAGAS project, which defines 8 metrics that cover various aspects of the RAG pipeline. Another inspiration for our project was the ARES paper, which shows that fine-tuning the LLM judges on synthetically generated outputs improves evaluation confidence.

Here is a short (3 minutes) video description of our project on Youtube. This was part of our submission for the hackathon. We provide some more information about our project in our blog post below.

We re-implemented the RAGAS metrics using LangChain Expression Language (LCEL) and applied them to (question, answer, context and ground truth) tuples from the AmnestyQA dataset to generate the scores for these metrics. My original reason for doing this, rather than using the using what RAGAS provided directly, was because I couldn't make them work properly with Claude. This was because Claude cannot read and write JSON as well as GPT-3 (it works better with XML), and RAGAS was developed using GPT-3. All the RAGAS metrics are prompt-based and transferrable across LLMs with minimal change, and the code is quite well written. I wasn't sure if I would encounter similar issues with Gemini, so it seemed easier to just re-implement the metrics from the ground up for Gemini using LCEL than try to figure out how to make RAGAS work with Gemini. However, as we will see shortly, it ended up being a good decision.

Next we re-implemented the metrics with DSPy. DSPy is a framework for optimizing LLM prompts. Unlike RAGAS, where we tell the LLM how to compute the metrics, with DSPy the general approach is to have very generic prompts and show the LLM what to do using few shot examples. The distinction is reminiscent of doing prediction using Rules Engines versus using Machine Learning. Extending the analogy a bit further, DSPy provides its BootstrapFewShotWithRandomSearch optimizer that allows you to search through its "hyperparameter space" of few shot examples, to find the best subset of examples to optimize the prompt with, with respect to some score metric you are optimizing for. In our case, we built the score metric to minimize the difference between the the score reported by the LCEL version of the metric and the score reporteed by the DSPy version. The result of this procedure are a set of prompts to generate the 8 RAG evaluation metrics that are optimized for the given domain.

To validate this claim, we generated histograms of scores for each metric using the LCEL and DSPy prompts, and compared how bimodal, or how tightly clustered around 0 and 1, they were. The intuition is that the more confident the LLM is about the evaluation, the more it will tend to deliver a confident judgment clustered around 0 or 1. In practice, we do see this happening in case of the DSPy prompts for all but 2 of the metrics, although the differences are not very large. This may be because we the AmnestyQA dataset is very small, only 20 questions.

To address the size of AmnestyQA dataset, Dave used the LLM to generate some more (question, context, answer, ground_truth) tuples given a question and answer pair from AmnestyQA and a Wikipedia retriever endpoint. The plan was for us to use this larger dataset for optimizing the DSPy prompts. However, rather than doing this completely unsupervised, we wanted to have a way for humans to validate and score the LCEL scores from these additional questions. We would then use these validated scores as the basis for optimizing the DSPy prompts for computing the various metrics.

This would require a web based tool that would allow humans to examine the output of each step of the LCEL metric score process. For example, the Faithfulness metric has two steps, the first is to extract facts from the answer, and the second is to provide a binary judgment of whether the context contains the fact. The score is computed by adding up the individual binary scores. The tool would allow us to view and update what facts were extracted in the first stage, and the binary output for each of the fact-context pairs. This is where implementing the RAGAS metrics on our own helped us, we refactored the code so the intermediate results were also available to the caller. Once the tool was in place, we would use it to validate our generated tuples and attempt to re-optimise the DSPy prompts. Mayank and Dave had started on this , but unfortunately we ran out of time before we could complete this step.

Another thing we noticed is that calculation of most of the metrics involves one or more subtasks to make some kind of binary (true / false) decision about a pair of strings. This is something that a smaller model, such as a T5 or a Sentence Transformer, could do quite easily, more predictably, faster, and at lower cost. As before, we could use extract the intermediate outputs from the LCEL metrics to create training data to do this. We could use DSPy and its BootstrapFindTune optimizer to fine-tune these smaller models, or fine-tune Sentence Transformers or BERT models for binary classification and hook them up into the evaluation pipeline.

Anyway, that was our project. Obviously, there is quite a bit of work remaining to make it into a viable product for LLM based evaluation using the strategy we laid out. But we believe we have demonstrated that this can be viable, that given sufficient training data (about 50-100 examples for the optimized prompt, and maybe 300-500 each for the binary classifiers), it should be possible to build metrics that are tailored to one's domain and that can deliver evaluation judgments with greater confidence than those built using simple prompt engineering. In case you are interested in exploring further, you can find our code and preliminary results at sujitpal/llm-rag-eval on GitHub.

Sunday, March 17, 2024

Hierarchical (and other) Indexes using LlamaIndex for RAG Content Enrichment

At our weekly This Week in Machine Learning (TWIML) meetings, (our leader and facilitataor) Darin Plutchok pointed out a LinkedIn blog post on Semantic Chunking that has been recently implemented in the LangChain framework. Unlike more traditional chunking approaches that use number of tokens or separator tokens as a guide, this one chunks groups of sentences into semantic units by breaking them when the (semantic) similarity between consecutive sentences (or sentence-grams) fall below some predefined threshold. I had tried it earlier (pre-LangChain) and while results were reasonable, it would need a lot of processing, so I went back to what I was using before.

I was also recently exploring LlamaIndex as part of the effort to familiarize myself with the GenAI ecosystem. LlamaIndex supports hierarchical indexes natively, meaning it provides the data structures that make building them easier and more natural. Unlike the typical RAG index, which are just a sequence of chunks (and their vectors), hierarchical indexes would cluster chunks into parent chunks, and parent chunks into grandparent chunks, and so on. A parent chunk would generally inherit or merge most of the metadata from its children, and its text would be a summary of its children's text contents. To illustrate my point about LlamaIndex data structures having natural support for this kind of setup, here are the definitions of the LlamaIndex TextNode (the LlamaIndex Document object is just a child of TextNode with an additional doc_id: str field) and the LangChain Document. Of particular interest is the relationships field, which allows pointers to other chunks using named relationships PARENT, CHILD, NEXT, PREVIOUS, SOURCE, etc. Arguably, the LlamaIndex TextNode can be represented more generally and succintly by the LangChain Document, but the hooks do help to support hierarchical indexing more naturally.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
# this is a LlamaIndex TextNode
class TextNode:
  id_: str = None
  embedding: Optional[List[float]] = None
  extra_info: Dict[str, Any]
  excluded_embed_metadata_keys: List[str] = None
  excluded_llm_metadata_keys: List[str] = None
  relationships: Dict[NodeRelationship, Union[RelatedNodeInfo, List[RelatedNodeInfo]] = None
  text: str
  start_char_idx: Optional[int] = None
  end_char_idx: Optional[int] = None
  text_template: str = "{metadata_str}\n\n{content}"
  metadata_template: str = "{key}: {value}",
  metadata_separator = str = "\n"

# and this is a LangChain Document
class Document:
  page_content: str
  metadata: Dict[str, Any]

In any case, having discovered the hammer that is LlamaIndex, I began to see a lot of potential hierarchical indexes nails. One such nail that occurred to me was to use Semantic Chunking to cluster consecutive chunks rather than sentences (or sentence-grams), and then create parents nodes from these chunk clusters. Instead of computing cosine similarity between consecutive sentence vectors to build up chunks, we compute cosine similarity across consecutive chunk vectors and split them up into clusters based on some similarity threshold, i.e. if the similarity drops below the threshold, we terminate the cluster and start a new one.

Both LangChain and LlamaIndex have implementations of Semantic Chunking (for sentence clustering into chunks, not chunk clustering into parent chunks). LangChain's Semantic Chunking allows you to set the threshold using percentiles, standard deviation and inter-quartile range, while the LlamaIndex implementation supports only the percentile threshold. But intuitively, here's how you could get an idea of the percentile threshold to use -- thresholds for the other methods can be computed similarly. Assume your content has N chunks and K clusters (based on your understanding of the data or from other estimates), then assuming a uniform distribution, there would be N/K chunks in each cluster. If N/K is approximately 20%, then your percentile threshold would be approximately 80.

LlamaIndex provides an IngestionPipeline which takes a list of TransformComponent objects. My pipeline looks something like below. The last component is a custom subclass of TransformComponent, all you need to do is to override it's __call__ method, which takes a List[TextNode] and returns a List[TextNode].

1
2
3
4
5
6
7
8
transformations = [
    text_splitter: SentenceSplitter,
    embedding_generator: HuggingFaceEmbedding,
    summary_node_builder: SemanticChunkingSummaryNodeBuilder
]
ingestion_pipeline = IngestionPipeline(transformations=transformations)
docs = SimpleDirectoryReader("/path/to/input/docs")
nodes = ingestion_pipeline.run(documents=docs)

My custom component takes the desired cluster size K during construction. It uses the vectors computed by the (LlamaIndex provided) HuggingFaceEmbedding component to compute similarities between consecutive vectors and uses K to compute a threshold to use. It then uses the threshold to cluster the chunks, resulting in a list of list of chunks List[List[TextNode]]. For each cluster, we create a summary TextNode and set its CHILD relationships to the cluster nodes, and the PARENT relationship of each child in the cluster to this new summary node. The text of the child nodes are first condensed using extractive summarization, then these condensed summaries are further summarized into one final summary using abstractive summarization. I used bert-extractive-summarizer with bert-base-uncased for the first and a HuggingFace summarization pipeline with facebook/bert-large-cnn for the second. I suppose I could have used an LLM for the second step, but it would have taken more time to build the index, and I have been experimenting with ideas described in the DeepLearning.AI course Open Source Models with HuggingFace.

Finally, I recalculate the embeddings for the summary nodes -- I ran the summary node texts through the HuggingFaceEmbedding, but I guess I could have done some aggregation (mean-pool / max-pool) on the child vectors as well.

Darin also pointed out another instance of Hierarchical Index proposed via the RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval and described in detail by the authors in this LlamaIndex webinar. This is a bit more radical than my idea of using semantic chunking to cluster consecutive chunks, in that it allows clustering of chunks across the entire corpus. One other significant difference is that it allows for soft-clustering, meaning a chunk can be a member of more than one chunk. They first reduce the dimensionality of the vector space using UMAP (Uniform Manifold Approximation and Projection) and then apply Gaussian Mixture Model (GMM) to do the soft clustering. To find the optimum number of clusters K for the GMM, one can use a combination of AIC (Aikake Information Criterion) and BIC (Bayesian Information Criterion).

In my case, when training the GMM, the AIC kept decreasing as the number of clusters increased, and the BIC had its minimum value for K=10, which corresponds roughly to the 12 chapters in my Snowflake book (my test corpus). But there was a lot of overlap, which would force me to implement some sort of logic to take advantage of the soft clustering, which I didn't want to do, since I wanted to reuse code from my earlier Semantic Chunking node builder component. Ultimately, I settled on 90 clusters by using my original intuition to compute K, and the resulting clusters seem reasonably well separated as seen below.

Using the results of the clustering, I built this also as another custom LlamaIndex TransformComponent for hierarchical indexing. This implementation differs from the previous one only in the way it assigns nodes to clusters, all other details with respect to text summarization and metadata merging are identical.

For both these indexes, we have a choice to maintain the index as hierarchical, and decide which layer(s) to query based on the question, or add the summary nodes into the same level as the other chunks, and let vector similarity surface them when queries deal with cross-cutting issues that may be found together in these nodes. The RAPTOR paper reports that they don't see a significant gain using the first approach over the second. Because my query functionality is LangChain based, my approach has been to generate the nodes and then reformat them into LangChain Document objects and use LCEL to query the index and generate answers, so I haven't looked into querying from a hierarchical index at all.

Looking back on this work, I am reminded of similar choices when designing traditional search pipelines. Often there is a choice between building functionality into the index to support a cheaper query implementation, or building the logic into the query pipeline that may be more expensive but also more flexible. I think LlamaIndex started with the first approach (as evidenced by their blog posts Chunking Strategies for Large Language Models Part I and Evaluating Ideal Chunk Sizes for RAG Systems using LlamaIndex) while LangChain started with the second, even though nowadays there is a lot of convergence between the two frameworks.

Saturday, February 24, 2024

Thoughts on using LangChain LCEL with Claude

I got into Natural Language Processing (NLP) and Machine Learning (ML) through Search. And this led me into Generative AI (GenAI), which led me back to Search via Retrieval Augmented Generation (RAG). RAG started out relatively simple -- take a query, generate search results, use search results as context for a Large Language Model (LLM) to generate an abstractive summary of the results. Back when I started on my first "official" GenAI project middle of last year, there were not too many frameworks to support building GenAI components (at least not the prompt based ones), except maybe LangChain, which was just starting out. But prompting as a concept is not too difficult to understand and implement, so thats what we did at the time.

I did have plans to use LangChain in my project once it became more stable, so I started out building my components to be "langchain compliant". But that turned out to be a bad idea as LangChain continued its exponential (and from the outside at least, somewhat haphazard) growth and showed no signs of stabilizing. At one point, LangChain users were advised to make pip install -U langchain part of their daily morning routine! So anyway, we ended up building up our GenAI application by hooking up third party components with our own (non-framework) code, using Anthropic's Claude-v2 as our LLM, ElasticSearch as our lexical / vector document store and PostgreSQL as our conversational buffer.

While I continue to believe that the decision to go with our own code made more sense than trying to jump on the LangChain (or Semantic Kernel, or Haystack, or some other) train, I do regret it in some ways. A collateral benefit for people who adopted and stuck with LangChain were the ready-to-use implementations of cutting-edge RAG and GenAI techniques that the community implemented at almost the same pace as they were being proposed in academic papers. For the subset of these people that were even slightly curious about how these implementations worked, this offered a ringside view into the latest advances in the field and a chance to stay current with it, with minimal effort.

So anyway, in an attempt to replicate this benefit for myself (going forward at least), I decided to learn LangChain by doing a small side project. Earlier I needed to learn to use Snowflake for something else and had their free O'Reilly book on disk, so I converted it to text, chunked it, and put it into a Chroma vector store. I then tried to implement examples from the DeepLearning.AI courses LangChain: Chat with your Data and LangChain for LLM Application Development. The big difference is that the course examples use OpenAI's GPT-3 as their LLM whereas I use Claude-2 on AWS Bedrock in mine. In this post, I share the issues I faced and my solutions, hopefully this can help guide others in similar situations.

Couple of observations here. First, the granularity of GenAI components is necessarily larger than traditional software components, and this means application details that the developer of the component was working on can leak into the component itself (mostly through the prompt). To a user of the component, this can manifest as subtle bugs. Fortunately, LangChain developers seem to have also noticed this and have come up with the LangChain Expression Language (LCEL), a small set of reusable components that can be composed to create chains from the ground up. They have also marked a large number of Chains as Legacy Chains (to be converted to LCEL chains in the future).

Second, most of the components (or chains, since that is LangChain's central abstraction) are developed against OpenAI GPT-3 (or its chat version GPT-3.5 Turbo) whose strengths and weaknesses may be different from those of your LLM. For example, OpenAI is very good at generating JSON output, whereas Claude is better at generating XML. I have also seen that Claude can terminate XML / JSON output mid-output unless forced to complete using stop_sequences. Yhis doesn't seem to be a problem GPT-3 users have observed -- when I mentioned this problem and the fix, I drew a blank on both counts.

To address the first issue, my general approach in trying to re-implement these examples has been to use LCEL to build my chains from scratch. I attempt to leverage the expertise available in LangChain by looking in the code or running the existing LangChain chain with langchain.debug set to True. Doing this helps me see the prompt being used and the flow, which I can use to adapt the prompt and flow for my LCEL chain. To address the second issue, I play to Claude's strengths by specifying XML output format in my prompts and parsing them as Pydantic objects for data transfer across chains.

The example application I will use to illustrate these techniques here is derived from the Evaluation lesson from the LangChain for LLM Application Development course, and is illustrated in the diagram below. The application takes a chunk of text as input, and uses the Question Generation chain to generate multiple question-answer pairs from it. The questions and the original content are fed into the Question Answering chain, which uses the question to generate additional context from a vector retriever, and uses all three to generate an answer. The answer generated from the Question Generation chain and the answer generated from the Question Answering chain are fed into a Question Generation Evaluation chain, where the LLM grades one against the other, and generates an aggregate score for the questions generated from the chunk.

Each chain in this pipeline is actually quite simple, they take one or more inputs and generates a block of XML. All the chains are structured as follows:

1
2
3
from langchain_core.output_parsers import StrOutputParser

chain = prompt | model | StrOutputParser()

And all our prompts follow the same general format. Here is the prompt for the Evaluation chain (the third one) which I adapted from the QAEvalChain used in the lesson notebook. Developing from scratch using LCEL gives me the chance to use Claude's Human / Assistant format (see LangChain Guidelines for Anthropic) rather than depend on the generic prompt that happens to work well for GPT-3.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
Human: You are a teacher grading a quiz.

You are given a question, the context the question is about, and the student's 
answer.

QUESTION: {question}
CONTEXT: {context}
STUDENT ANSWER: {predicted_answer}
TRUE ANSWER: {generated_answer}

You are to score the student's answer as either CORRECT or INCORRECT, based on the 
context.

Write out in a step by step manner your reasoning to be sure that your conclusion 
is correct. Avoid simply stating the correct answer at the outset.

Please provide your response in the following format:

<result>
    <qa_eval>
        <question>the question here</question>
        <student_answer>the student's answer here</student_answer>
        <true_answer>the true answer here</true_answer>
        <explanation>step by step reasoning here</explanation>
        <grade>CORRECT or INCORRECT here</grade>
    </qa_eval>
</result>

Grade the student answers based ONLY on their factual accuracy. Ignore differences in 
punctuation and phrasing between the student answer and true answer. It is OK if the 
student answer contains more information than the true answer, as long as it does not 
contain any conflicting statements.

Assistant:

In addition, I specify the formatting instructions explicitly in the prompt instead of using the canned ones from XMLOutputParser or PydanticOutputParser via get_formatting_instructions(), which are comparatively quite generic and sub-optimal. By convention, the outermost tag in my format is always <result>...</result>. The qa_eval tag inside result has a corresponding Pydantic class analog declared in the code as follows:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
from pydantic import BaseModel, Field

class QAEval(BaseModel):
    question: str = Field(alias="question", description="question text")
    student_answer: str = Field(alias="student_answer",
                                description="answer predicted by QA chain")
    true_answer: str = Field(alias="true_answer",
                             description="answer generated by QG chain")
    explanation: str = Field(alias="explanation",
                             description="chain of thought for grading")
    grade: str = Field(alias="grade",
                       description="LLM grade CORRECT or INCORRECT")

After the StrOutputParser extracts the LLM output into a string, it is first passed through a regular expression to remove any content outside the <result>...</result>, then convert it into the QAEval Pydantic object using the following code. This allows us to keep object manipulation between chains independent of the output format, as well as negate any need for format specific parsing.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import re
import xmltodict

from pydantic import Field
from pydantic.generics import GenericModel
from typing import Generic, List, Tuple, TypeVar

T = TypeVar("T")

class Result(GenericModel, Generic[T]):
    value: T = Field(alias="result")

def parse_response(response):
    response = response.strip()
    start_tag, end_tag = "<result>", "</result>"
    is_valid = response.startswith(start_tag) and response.endswith(end_tag)
    if not is_valid:
        pattern = f"(?:{start_tag})(.*)(?:{end_tag})"
        p = re.compile(pattern, re.DOTALL)
        m = p.search(response)
        if m is not None:
            response = start_tag + m.group(1) + end_tag
    resp_dict = xmltodict.parse(response)
    result = Result(**resp_dict)
    return result

# example call
response = chain.invoke(
    "question": "the question",
    "context": "the context",
    "predicted_answer": "the predicted answer",
    "generated_answer": "the generated answer"
})
result = parse_response(response)
qa_eval = result.value["qa_eval"]

One downside to this approach is that it uses the current version of the Pydantic toolkit (v2) whereas LangChain still uses Pydantic V1 internally, as descibed in LangChain's Pydantic compatibility page. This is why this conversion needs to be outside LangChain and in the application code. Ideally, I would like this to be part of a subclass of PydanticOutputParser where the formatting_instructions could be generated from the class definition as a nice side effect, but that would mean more work than I am prepared to do at this point :-). Meanwhile, this seems like a decent compromise.

Thats all I had for today. Thank you for staying with me so far, and hope you found this useful!