Use AutoAI RAG with custom foundation model.ipynb

Kernel: autoai_rag

AutoAI RAG experiment with custom foundation model.

Disclaimers

• Use only Projects and Spaces that are available in the watsonx context.

Notebook content

This notebook demonstrates how to deploy custom foundation model and use this model in AutoAI RAG experiment. The data used in this notebook is from the Granite Code Models paper.

Some familiarity with Python is helpful. This notebook uses Python 3.12.

Learning goal

The learning goals of this notebook are:

• How to deploy your own foundation models with huggingface hub

• Create an AutoAI RAG job that will find the best RAG pattern based on custom foundation model used during the experiment

Contents

This notebook contains the following parts:

• Set up the environment

• Prerequisites

• Create API Client instance.

• Download custom model from hugging face

• Deploy the model

• Prepare the data for the AutoAI RAG experiment

• Run the AutoAI RAG experiment

• Query generated pattern locally

• Summary

Set up the environment

%pip install -U wget | tail -n 1
%pip install -U 'ibm-watsonx-ai[rag]>=1.3.26' | tail -n 1

Requirement already satisfied: wget in /Users/michalsteczko/anaconda3/envs/autoai_rag/lib/python3.11/site-packages (3.2)
Note: you may need to restart the kernel to use updated packages.
Requirement already satisfied: pyasn1<0.7.0,>=0.4.6 in /Users/michalsteczko/anaconda3/envs/autoai_rag/lib/python3.11/site-packages (from pyasn1-modules>=0.2.1->google-auth>=1.0.1->kubernetes>=28.1.0->chromadb!=0.5.4,!=0.5.5,<0.6.0,>=0.4.0->langchain-chroma==0.1.4->ibm-watsonx-ai[rag]>=1.3.26) (0.6.1)
Note: you may need to restart the kernel to use updated packages.

Prerequisites

Please fill below values to be able to move forward:

• URL - url which points to your CPD instance

• USERNAME - username to your CPD instance

• PASSWORD - password to your CPD instance associated with your username.

• INSTANCE_ID - your CPD instance ID

• VERSION - CPD version which your instance supports, but it has to be at least 5.2

• PROJECT_ID - ID of the project associated with your username and instance.

URL = "PUT YOUR CPD INSTANCE URL HERE"
USERNAME = "PUT YOUR USERNAME HERE"
PASSWORD = "PUT YOUR PASSWORD HERE"
INSTANCE_ID = "PUT YOUR INSTANCE ID HERE"
VERSION = "5.2"

PROJECT_ID = "PUT YOUR PROJECT ID HERE"

BUCKET_BENCHMARK_JSON_FILE_PATH = "benchmark.json"

Create API Client instance.

This client will allow us to connect with the IBM services.

from ibm_watsonx_ai import APIClient, Credentials

credentials = Credentials(
    url=URL,
    username=USERNAME,
    password=PASSWORD,
    instance_id=INSTANCE_ID,
    version=VERSION,
    verify=False,
)

client = APIClient(credentials=credentials, project_id=PROJECT_ID)

Create custom PVC for custom foundation model

Once you have our model files in our COS bucket you need to create our custom model specification.
Please note that this is beyond the scope of this notebook, refer to this documentation to be able to use your custom model on your CPD cluster.

Deploy the model

Check the docs to avoid any problems during model deployment here.

Create custom model repository

from ibm_watsonx_ai.foundation_models import get_custom_model_specs

custom_model_spec = get_custom_model_specs(api_client=client, limit=1)
model_id = custom_model_spec["resources"][0]["model_id"]
model_id

'deepseek-r1-distill-llama-8b'

sw_spec_id = client.software_specifications.get_id_by_name("watsonx-cfm-caikit-1.0")
sw_metadata = {
    client.repository.ModelMetaNames.NAME: "My custom deployment",
    client.repository.ModelMetaNames.SOFTWARE_SPEC_ID: sw_spec_id,
    client.repository.ModelMetaNames.TYPE: client.repository.ModelAssetTypes.CUSTOM_FOUNDATION_MODEL_1_0,
}

stored_model_details = client.repository.store_model(
    model=model_id, meta_props=sw_metadata
)
stored_model_asset_id = client.repository.get_model_id(stored_model_details)
client.repository.list(framework_filter="custom_foundation_model_1.0")

Create custom foundation model hardware specification

hardware_spec_meta_props = {
    client.hardware_specifications.ConfigurationMetaNames.NAME: "Custom GPU hw spec",
    client.hardware_specifications.ConfigurationMetaNames.NODES: {
        "cpu": {"units": "2"},
        "mem": {"size": "128Gi"},
        "gpu": {"num_gpu": 1},
    },
}

hw_spec_details = client.hardware_specifications.store(hardware_spec_meta_props)

hw_spec_details = client.hardware_specifications.get_details(
    client.hardware_specifications.get_id_by_name("Custom GPU hw spec")
)
hw_spec_id = client.hardware_specifications.get_id(hw_spec_details)

Perform custom model deployment

MAX_SEQUENCE_LENGTH = 32_000
MAX_NEW_TOKENS = 1000
MIN_NEW_TOKENS = 1
MAX_BATCH_SIZE = 1024
metadata = {
    client.deployments.ConfigurationMetaNames.NAME: "My custom foundation model",
    client.deployments.ConfigurationMetaNames.DESCRIPTION: "My custom foundation model",
    client.deployments.ConfigurationMetaNames.ONLINE: {},
    # optionally overwrite model parameters here
    client.deployments.ConfigurationMetaNames.HARDWARE_SPEC: {
        "name": "Custom GPU hw spec"
    },
    client.deployments.ConfigurationMetaNames.FOUNDATION_MODEL: {
        "max_sequence_length": MAX_SEQUENCE_LENGTH,
        "max_new_tokens": MAX_NEW_TOKENS,
        "max_batch_size": MAX_BATCH_SIZE,
    },
    client.deployments.ConfigurationMetaNames.SERVING_NAME: "pllum_12b_instruct",  # must be unique
}
deployment_details = client.deployments.create(stored_model_asset_id, metadata)
deployment_id = deployment_details["metadata"]["id"]

######################################################################################

Synchronous deployment creation for id: 'e2a38f94-31c5-4106-880b-f41423fd42db' started

######################################################################################


initializing
Note: online_url and serving_urls are deprecated and will be removed in a future release. Use inference instead.
..........
ready


-----------------------------------------------------------------------------------------------
Successfully finished deployment creation, deployment_id='30702554-42cc-43e8-82a5-5dc7826d719d'
-----------------------------------------------------------------------------------------------

Prepare the data for the AutoAI RAG experiment

Download granite_code_models.pdf document

import wget

data_url = "https://arxiv.org/pdf/2405.04324"

byom_input_filename = "granite_code_models.pdf"

wget.download(data_url, byom_input_filename)

'granite_code_models.pdf'

Create data asset with your training data

document_asset_details = client.data_assets.create(name=byom_input_filename, file_path=byom_input_filename)

document_asset_id = client.data_assets.get_id(document_asset_details)
document_asset_id

Creating data asset...
SUCCESS

'fb05d631-8c5e-4641-bc72-6633ba06cbe2'

from ibm_watsonx_ai.helpers import DataConnection

input_data_references = [DataConnection(data_asset_id=document_asset_id)]

Create your own benchmark.json file to ask the questions related to the document

import json

local_benchmark_json_filename = "benchmark.json"

