Blog: Add post on leveraging Katib for efficient RAG optimization.

Signed-off-by: Varsha Prasad Narsing <[email protected]>

Author: varshaprasad96

_posts/2025-02-21-katib-rag-optimization.md

@@ -0,0 +1,224 @@
+---
+toc: true
+layout: post
+categories: [katib]
+description: "Leveraging Katib for efficient RAG optimization."
+comments: true
+title: "Optimizing RAG Pipelines with Katib: Hyperparameter Tuning for Better Retrieval & Generation"
+author: "Varsha Prasad Narsing" (@varshaprasad96)
+---
+
+# Introduction
+
+As machine learning models become more sophisticated, optimising their performance remains a critical challenge. Kubeflow provides a robust component, [KATIB][Katib], designed for hyperparameter optimization and neural architecture search. As a part of the Kubeflow ecosystem, KATIB enables scalable, automated tuning of underlying machine learning models, reducing the manual effort required for parameter selection while improving model performance across diverse ML workflows.
+
+With Retrieval-Augmented Generation ([RAG][rag]) becoming an increasingly popular approach for improving search and retrieval quality, optimizing its parameters is essential to achieving high-quality results. RAG pipelines involve multiple hyperparameters that influence retrieval accuracy, hallucination reduction, and language generation quality. In this blog, we will explore how KATIB can be leveraged to fine-tune a RAG pipeline, ensuring optimal performance by systematically adjusting key hyperparameters.
+
+# Let's Get Started!
+
+## STEP 1: Setup
+
+Since compute resources are scarcer than a perfectly labeled dataset :), we’ll use a lightweight KinD cluster to run this example locally. Rest assured, this setup can seamlessly scale to larger clusters by increasing the dataset size and the number of hyperparameters to tune.
+
+To get started, we'll first install the KATIB controller in our cluster by following the steps outlined [here][katib_installation].
+
+## STEP 2: Implementing RAG pipeline
+
+In this implementation, we use a retriever model to fetch relevant documents based on a query and a generator model to produce coherent text responses.
+
+### Implementation Details:
+
+1. Retriever: Sentence Transformer & FAISS Index
+   * A SentenceTransformer model (paraphrase-MiniLM-L6-v2) encodes predefined documents into vector representations.
+   * FAISS is used to index these document embeddings and perform efficient similarity searches to retrieve the most relevant documents.
+2. Generator: Pre-trained GPT-2 Model
+   * A Hugging Face GPT-2 text generation pipeline (which can be replaced with any other model) is used to generate responses based on the retrieved documents. I chose GPT-2 for this example as it is lightweight enough to run on my local machine while still generating coherent responses.
+3. Query Processing & Response Generation
+   * When a query is submitted, the retriever encodes it and searches the FAISS index for the top-k most similar documents.
+   * These retrieved documents are concatenated to form a context, which is then passed to the GPT-2 model to generate a response.
+4. Evaluation: [BLEU][bleu] (Bilingual Evaluation Understudy) Score Calculation
+   * To assess the quality of generated responses, we use the BLEU score, a popular metric for evaluating text generation.
+   * The evaluate function takes a query, retrieves documents, generates a response, and compares it against a ground-truth reference to compute a BLEU score with smoothing functions from the nltk library.
+ 
+Below is the script implementing this RAG pipeline:
+
+```python
+from sentence_transformers import SentenceTransformer
+from transformers import pipeline
+import numpy as np
+import faiss
+
+# Retriever: Pre-trained SentenceTransformer
+retriever_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
+
+documents = fetch_documents()
+doc_embeddings = retriever_model.encode(documents)
+
+# Build FAISS index
+index = faiss.IndexFlatL2(doc_embeddings.shape[1])
+index.add(np.array(doc_embeddings))
+
+# Generator: Pre-trained GPT-2
+generator = pipeline("text-generation", model="gpt2", tokenizer="gpt2")
+
+def rag_pipeline_execute(query, top_k=1, temperature=1.0):
+    query_embedding = retriever_model.encode([query])
+    distances, indices = index.search(query_embedding, top_k)
+    retrieved_docs = [documents[i] for i in indices[0]]
+
+    context = " ".join(retrieved_docs)
+    generated = generator(context, max_length=50, temperature=temperature, num_return_sequences=1)
+    return generated[0]['generated_text']
+
+def fetch_documents():
+    return [
+        # Retun the list of documents...
+    ]
+```
+
+#### Evaluating the Pipeline - BLEU Score Calculation
+The `evaluate` function measures the quality of generated text using the BLEU score. This function runs the RAG pipeline for a given query, compares the output against a reference answer, and computes the BLEU score with smoothing techniques.
+
+```python
+import argparse
+from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
+from rag import rag_pipeline_execute
+
+# Simulated RAG pipeline (simplified for example)
+def rag_pipeline(query, top_k, temperature):
+    return rag_pipeline_execute(query, top_k=top_k, temperature=temperature)
+
+# Evaluate BLEU score
+def evaluate(query, ground_truth, top_k, temperature):
+    # Get the RAG pipeline response
+    response = rag_pipeline(query, top_k, temperature)
+
+    # Tokenize the response and ground truth
+    reference = [ground_truth.split()]  # Reference should be a list of tokens
+    candidate = response.split()  # Candidate is the generated response tokens
+
+    # Apply smoothing to the BLEU score
+    smoothie = SmoothingFunction().method1  # Use method1 for smoothing
+    bleu_score = sentence_bleu(reference, candidate, smoothing_function=smoothie)
+
+    return bleu_score
+
+
+if __name__ == "__main__":
+    # Parse command-line arguments
+    parser = argparse.ArgumentParser(description="Evaluate BLEU score for a query using RAG pipeline")
+    parser.add_argument("--top_k", type=int, required=True, help="Number of top documents to retrieve")
+    parser.add_argument("--temperature", type=float, required=True, help="Temperature for the generator")
+    args = parser.parse_args()
+
+    #TODO: The quries and ground truth against which the BLEU score needs to be evaluated. They can be provided in the script below or loaded from an external volume. 
+    query = ""
+    ground_truth = " "
+
+    # Call evaluate with arguments from the command line
+    bleu_score = evaluate(query, ground_truth, args.top_k, args.temperature)
+    print(f"BLEU={bleu_score}")
+```
+*Note*: Make sure to return the result in the format of `<parameter>=<value>` for Katib's metrics collector to be able to utilize it. More ways to configure the output are available in [Kubeflow Metrics Collector][Katib_metrics_collector] guide.
+
+The above scripts will be used to build our `rag-pipeline` image, that will be specified in the `Experiment`.
+
+## STEP 3: Run a KATIB Experiment
+
+To optimize the RAG pipeline’s hyperparameters, let's define an `Experiment` CR. An experiment will define a single optimization run. More details on this API is available in the [documentation][katib_api].
+
+```yaml
+apiVersion: "kubeflow.org/v1beta1"
+kind: Experiment
+metadata:
+  name: rag-tuning-experiment
+  namespace: kubeflow
+spec:
+  objective:
+    type: maximize  # Ensures that Katib tries to maximize the BLEU score.
+    goal: 0.8  # Sets the desired BLEU score threshold.
+    objectiveMetricName: BLEU  # This should match with what we have provided in the script. 
+  algorithm:
+    algorithmName: grid
+  parameters:
+    - name: top_k # Controls the number of retrieved documents.
+      parameterType: int
+      feasibleSpace:
+        min: "1"
+        max: "5"
+        step: "1"  # Adding a step for discrete search
+    - name: temperature # Influences randomness in text generation.
+      parameterType: double
+      feasibleSpace:
+        min: "0.5"
+        max: "1.0"
+        step: "0.1"  # Adding a step for temperature
+  metricsCollectorSpec: #  Specifies how Katib collects experiment results.
+    collector:  
+      kind: StdOut  #  Tells Katib to extract metrics from standard output logs.
+  trialTemplate:
+    primaryContainerName: training-container
+    trialParameters:  #  Map hyperparameters (top_k and temperature) to their respective references in the job spec.
+      - name: top_k
+        description: Number of top documents to retrieve
+        reference: top_k
+      - name: temperature
+        description: Temperature for text generation
+        reference: temperature
+    trialSpec:
+      apiVersion: batch/v1
+      kind: Job
+      spec:
+        template:
+          spec:
+            containers:
+              - name: training-container
+                image: rag-pipeline:latest
+                command:
+                  - "python"
+                  - "/app/optimization-script.py"
+                  - "--top_k=${trialParameters.top_k}"
+                  - "--temperature=${trialParameters.temperature}"
+                resources:
+                  limits:
+                    cpu: "1"
+                    memory: "2Gi"
+            restartPolicy: Never
+```
+
+On applying this yaml on cluster, we can see our optimization script in action.
+
+```commandline
+kubectl get experiments.kubeflow.org -n kubeflow
+NAME                    TYPE      STATUS   AGE
+rag-tuning-experiment   Running   True     10m
+```
+
+We can also see the trials being run to search for the optimized parameter:
+
+```commandline
+kubectl get trials --all-namespaces
+NAMESPACE   NAME                             TYPE      STATUS   AGE
+kubeflow    rag-tuning-experiment-7wskq9b9   Running   True     10m
+kubeflow    rag-tuning-experiment-cll6bt4z   Running   True     10m
+kubeflow    rag-tuning-experiment-hzxrzq2t   Running   True     10m
+```
+
+The list of completed trials and their results will be shown in the UI as below. Steps to access Katib UI are available [here][katib_ui].:  
+
+![completed_runs](../images/2025-02-21-katib-rag-optimization/katib_experiment_run.jpeg)
+![trial details](../images/2025-02-21-katib-rag-optimization/katib_ui.jpeg)
+
+# Conclusion
+
+In this experiment, we leveraged Kubeflow Katib to optimize a Retrieval-Augmented Generation (RAG) pipeline, systematically tuning key hyperparameters like top_k and temperature to enhance retrieval precision and generative response quality.
+
+For anyone working with RAG systems or hyperparameter optimization, Katib is a game-changer—enabling scalable, efficient, and intelligent tuning of machine learning models! In this experiment. Hope this helps you streamline hyperparameter tuning and unlock new efficiencies in your ML workflows!
+
+[Katib]: https://www.kubeflow.org/docs/components/katib/
+[rag]: https://en.wikipedia.org/wiki/Retrieval-augmented_generation
+[katib_installation]: https://www.kubeflow.org/docs/components/katib/installation/
+[bleu]: https://huggingface.co/spaces/evaluate-metric/bleu
+[Katib_metrics_collector]: https://www.kubeflow.org/docs/components/katib/user-guides/metrics-collector/#pull-based-metrics-collector
+[katib_api]: https://www.kubeflow.org/docs/components/katib/reference/architecture/#experiment
+[katib_ui]: https://www.kubeflow.org/docs/components/katib/user-guides/katib-ui/