benchmarking_data = [
    {
        "question": "What are the two main variants of Granite Code models?",
        "correct_answer": "The two main variants are Granite Code Base and Granite Code Instruct.",
        "correct_answer_document_ids": [byom_input_filename],
    },
    {
        "question": "What is the purpose of Granite Code Instruct models?",
        "correct_answer": "Granite Code Instruct models are finetuned for instruction-following tasks using datasets like CommitPack, OASST, HelpSteer, and synthetic code instruction datasets, aiming to improve reasoning and instruction-following capabilities.",
        "correct_answer_document_ids": [byom_input_filename],
    },
    {
        "question": "What is the licensing model for Granite Code models?",
        "correct_answer": "Granite Code models are released under the Apache 2.0 license, ensuring permissive and enterprise-friendly usage.",
        "correct_answer_document_ids": [byom_input_filename],
    },
]

with open(local_benchmark_json_filename, mode="w", encoding="utf-8") as fp:
    json.dump(benchmarking_data, fp, indent=4)

Create data asset with benchmark.json file

test_asset_details = client.data_assets.create(name=local_benchmark_json_filename, file_path=local_benchmark_json_filename)

test_asset_id = client.data_assets.get_id(test_asset_details)
test_asset_id

Creating data asset...
SUCCESS

'fc14fca8-93c4-4b5e-a72f-d19fcc0aeaf4'

test_data_references = [DataConnection(data_asset_id=test_asset_id)]

Run the AutoAI RAG experiment

Provide the input information for AutoAI RAG optimizer:

• custom_prompt_template_text - custom prompt template text which will be used to query your own foundation model

• custom_context_template_text - custom context template text which will be used to query your own foundation model

• name - experiment name

• description - experiment description

• max_number_of_rag_patterns - maximum number of RAG patterns to create

• optimization_metrics - target optimization metrics

from ibm_watsonx_ai.experiment import AutoAI
from ibm_watsonx_ai.foundation_models.schema import (
    AutoAIRAGCustomModelConfig,
    AutoAIRAGModelParams,
)

experiment = AutoAI(credentials, project_id=PROJECT_ID)

custom_prompt_template_text = (
    "Answer my question {question} related to these documents {reference_documents}."
)
custom_context_template_text = "My document {document}"

parameters = AutoAIRAGModelParams(max_sequence_length=32_000)
custom_foundation_model_config = AutoAIRAGCustomModelConfig(
    deployment_id=deployment_id,
    project_id=PROJECT_ID,
    prompt_template_text=custom_prompt_template_text,
    context_template_text=custom_context_template_text,
    parameters=parameters,
)

rag_optimizer = experiment.rag_optimizer(
    name="AutoAI RAG - Custom foundation model experiment",
    description="AutoAI RAG experiment with custom foundation model.",
    max_number_of_rag_patterns=4,
    optimization_metrics=["faithfulness"],
    foundation_models=[custom_foundation_model_config],
)

rag_optimizer.run(
    test_data_references=test_data_references,
    input_data_references=input_data_references,
    background_mode=False
)

##############################################

Running '7878e0b8-94a6-4e9a-b1b4-f3c62fa70786'

##############################################


pending....
running...........................................................................................................................................
completed
Training of '7878e0b8-94a6-4e9a-b1b4-f3c62fa70786' finished successfully.

{'entity': {'hardware_spec': {'id': 'a6c4923b-b8e4-444c-9f43-8a7ec3020110',
   'name': 'L'},
  'input_data_references': [{'location': {'href': '/v2/assets/fb05d631-8c5e-4641-bc72-6633ba06cbe2?project_id=19a157bf-5c1c-4773-812a-520d8068d84a',
     'id': 'fb05d631-8c5e-4641-bc72-6633ba06cbe2'},
    'type': 'data_asset'}],
  'parameters': {'constraints': {'generation': {'foundation_models': [{'context_template_text': 'My document {document}',
       'deployment_id': '30702554-42cc-43e8-82a5-5dc7826d719d',
       'parameters': {'max_sequence_length': 32000},
       'project_id': '19a157bf-5c1c-4773-812a-520d8068d84a',
       'prompt_template_text': 'Answer my question {question} related to these documents {reference_documents}.'}]},
    'max_number_of_rag_patterns': 4},
   'optimization': {'metrics': ['faithfulness']},
   'output_logs': True},
  'results': [{'context': {'iteration': 0,
     'max_combinations': 40,
     'rag_pattern': {'composition_steps': ['model_selection',
       'chunking',
       'embeddings',
       'retrieval',
       'generation'],
      'duration_seconds': 193,
      'location': {'evaluation_results': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern1/evaluation_results.json',
       'indexing_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern1/indexing_inference_notebook.ipynb',
       'inference_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern1/indexing_inference_notebook.ipynb',
       'inference_service_code': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern1/inference_ai_service.gz',
       'inference_service_metadata': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern1/inference_service_metadata.json'},
      'name': 'Pattern1',
      'settings': {'chunking': {'chunk_overlap': 128,
        'chunk_size': 512,
        'method': 'recursive'},
       'embeddings': {'model_id': 'ibm/slate-125m-english-rtrvr',
        'truncate_input_tokens': 512,
        'truncate_strategy': 'left'},
       'generation': {'context_template_text': 'My document {document}',
        'deployment_id': '30702554-42cc-43e8-82a5-5dc7826d719d',
        'parameters': {'decoding_method': 'greedy',
         'max_new_tokens': 1000,
         'max_sequence_length': 32000,
         'min_new_tokens': 1},
        'prompt_template_text': 'Answer my question {question} related to these documents {reference_documents}.',
        'word_to_token_ratio': 2.0},
       'retrieval': {'method': 'simple', 'number_of_chunks': 3},
       'vector_store': {'datasource_type': 'chroma',
        'distance_metric': 'cosine',
        'index_name': 'autoai_rag_7878e0b8_20250627145642',
        'operation': 'upsert',
        'schema': {'fields': [{'description': 'text field',
           'name': 'text',
           'role': 'text',
           'type': 'string'},
          {'description': 'document name field',
           'name': 'document_id',
           'role': 'document_name',
           'type': 'string'},
          {'description': 'chunk starting token position in the source document',
           'name': 'start_index',
           'role': 'start_index',
           'type': 'number'},
          {'description': 'chunk number per document',
           'name': 'sequence_number',
           'role': 'sequence_number',
           'type': 'number'},
          {'description': 'vector embeddings',
           'name': 'vector',
           'role': 'vector_embeddings',
           'type': 'array'}],
         'id': 'autoai_rag_1.0',
         'name': 'Document schema using open-source loaders',
         'type': 'struct'}}},
      'settings_importance': {'chunking': [{'importance': 0.125,
         'parameter': 'chunk_size'},
        {'importance': 0.125, 'parameter': 'chunk_overlap'}],
       'embeddings': [{'importance': 0.125, 'parameter': 'embedding_model'}],
       'generation': [{'importance': 0.125, 'parameter': 'foundation_model'}],
       'retrieval': [{'importance': 0.125, 'parameter': 'retrieval_method'},
        {'importance': 0.125, 'parameter': 'window_size'},
        {'importance': 0.125, 'parameter': 'number_of_chunks'}]}},
     'software_spec': {}},
    'metrics': {'test_data': [{'ci_high': 0.7937,
       'ci_low': 0.0,
       'mean': 0.5079,
       'metric_name': 'answer_correctness'},
      {'ci_high': 0.1414,
       'ci_low': 0.0,
       'mean': 0.072,
       'metric_name': 'faithfulness'},
      {'mean': 1.0, 'metric_name': 'context_correctness'}]}},
   {'context': {'iteration': 1,
     'max_combinations': 40,
     'rag_pattern': {'composition_steps': ['model_selection',
       'chunking',
       'embeddings',
       'retrieval',
       'generation'],
      'duration_seconds': 266,
      'location': {'evaluation_results': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/evaluation_results.json',
       'indexing_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/indexing_inference_notebook.ipynb',
       'inference_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/indexing_inference_notebook.ipynb',
       'inference_service_code': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/inference_ai_service.gz',
       'inference_service_metadata': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/inference_service_metadata.json'},
      'name': 'Pattern2',
      'settings': {'chunking': {'chunk_overlap': 256,
        'chunk_size': 512,
        'method': 'recursive'},
       'embeddings': {'model_id': 'ibm/slate-125m-english-rtrvr',
        'truncate_input_tokens': 512,
        'truncate_strategy': 'left'},
       'generation': {'context_template_text': 'My document {document}',
        'deployment_id': '30702554-42cc-43e8-82a5-5dc7826d719d',
        'parameters': {'decoding_method': 'greedy',
         'max_new_tokens': 1000,
         'max_sequence_length': 32000,
         'min_new_tokens': 1},
        'prompt_template_text': 'Answer my question {question} related to these documents {reference_documents}.',
        'word_to_token_ratio': 2.0},
       'retrieval': {'method': 'window',
        'number_of_chunks': 3,
        'window_size': 1},
       'vector_store': {'datasource_type': 'chroma',
        'distance_metric': 'cosine',
        'index_name': 'autoai_rag_7878e0b8_20250627150244',
        'operation': 'upsert',
        'schema': {'fields': [{'description': 'text field',
           'name': 'text',
           'role': 'text',
           'type': 'string'},
          {'description': 'document name field',
           'name': 'document_id',
           'role': 'document_name',
           'type': 'string'},
          {'description': 'chunk starting token position in the source document',
           'name': 'start_index',
           'role': 'start_index',
           'type': 'number'},
          {'description': 'chunk number per document',
           'name': 'sequence_number',
           'role': 'sequence_number',
           'type': 'number'},
          {'description': 'vector embeddings',
           'name': 'vector',
           'role': 'vector_embeddings',
           'type': 'array'}],
         'id': 'autoai_rag_1.0',
         'name': 'Document schema using open-source loaders',
         'type': 'struct'}}},
      'settings_importance': {'chunking': [{'importance': 0.0,
         'parameter': 'chunk_size'},
        {'importance': 0.0, 'parameter': 'chunk_overlap'}],
       'embeddings': [{'importance': 0.0, 'parameter': 'embedding_model'}],
       'generation': [{'importance': 0.0, 'parameter': 'foundation_model'}],
       'retrieval': [{'importance': 1.0, 'parameter': 'retrieval_method'},
        {'importance': 0.0, 'parameter': 'window_size'},
        {'importance': 0.0, 'parameter': 'number_of_chunks'}]}},
     'software_spec': {}},
    'metrics': {'test_data': [{'ci_high': 1.0,
       'ci_low': 0.0,
       'mean': 0.4815,
       'metric_name': 'answer_correctness'},
      {'ci_high': 0.3475,
       'ci_low': 0.0,
       'mean': 0.2144,
       'metric_name': 'faithfulness'},
      {'mean': 1.0, 'metric_name': 'context_correctness'}]}},
   {'context': {'iteration': 2,
     'max_combinations': 40,
     'rag_pattern': {'composition_steps': ['model_selection',
       'chunking',
       'embeddings',
       'retrieval',
       'generation'],
      'duration_seconds': 180,
      'location': {'evaluation_results': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern3/evaluation_results.json',
       'indexing_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern3/indexing_inference_notebook.ipynb',
       'inference_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern3/indexing_inference_notebook.ipynb',
       'inference_service_code': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern3/inference_ai_service.gz',
       'inference_service_metadata': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern3/inference_service_metadata.json'},
      'name': 'Pattern3',
      'settings': {'chunking': {'chunk_overlap': 256,
        'chunk_size': 1024,
        'method': 'recursive'},
       'embeddings': {'model_id': 'ibm/slate-125m-english-rtrvr',
        'truncate_input_tokens': 512,
        'truncate_strategy': 'left'},
       'generation': {'context_template_text': 'My document {document}',
        'deployment_id': '30702554-42cc-43e8-82a5-5dc7826d719d',
        'parameters': {'decoding_method': 'greedy',
         'max_new_tokens': 1000,
         'max_sequence_length': 32000,
         'min_new_tokens': 1},
        'prompt_template_text': 'Answer my question {question} related to these documents {reference_documents}.',
        'word_to_token_ratio': 2.0},
       'retrieval': {'method': 'simple', 'number_of_chunks': 5},
       'vector_store': {'datasource_type': 'chroma',
        'distance_metric': 'cosine',
        'index_name': 'autoai_rag_7878e0b8_20250627150710',
        'operation': 'upsert',
        'schema': {'fields': [{'description': 'text field',
           'name': 'text',
           'role': 'text',
           'type': 'string'},
          {'description': 'document name field',
           'name': 'document_id',
           'role': 'document_name',
           'type': 'string'},
          {'description': 'chunk starting token position in the source document',
           'name': 'start_index',
           'role': 'start_index',
           'type': 'number'},
          {'description': 'chunk number per document',
           'name': 'sequence_number',
           'role': 'sequence_number',
           'type': 'number'},
          {'description': 'vector embeddings',
           'name': 'vector',
           'role': 'vector_embeddings',
           'type': 'array'}],
         'id': 'autoai_rag_1.0',
         'name': 'Document schema using open-source loaders',
         'type': 'struct'}}},
      'settings_importance': {'chunking': [{'importance': 0.18471761,
         'parameter': 'chunk_size'},
        {'importance': 0.2972388, 'parameter': 'chunk_overlap'}],
       'embeddings': [{'importance': 0.0, 'parameter': 'embedding_model'}],
       'generation': [{'importance': 0.0, 'parameter': 'foundation_model'}],
       'retrieval': [{'importance': 0.35138428,
         'parameter': 'retrieval_method'},
        {'importance': 0.061107427, 'parameter': 'window_size'},
        {'importance': 0.10555187, 'parameter': 'number_of_chunks'}]}},
     'software_spec': {}},
    'metrics': {'test_data': [{'ci_high': 0.4545,
       'ci_low': 0.0,
       'mean': 0.2256,
       'metric_name': 'answer_correctness'},
      {'ci_high': 0.3934,
       'ci_low': 0.0323,
       'mean': 0.1634,
       'metric_name': 'faithfulness'},
      {'mean': 1.0, 'metric_name': 'context_correctness'}]}}],
  'results_reference': {'location': {'path': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results',
    'training': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786',
    'training_status': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/training-status.json',
    'training_log': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/output.log',
    'assets_path': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/assets'},
   'type': 'fs'},
  'status': {'completed_at': '2025-06-27T15:10:12.169Z',
   'message': {'level': 'warning',
    'text': 'AAR021W: Warnings:\n AR1128W: Unable to retrieve chunks and generate answers for deployment with ID: 30702554-42cc-43e8-82a5-5dc7826d719d.'},
   'running_at': '2025-06-27T14:56:38.000Z',
   'state': 'completed',
   'step': 'generation'},
  'test_data_references': [{'location': {'href': '/v2/assets/fc14fca8-93c4-4b5e-a72f-d19fcc0aeaf4?project_id=19a157bf-5c1c-4773-812a-520d8068d84a',
     'id': 'fc14fca8-93c4-4b5e-a72f-d19fcc0aeaf4'},
    'type': 'data_asset'}],
  'timestamp': '2025-06-27T15:10:17.754Z'},
 'metadata': {'created_at': '2025-06-27T14:56:03.343Z',
  'description': 'AutoAI RAG experiment with custom foundation model.',
  'id': '7878e0b8-94a6-4e9a-b1b4-f3c62fa70786',
  'modified_at': '2025-06-27T15:10:12.177Z',
  'name': 'AutoAI RAG - Custom foundation model experiment',
  'project_id': '19a157bf-5c1c-4773-812a-520d8068d84a'}}

rag_optimizer.get_run_details()

{'entity': {'hardware_spec': {'id': 'a6c4923b-b8e4-444c-9f43-8a7ec3020110',
   'name': 'L'},
  'input_data_references': [{'location': {'href': '/v2/assets/fb05d631-8c5e-4641-bc72-6633ba06cbe2?project_id=19a157bf-5c1c-4773-812a-520d8068d84a',
     'id': 'fb05d631-8c5e-4641-bc72-6633ba06cbe2'},
    'type': 'data_asset'}],
  'parameters': {'constraints': {'generation': {'foundation_models': [{'context_template_text': 'My document {document}',
       'deployment_id': '30702554-42cc-43e8-82a5-5dc7826d719d',
       'parameters': {'max_sequence_length': 32000},
       'project_id': '19a157bf-5c1c-4773-812a-520d8068d84a',
       'prompt_template_text': 'Answer my question {question} related to these documents {reference_documents}.'}]},
    'max_number_of_rag_patterns': 4},
   'optimization': {'metrics': ['faithfulness']},
   'output_logs': True},
  'results': [{'context': {'iteration': 0,
     'max_combinations': 40,
     'rag_pattern': {'composition_steps': ['model_selection',
       'chunking',
       'embeddings',
       'retrieval',
       'generation'],
      'duration_seconds': 193,
      'location': {'evaluation_results': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern1/evaluation_results.json',
       'indexing_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern1/indexing_inference_notebook.ipynb',
       'inference_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern1/indexing_inference_notebook.ipynb',
       'inference_service_code': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern1/inference_ai_service.gz',
       'inference_service_metadata': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern1/inference_service_metadata.json'},
      'name': 'Pattern1',
      'settings': {'chunking': {'chunk_overlap': 128,
        'chunk_size': 512,
        'method': 'recursive'},
       'embeddings': {'model_id': 'ibm/slate-125m-english-rtrvr',
        'truncate_input_tokens': 512,
        'truncate_strategy': 'left'},
       'generation': {'context_template_text': 'My document {document}',
        'deployment_id': '30702554-42cc-43e8-82a5-5dc7826d719d',
        'parameters': {'decoding_method': 'greedy',
         'max_new_tokens': 1000,
         'max_sequence_length': 32000,
         'min_new_tokens': 1},
        'prompt_template_text': 'Answer my question {question} related to these documents {reference_documents}.',
        'word_to_token_ratio': 2.0},
       'retrieval': {'method': 'simple', 'number_of_chunks': 3},
       'vector_store': {'datasource_type': 'chroma',
        'distance_metric': 'cosine',
        'index_name': 'autoai_rag_7878e0b8_20250627145642',
        'operation': 'upsert',
        'schema': {'fields': [{'description': 'text field',
           'name': 'text',
           'role': 'text',
           'type': 'string'},
          {'description': 'document name field',
           'name': 'document_id',
           'role': 'document_name',
           'type': 'string'},
          {'description': 'chunk starting token position in the source document',
           'name': 'start_index',
           'role': 'start_index',
           'type': 'number'},
          {'description': 'chunk number per document',
           'name': 'sequence_number',
           'role': 'sequence_number',
           'type': 'number'},
          {'description': 'vector embeddings',
           'name': 'vector',
           'role': 'vector_embeddings',
           'type': 'array'}],
         'id': 'autoai_rag_1.0',
         'name': 'Document schema using open-source loaders',
         'type': 'struct'}}},
      'settings_importance': {'chunking': [{'importance': 0.125,
         'parameter': 'chunk_size'},
        {'importance': 0.125, 'parameter': 'chunk_overlap'}],
       'embeddings': [{'importance': 0.125, 'parameter': 'embedding_model'}],
       'generation': [{'importance': 0.125, 'parameter': 'foundation_model'}],
       'retrieval': [{'importance': 0.125, 'parameter': 'retrieval_method'},
        {'importance': 0.125, 'parameter': 'window_size'},
        {'importance': 0.125, 'parameter': 'number_of_chunks'}]}},
     'software_spec': {}},
    'metrics': {'test_data': [{'ci_high': 0.7937,
       'ci_low': 0.0,
       'mean': 0.5079,
       'metric_name': 'answer_correctness'},
      {'ci_high': 0.1414,
       'ci_low': 0.0,
       'mean': 0.072,
       'metric_name': 'faithfulness'},
      {'mean': 1.0, 'metric_name': 'context_correctness'}]}},
   {'context': {'iteration': 1,
     'max_combinations': 40,
     'rag_pattern': {'composition_steps': ['model_selection',
       'chunking',
       'embeddings',
       'retrieval',
       'generation'],
      'duration_seconds': 266,
      'location': {'evaluation_results': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/evaluation_results.json',
       'indexing_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/indexing_inference_notebook.ipynb',
       'inference_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/indexing_inference_notebook.ipynb',
       'inference_service_code': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/inference_ai_service.gz',
       'inference_service_metadata': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/inference_service_metadata.json'},
      'name': 'Pattern2',
      'settings': {'chunking': {'chunk_overlap': 256,
        'chunk_size': 512,
        'method': 'recursive'},
       'embeddings': {'model_id': 'ibm/slate-125m-english-rtrvr',
        'truncate_input_tokens': 512,
        'truncate_strategy': 'left'},
       'generation': {'context_template_text': 'My document {document}',
        'deployment_id': '30702554-42cc-43e8-82a5-5dc7826d719d',
        'parameters': {'decoding_method': 'greedy',
         'max_new_tokens': 1000,
         'max_sequence_length': 32000,
         'min_new_tokens': 1},
        'prompt_template_text': 'Answer my question {question} related to these documents {reference_documents}.',
        'word_to_token_ratio': 2.0},
       'retrieval': {'method': 'window',
        'number_of_chunks': 3,
        'window_size': 1},
       'vector_store': {'datasource_type': 'chroma',
        'distance_metric': 'cosine',
        'index_name': 'autoai_rag_7878e0b8_20250627150244',
        'operation': 'upsert',
        'schema': {'fields': [{'description': 'text field',
           'name': 'text',
           'role': 'text',
           'type': 'string'},
          {'description': 'document name field',
           'name': 'document_id',
           'role': 'document_name',
           'type': 'string'},
          {'description': 'chunk starting token position in the source document',
           'name': 'start_index',
           'role': 'start_index',
           'type': 'number'},
          {'description': 'chunk number per document',
           'name': 'sequence_number',
           'role': 'sequence_number',
           'type': 'number'},
          {'description': 'vector embeddings',
           'name': 'vector',
           'role': 'vector_embeddings',
           'type': 'array'}],
         'id': 'autoai_rag_1.0',
         'name': 'Document schema using open-source loaders',
         'type': 'struct'}}},
      'settings_importance': {'chunking': [{'importance': 0.0,
         'parameter': 'chunk_size'},
        {'importance': 0.0, 'parameter': 'chunk_overlap'}],
       'embeddings': [{'importance': 0.0, 'parameter': 'embedding_model'}],
       'generation': [{'importance': 0.0, 'parameter': 'foundation_model'}],
       'retrieval': [{'importance': 1.0, 'parameter': 'retrieval_method'},
        {'importance': 0.0, 'parameter': 'window_size'},
        {'importance': 0.0, 'parameter': 'number_of_chunks'}]}},
     'software_spec': {}},
    'metrics': {'test_data': [{'ci_high': 1.0,
       'ci_low': 0.0,
       'mean': 0.4815,
       'metric_name': 'answer_correctness'},
      {'ci_high': 0.3475,
       'ci_low': 0.0,
       'mean': 0.2144,
       'metric_name': 'faithfulness'},
      {'mean': 1.0, 'metric_name': 'context_correctness'}]}},
   {'context': {'iteration': 2,
     'max_combinations': 40,
     'rag_pattern': {'composition_steps': ['model_selection',
       'chunking',
       'embeddings',
       'retrieval',
       'generation'],
      'duration_seconds': 180,
      'location': {'evaluation_results': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern3/evaluation_results.json',
       'indexing_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern3/indexing_inference_notebook.ipynb',
       'inference_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern3/indexing_inference_notebook.ipynb',
       'inference_service_code': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern3/inference_ai_service.gz',
       'inference_service_metadata': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern3/inference_service_metadata.json'},
      'name': 'Pattern3',
      'settings': {'chunking': {'chunk_overlap': 256,
        'chunk_size': 1024,
        'method': 'recursive'},
       'embeddings': {'model_id': 'ibm/slate-125m-english-rtrvr',
        'truncate_input_tokens': 512,
        'truncate_strategy': 'left'},
       'generation': {'context_template_text': 'My document {document}',
        'deployment_id': '30702554-42cc-43e8-82a5-5dc7826d719d',
        'parameters': {'decoding_method': 'greedy',
         'max_new_tokens': 1000,
         'max_sequence_length': 32000,
         'min_new_tokens': 1},
        'prompt_template_text': 'Answer my question {question} related to these documents {reference_documents}.',
        'word_to_token_ratio': 2.0},
       'retrieval': {'method': 'simple', 'number_of_chunks': 5},
       'vector_store': {'datasource_type': 'chroma',
        'distance_metric': 'cosine',
        'index_name': 'autoai_rag_7878e0b8_20250627150710',
        'operation': 'upsert',
        'schema': {'fields': [{'description': 'text field',
           'name': 'text',
           'role': 'text',
           'type': 'string'},
          {'description': 'document name field',
           'name': 'document_id',
           'role': 'document_name',
           'type': 'string'},
          {'description': 'chunk starting token position in the source document',
           'name': 'start_index',
           'role': 'start_index',
           'type': 'number'},
          {'description': 'chunk number per document',
           'name': 'sequence_number',
           'role': 'sequence_number',
           'type': 'number'},
          {'description': 'vector embeddings',
           'name': 'vector',
           'role': 'vector_embeddings',
           'type': 'array'}],
         'id': 'autoai_rag_1.0',
         'name': 'Document schema using open-source loaders',
         'type': 'struct'}}},
      'settings_importance': {'chunking': [{'importance': 0.18471761,
         'parameter': 'chunk_size'},
        {'importance': 0.2972388, 'parameter': 'chunk_overlap'}],
       'embeddings': [{'importance': 0.0, 'parameter': 'embedding_model'}],
       'generation': [{'importance': 0.0, 'parameter': 'foundation_model'}],
       'retrieval': [{'importance': 0.35138428,
         'parameter': 'retrieval_method'},
        {'importance': 0.061107427, 'parameter': 'window_size'},
        {'importance': 0.10555187, 'parameter': 'number_of_chunks'}]}},
     'software_spec': {}},
    'metrics': {'test_data': [{'ci_high': 0.4545,
       'ci_low': 0.0,
       'mean': 0.2256,
       'metric_name': 'answer_correctness'},
      {'ci_high': 0.3934,
       'ci_low': 0.0323,
       'mean': 0.1634,
       'metric_name': 'faithfulness'},
      {'mean': 1.0, 'metric_name': 'context_correctness'}]}}],
  'results_reference': {'location': {'path': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results',
    'training': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786',
    'training_status': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/training-status.json',
    'training_log': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/output.log',
    'assets_path': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/assets'},
   'type': 'fs'},
  'status': {'completed_at': '2025-06-27T15:10:12.169Z',
   'message': {'level': 'warning',
    'text': 'AAR021W: Warnings:\n AR1128W: Unable to retrieve chunks and generate answers for deployment with ID: 30702554-42cc-43e8-82a5-5dc7826d719d.'},
   'running_at': '2025-06-27T14:56:38.000Z',
   'state': 'completed',
   'step': 'generation'},
  'test_data_references': [{'location': {'href': '/v2/assets/fc14fca8-93c4-4b5e-a72f-d19fcc0aeaf4?project_id=19a157bf-5c1c-4773-812a-520d8068d84a',
     'id': 'fc14fca8-93c4-4b5e-a72f-d19fcc0aeaf4'},
    'type': 'data_asset'}],
  'timestamp': '2025-06-27T15:10:17.754Z'},
 'metadata': {'created_at': '2025-06-27T14:56:03.343Z',
  'description': 'AutoAI RAG experiment with custom foundation model.',
  'id': '7878e0b8-94a6-4e9a-b1b4-f3c62fa70786',
  'modified_at': '2025-06-27T15:10:12.177Z',
  'name': 'AutoAI RAG - Custom foundation model experiment',
  'project_id': '19a157bf-5c1c-4773-812a-520d8068d84a'}}

summary = rag_optimizer.summary()
summary

best_pattern_name = summary.index.values[0]
print("Best pattern is:", best_pattern_name)

best_pattern = rag_optimizer.get_pattern()

Best pattern is: Pattern2

rag_optimizer.get_pattern_details(pattern_name=best_pattern_name)

{'composition_steps': ['model_selection',
  'chunking',
  'embeddings',
  'retrieval',
  'generation'],
 'duration_seconds': 266,
 'location': {'evaluation_results': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/evaluation_results.json',
  'indexing_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/indexing_inference_notebook.ipynb',
  'inference_notebook': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/indexing_inference_notebook.ipynb',
  'inference_service_code': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/inference_ai_service.gz',
  'inference_service_metadata': '/projects/19a157bf-5c1c-4773-812a-520d8068d84a/assets/autorag/results/7878e0b8-94a6-4e9a-b1b4-f3c62fa70786/Pattern2/inference_service_metadata.json'},
 'name': 'Pattern2',
 'settings': {'chunking': {'chunk_overlap': 256,
   'chunk_size': 512,
   'method': 'recursive'},
  'embeddings': {'model_id': 'ibm/slate-125m-english-rtrvr',
   'truncate_input_tokens': 512,
   'truncate_strategy': 'left'},
  'generation': {'context_template_text': 'My document {document}',
   'deployment_id': '30702554-42cc-43e8-82a5-5dc7826d719d',
   'parameters': {'decoding_method': 'greedy',
    'max_new_tokens': 1000,
    'min_new_tokens': 1},
   'prompt_template_text': 'Answer my question {question} related to these documents {reference_documents}.',
   'word_to_token_ratio': 2.0},
  'retrieval': {'method': 'window', 'number_of_chunks': 3, 'window_size': 1},
  'vector_store': {'datasource_type': 'chroma',
   'distance_metric': 'cosine',
   'index_name': 'autoai_rag_7878e0b8_20250627150244',
   'operation': 'upsert',
   'schema': {'fields': [{'description': 'text field',
      'name': 'text',
      'role': 'text',
      'type': 'string'},
     {'description': 'document name field',
      'name': 'document_id',
      'role': 'document_name',
      'type': 'string'},
     {'description': 'chunk starting token position in the source document',
      'name': 'start_index',
      'role': 'start_index',
      'type': 'number'},
     {'description': 'chunk number per document',
      'name': 'sequence_number',
      'role': 'sequence_number',
      'type': 'number'},
     {'description': 'vector embeddings',
      'name': 'vector',
      'role': 'vector_embeddings',
      'type': 'array'}],
    'id': 'autoai_rag_1.0',
    'name': 'Document schema using open-source loaders',
    'type': 'struct'}}},
 'settings_importance': {'chunking': [{'importance': 0.0,
    'parameter': 'chunk_size'},
   {'importance': 0.0, 'parameter': 'chunk_overlap'}],
  'embeddings': [{'importance': 0.0, 'parameter': 'embedding_model'}],
  'generation': [{'importance': 0.0, 'parameter': 'foundation_model'}],
  'retrieval': [{'importance': 1.0, 'parameter': 'retrieval_method'},
   {'importance': 0.0, 'parameter': 'window_size'},
   {'importance': 0.0, 'parameter': 'number_of_chunks'}]}}

Query generated pattern locally

from ibm_watsonx_ai.deployments import RuntimeContext

runtime_context = RuntimeContext(api_client=client)
inference_service_function = best_pattern.inference_service(runtime_context)[0]

/Users/michalsteczko/anaconda3/envs/autoai_rag/lib/python3.11/site-packages/pypdf/_crypt_providers/_cryptography.py:32: CryptographyDeprecationWarning: ARC4 has been moved to cryptography.hazmat.decrepit.ciphers.algorithms.ARC4 and will be removed from cryptography.hazmat.primitives.ciphers.algorithms in 48.0.0.
  from cryptography.hazmat.primitives.ciphers.algorithms import AES, ARC4

question = "What training objectives are used for the granite models?"

context = RuntimeContext(
    api_client=client,
    request_payload_json={"messages": [{"role": "user", "content": question}]},
)


resp = inference_service_function(context)
resp

{'body': {'choices': [{'index': 0,
    'message': {'role': 'assistant',
     'content': ' The\nclusters are equipped with 100Gbps and 200Gbps HDR InfiniBand links, respectively.\nWe utilize NVIDIA’s Megatron-LM (Shoeybi et al., 2019; Narayanan et al., 2021) for\ndistributed training, which is optimized for large language models. We use the same Megatron\nLM framework for all our models, ensuring consistency in training infrastructure.\n4.5 Model Architecture\nThe architecture of the Granite Code models is based on the original transformer architecture\n(Douglas & Smith, 2019) with modifications for code modeling. The base model has 16\nlayers, 8 attention heads, and a token embedding dimension of 512. For the 3B model, we\nuse the standard transformer architecture with a multi-head attention mechanism. The 8B model\nincorporates Grouped-Query Attention (GQA) (Ainslie et al., 2023) to improve inference\nefficiency. The 20B model uses learned absolute position embeddings and Multi-Query\nAttention (Shazeer, 2019). The 34B model is built upon the 20B model with depth\nupscaling (Kim et al., 2024) to double the model depth, resulting in 88 layers.\nThe models are trained with different context lengths depending on their size: 2048, 4096,\n8192, and 8192 tokens respectively for 3B, 8B, 20B, and 34B models.\nFor the MLP block, we use GELU activation function (Hendrycks & Gimpel, 2023) for the 20B\nand 34B models, while using GLU (Shazeer, 2020) for the 3B and 8B models. For\nnormalization, we use LayerNorm (Ba et al., 2016) for all models except the 8B model,\nwhich uses RMSNorm (Zhang & Sennrich, 2019) for computational efficiency.\n5.1 Evaluation Protocol\nWe evaluate the Granite Code models on a comprehensive set of benchmarks including\nHumanEvalPack (Muennighoff et al., 2023), MBPP(+) (Austin et al., 2021; Liu et al.,\n2023a), RepoBench (Liu et al., 2023b), ReCode (Wang et al., 2022), and more. This\nset of benchmarks encompasses many different kinds of coding tasks beyond just code\nsynthesis in Python, e.g., code fixing, code explanation, code editing, code translation,\netc., across most major programming languages (Python, JavaScript, Java, Go, C++, Rust,\netc.). Our evaluation protocol includes both automated metrics and human evaluations.\nFor automated metrics, we use standard evaluation scripts that measure code synthesis,\nfixing, and explanation. For human evaluations, we conduct a series of studies where\nparticipants are asked to compare the output of Granite Code models with other open-source\ncode models on various coding tasks. We also perform ablation studies to understand the\nimpact of different architectural choices on model performance.\n5.2 Results\nOur findings reveal that among open-source models, the Granite Code models overall show\nvery strong performance across all model sizes and benchmarks (often outperforming other\nopen-source code models that are twice large compared to Granite). As an illustration,\nfigure 1 (top) shows a comparison of Granite-8B-Code-Base with other open-source base code\nLLMs, including recent high-performing general purpose base LLMs like Mistral (Jiang et al.,\n2023b) and LLama-3 (AI@Meta, 2024) on HumanEvalPack (Muennighoff et al., 2023). While\nCodeGemma and StarCoder2 perform reasonably well in generating code, they perform\nsignificantly worse on the code fixing and explanation variants of HumanEvalPack. On av-\nerage, Granite-8B-Code-Base achieves 73% accuracy on code fixing, compared to 60% for\nCodeGemma and 55% for StarCoder2. Similarly, on code explanation, Granite-8B-\nCode-Base achieves 70% accuracy, outperforming CodeGemma (58%) and StarCoder2\n(50%). These results demonstrate that Granite Code models are not only capable of\ngenerating high-quality code but also excelling in code-related tasks that require reasoning\nand understanding, such as code fixing and explanation.\n\n6.3 Conclusion\nIn conclusion, we present Granite Code models, a series of highly capable code LLMs designed\nto support enterprise software development across a wide range of coding tasks. Our\nresults show that Granite Code models achieve state-of-the-art performance across a variety\nof'},
    'reference_documents': [{'page_content': 'code, while others focus primarily on coding-related tasks (e.g. StarCoder (Li et al., 2023a;\nLozhkov et al., 2024), CodeGen (Nijkamp et al., 2023), CodeLlama (Rozi `ere et al., 2023), and\nCodeGemma (CodeGemma Team et al., 2024)).\nHowever, there remain important gaps in the current field of LLMs for code, especially in\nthe context of enterprise software development. First, while very large, generalist LLMs can\nachieve excellent coding performance, their size makes them expensive to deploy. Smaller\ncode-focused models (Li et al., 2023a; Lozhkov et al., 2024; Nijkamp et al., 2023; Rozi `ere et al.,\n2023; CodeGemma Team et al., 2024) can achieve excellent code generation performance in\na smaller and more flexible package, but performance in coding tasks beyond generation\n(e.g. fixing and explanation) can lag behind code generation performance.\nIn many enterprise contexts, code LLM adoption can be further complicated by factors\nbeyond the performance of the models. For instance, even open models are sometimes\nplagued by a lack of transparency about the data sources and data processing methods\nthat went into model, which can erode trust in models in mission critical and regulated\ncontexts. Furthermore, license terms in today’s open LLMs can encumber and complicate\nan enterprise’s ability to use a model.\nHere, we present Granite Code models, a series of highly capable code LLMs, designed to\nsupport enterprise software development across a wide range of coding tasks. Granite Code\nmodels has two main variants that we release in four different sizes (3B, 8B, 20B, and 34B):\n2\nIBM Granite Code Models\n•Granite Code Base: base foundation models for code-related tasks;\n•Granite Code Instruct: instruction following models finetuned using a combination\nof Git commits paired with human instructions and open-source synthetically\ngenerated code instruction datasets.\nThe base models in the series have been trained from scratch with a two-phase training\nstrategy. In phase 1, our model is trained on 3 to 4 trillion tokens sourced from 116 pro-\ngramming languages, ensuring a comprehensive understanding of programming languages\nand syntax. In phase 2, our model is further trained on 500 billion tokens with a carefully\ndesigned mixture of high-quality data from code and natural language domains to improve\nthe model’s ability to reason. We use the unsupervised language modeling objective to\ntrain the base models in both the phases of training. The instruct models are derived by\nfurther finetuning the above trained base models on a combination of a filtered variant of\nCommitPack (Muennighoff et al., 2023), natural language instruction following datasets\n(OASST (K ¨opf et al., 2023), HelpSteer (Wang et al., 2023)) and open-source math datasets\n(MathInstruct (Yue et al., 2023) and MetaMathQA (Yu et al., 2023)), including synthetically\ngenerated code datasets for improving instruction following and reasoning capabilities.\nWe conduct extensive evaluations of our code LLMs on a comprehensive set of benchmarks,\nincluding HumanEvalPack (Muennighoff et al., 2023), MBPP(+) (Austin et al., 2021; Liu\net al., 2023a), RepoBench (Liu et al., 2023b), ReCode (Wang et al., 2022), and more. This set of\nbenchmarks encompasses many different kinds of coding tasks beyond just code synthesis\nin Python, e.g., code fixing, code explanation, code editing, code translation, etc., across\nmost major programming languages (Python, JavaScript, Java, Go, C++, Rust, etc.).\nOur findings reveal that among open-source models, the Granite Code models overall show\nvery strong performance across all model sizes and benchmarks (often outperforming other\nopen-source code models that are twice large compared to Granite). As an illustration, fig-\nure 1 (top) shows a comparison of Granite-8B-Code-Base with other open-source base code\nLLMs, including recent high-performing general purpose base LLMs like Mistral (Jiang et al.,\n2023b) and LLama-3 (AI@Meta, 2024) on HumanEvalPack (Muennighoff et al., 2023). While\nCodeGemma and StarCoder2 perform reasonably well in generating code, they perform\nsignificantly worse on the code fixing and explanation variants of HumanEvalPack. On av-',
      'metadata': {'sequence_number': [6, 7, 8, 9, 10],
       'document_id': 'granite_code_models.pdf'}},
     {'page_content': 'activation function (Ramachandran et al., 2017) with GLU (Shazeer, 2020) for the MLP , also\ncommonly referred to as swiglu. For normalization, we use RMSNorm (Zhang & Sennrich,\n2019) since it’s computationally more efficient than LayerNorm (Ba et al., 2016). The 3B\nmodel is trained with a context length of 2048 tokens.\n8B: The 8B model has a similar architecture as the 3B model with the exception of using\nGrouped-Query Attention (GQA) (Ainslie et al., 2023). Using GQA offers a better tradeoff\nbetween model performance and inference efficiency at this scale. We train the 8B model\nwith a context length of 4096 tokens.\n20B: The 20B code model is trained with learned absolute position embeddings. We use\nMulti-Query Attention (Shazeer, 2019) during training for efficient downstream inference.\nFor the MLP block, we use the GELU activation function (Hendrycks & Gimpel, 2023). For\nnormalizing the activations, we use LayerNorm (Ba et al., 2016). This model is trained with\na context length of 8192 tokens.\n34B: To train the 34B model, we follow the approach by Kim et al. for depth upscaling of\nthe 20B model. Specifically, we first duplicate the 20B code model with 52 layers and then\n5https://www.clamav.net/\n5\nIBM Granite Code Models\nFigure 2: An overview of depth upscaling (Kim et al., 2024) for efficient training of Granite-\n34B-Code. We utilize the 20B model after 1.6T tokens to start training of 34B model with the\nsame code pretraining data without any changes to the training and inference framework.\nremove final 8 layers from the original model and initial 8 layers from its duplicate to form\ntwo models. Finally, we concatenate both models to form Granite-34B-Code model with\n88 layers (see Figure 2 for an illustration). After the depth upscaling, we observe that the\ndrop in performance compared to 20B model is pretty small contrary to what is observed by\nKim et al.. This performance is recovered pretty quickly after we continue pretraining of the\nupscaled 34B model. Similar, to 20B, we use a 8192 token context during pretraining.\n4 Pretraining\nIn this section, we provide details on two phase training (Sec. 4.1), training objectives\n(Sec. 4.2), optimization (Sec. 4.3) and infrastructure (Sec. 4.4) used in pretraining the models.\n4.1 Two Phase Training\nGranite Code models are trained on 3.5T to 4.5T tokens of code data and natural language\ndatasets related to code. Data is tokenized via byte pair encoding (BPE, (Sennrich et al.,\n2015)), employing the same tokenizer as StarCoder (Li et al., 2023a). Following (Shen et al.,\n2024; Hu et al., 2024), we utilize high-quality data with two phases of training as follows.\n•Phase 1 (code only training) : During phase 1, both 3B and 8B models are trained for\n4 trillion tokens of code data comprising 116 languages. The 20B parameter model\nis trained on 3 trillion tokens of code. The 34B model is trained on 1.4T tokens after\nthe depth upscaling which is done on the 1.6T checkpoint of 20B model.\n•Phase 2 (code + language training) : In phase 2, we include additional high-quality\npublicly available data from various domains, including technical, mathematics,\nand web documents, to further improve the model’s performance in reasoning and\nproblem solving skills, which are essential for code generation. We train all our\nmodels for 500B tokens (80% code and 20% language data) in phase 2 training.\n4.2 Training Objective\nFor training of all our models, we use the causal language modeling objective and Fill-In-\nthe-Middle (FIM) (Bavarian et al., 2022) objective. The FIM objective is tasked to predict\ninserted tokens with the given context and subsequent text. We train our models to work\nwith both PSM (Prefix-Suffix-Middle) and SPM (Suffix-Prefix-Middle) modes, with relevant\nformatting control tokens, same as StarCoder (Li et al., 2023a).\nThe overall loss is computed as a weighted combination of the 2 objectives:\nL=αLCLM + (1−α)LFIM (1)\nWe emperically set α=0.5 during training and find that this works well in practice leading\nto SOTA performance on both code completion and code infilling tasks. It should be\n6\nIBM Granite Code Models',
      'metadata': {'sequence_number': [21, 22, 23, 24, 25],
       'document_id': 'granite_code_models.pdf'}},
     {'page_content': 'datasets related to code. Data is tokenized via byte pair encoding (BPE, (Sennrich et al.,\n2015)), employing the same tokenizer as StarCoder (Li et al., 2023a). Following (Shen et al.,\n2024; Hu et al., 2024), we utilize high-quality data with two phases of training as follows.\n•Phase 1 (code only training) : During phase 1, both 3B and 8B models are trained for\n4 trillion tokens of code data comprising 116 languages. The 20B parameter model\nis trained on 3 trillion tokens of code. The 34B model is trained on 1.4T tokens after\nthe depth upscaling which is done on the 1.6T checkpoint of 20B model.\n•Phase 2 (code + language training) : In phase 2, we include additional high-quality\npublicly available data from various domains, including technical, mathematics,\nand web documents, to further improve the model’s performance in reasoning and\nproblem solving skills, which are essential for code generation. We train all our\nmodels for 500B tokens (80% code and 20% language data) in phase 2 training.\n4.2 Training Objective\nFor training of all our models, we use the causal language modeling objective and Fill-In-\nthe-Middle (FIM) (Bavarian et al., 2022) objective. The FIM objective is tasked to predict\ninserted tokens with the given context and subsequent text. We train our models to work\nwith both PSM (Prefix-Suffix-Middle) and SPM (Suffix-Prefix-Middle) modes, with relevant\nformatting control tokens, same as StarCoder (Li et al., 2023a).\nThe overall loss is computed as a weighted combination of the 2 objectives:\nL=αLCLM + (1−α)LFIM (1)\nWe emperically set α=0.5 during training and find that this works well in practice leading\nto SOTA performance on both code completion and code infilling tasks. It should be\n6\nIBM Granite Code Models\nnoted that the FIM objective is only used during pretraining, however we drop it during\ninstruction finetuning i.e we set α=1.\n4.3 Optimization\nWe use AdamW optimizer (Kingma & Ba, 2017) with β1=0.9,β2=0.95 and weight decay\nof 0.1 for training all our Granite code models. For the phase-1 pretraining, the learning\nrate follows a cosine schedule starting from 3 ×10−4which decays to 3 ×10−5with an\ninitial linear warmup step of 2k iterations. For phase-2 pretraining, we start from 3 ×10−4\n(1.5×10−4for 20B and 34B models) and adopt an exponential decay schedule to anneal it\nto 10% of the initial learning rate. We use a batch size of 4M-5M tokens depending on the\nmodel size during both phases of pretraining.\nTo accelerate training, we use FlashAttention 2 (Dao et al., 2022; Dao, 2023), the persistent\nlayernorm kernel, Fused RMSNorm kernel (depending on the model) and the Fused Adam\nkernel available in NVIDIA’s Apex library. We use a custom fork of NVIDIA’s Megatron-\nLM (Shoeybi et al., 2019; Narayanan et al., 2021) for distributed training of all our models.\nWe train with a mix of 3D parallelism: tensor parallel, pipeline parallel and data parallel.\nWe also use sequence parallelism (Korthikanti et al., 2023) for reducing the activation\nmemory consumption of large context length during training. We use Megatron’s distributed\noptimizer with mixed precision training (Micikevicius et al., 2018) in BF16 (Kalamkar et al.,\n2019) with gradient all-reduce and gradient accumulation in FP32 for training stability.\n4.4 Infrastructure\nWe train the Granite Code models using IBM’s two supercomputing clusters, namely Vela\nand Blue Vela, outfitted with NVIDIA A100 and H100 GPUs, respectively. In the Vela\nA100 GPU cluster, each node has 2 ×Intel Xeon Scalable Processors with 8 ×80GB A100\nGPUs connected to each other by NVLink and NVSwitch. The Vela cluster adopts RoCE\n(RDMA over Converged Ethernet) and GDR (GPU-direct RDMA) for high-performance\nnetworking. Similarly, each node in Blue Vela cluster consists of dual 48-core Intel processors\nwith 8 ×80GB H100 GPUs. Blue Vela employs 3.2Tbps InfiniBand interconnect to facilitate\nseamless communication between nodes, known for their high throughput and low latency.',
      'metadata': {'sequence_number': [24, 25, 26, 27, 28],
       'document_id': 'granite_code_models.pdf'}}]}]}}

print(inference_service_function(context)["body"]["choices"][0]["message"]["content"])

 The
clusters are equipped with 100Gbps and 200Gbps HDR InfiniBand links, respectively.
We utilize NVIDIA’s Megatron-LM (Shoeybi et al., 2019; Narayanan et al., 2021) for
distributed training, which is optimized for large language models. We use the same Megatron
LM framework for all our models, ensuring consistency in training infrastructure.
4.5 Model Architecture
The architecture of the Granite Code models is based on the original transformer architecture
(Douglas & Smith, 2019) with modifications for code modeling. The base model has 16
layers, 8 attention heads, and a token embedding dimension of 512. For the 3B model, we
use the standard transformer architecture with a multi-head attention mechanism. The 8B model
incorporates Grouped-Query Attention (GQA) (Ainslie et al., 2023) to improve inference
efficiency. The 20B model uses learned absolute position embeddings and Multi-Query
Attention (Shazeer, 2019). The 34B model is built upon the 20B model with depth
upscaling (Kim et al., 2024) to double the model depth, resulting in 88 layers.
The models are trained with different context lengths depending on their size: 2048, 4096,
8192, and 8192 tokens respectively for 3B, 8B, 20B, and 34B models.
For the MLP block, we use GELU activation function (Hendrycks & Gimpel, 2023) for the 20B
and 34B models, while using GLU (Shazeer, 2020) for the 3B and 8B models. For
normalization, we use LayerNorm (Ba et al., 2016) for all models except the 8B model,
which uses RMSNorm (Zhang & Sennrich, 2019) for computational efficiency.
5.1 Evaluation Protocol
We evaluate the Granite Code models on a comprehensive set of benchmarks including
HumanEvalPack (Muennighoff et al., 2023), MBPP(+) (Austin et al., 2021; Liu et al.,
2023a), RepoBench (Liu et al., 2023b), ReCode (Wang et al., 2022), and more. This
set of benchmarks encompasses many different kinds of coding tasks beyond just code
synthesis in Python, e.g., code fixing, code explanation, code editing, code translation,
etc., across most major programming languages (Python, JavaScript, Java, Go, C++, Rust,
etc.). Our evaluation protocol includes both automated metrics and human evaluations.
For automated metrics, we use standard evaluation scripts that measure code synthesis,
fixing, and explanation. For human evaluations, we conduct a series of studies where
participants are asked to compare the output of Granite Code models with other open-source
code models on various coding tasks. We also perform ablation studies to understand the
impact of different architectural choices on model performance.
5.2 Results
Our findings reveal that among open-source models, the Granite Code models overall show
very strong performance across all model sizes and benchmarks (often outperforming other
open-source code models that are twice large compared to Granite). As an illustration,
figure 1 (top) shows a comparison of Granite-8B-Code-Base with other open-source base code
LLMs, including recent high-performing general purpose base LLMs like Mistral (Jiang et al.,
2023b) and LLama-3 (AI@Meta, 2024) on HumanEvalPack (Muennighoff et al., 2023). While
CodeGemma and StarCoder2 perform reasonably well in generating code, they perform
significantly worse on the code fixing and explanation variants of HumanEvalPack. On av-
erage, Granite-8B-Code-Base achieves 73% accuracy on code fixing, compared to 60% for
CodeGemma and 55% for StarCoder2. Similarly, on code explanation, Granite-8B-
Code-Base achieves 70% accuracy, outperforming CodeGemma (58%) and StarCoder2
(50%). These results demonstrate that Granite Code models are not only capable of
generating high-quality code but also excelling in code-related tasks that require reasoning
and understanding, such as code fixing and explanation.

6.3 Conclusion
In conclusion, we present Granite Code models, a series of highly capable code LLMs designed
to support enterprise software development across a wide range of coding tasks. Our
results show that Granite Code models achieve state-of-the-art performance across a variety
of

Summary

You successfully completed this notebook!

You learned how to use AutoAI RAG with your own foundation model.

Check out our Online Documentation for more samples, tutorials, documentation, how-tos, and blog posts.

Author:

Michał Steczko, Software Engineer at watsonx.ai.

Copyright © 2025 IBM. This notebook and its source code are released under the terms of the MIT License.