GitHub Repository: ibm/watson-machine-learning-samples
Path: blob/master/cpd5.2/notebooks/python_sdk/experiments/autoai_rag/Use AutoAI RAG with watsonx Text Extraction service.ipynb
⁶⁴¹² views

Kernel: Python 3.11

Use AutoAI RAG with watsonx Text Extraction service

Disclaimers

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

Notebook content

This notebook demonstrates how to process data using the IBM watsonx.ai Text Extraction service and use the result in an 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.11.

Learning goal

The learning goals of this notebook are:

Process data using the IBM watsonx.ai Text Extraction service
Create an AutoAI RAG job that will find the best RAG pattern based on processed data

This notebook contains the following parts:

Set up the environment

Before you use the sample code in this notebook, you must perform the following setup tasks:

Contact your Cloud Pak for Data administrator and ask them for your account credentials

Install and import the required modules and dependencies

In [1]:

!pip install -U 'ibm-watsonx-ai[rag]>=1.3.13' | tail -n 1

Connect to WML

Authenticate the Watson Machine Learning service on IBM Cloud Pak for Data. You need to provide the platform url, your username, and your api_key.

url - url which points to your CPD instance.
username - username to your CPD instance.
api_key - api_key to your CPD instance.

In [2]:

url = "PASTE YOUR CPD INSTANCE URL HERE"
api_key = "PASTE YOUR CPD INSTANCE API KEY HERE"
username = "PASTE YOUR CPD INSTANCE USERNAME HERE"

In [3]:

from ibm_watsonx_ai import Credentials

credentials = Credentials(
    username=username,
    api_key=api_key,
    url=url,
    instance_id="openshift",
    version="5.2"
)

Alternatively, you can use your username and password to authenticate WML services.

credentials = Credentials(
    username=***,
    password=***,
    url=***,
    instance_id="openshift",
    version="5.2"
)

Create an instance of APIClient with authentication details

In [4]:

from ibm_watsonx_ai import APIClient

client = APIClient(credentials)

Working with spaces

First, you need to create a space for your work. If you do not have a space already created, you can use {PLATFORM_URL}/ml-runtime/spaces?context=icp4data to create one.

Click New Deployment Space
Create an empty space
Go to the space Settings tab
Copy the space_id and paste it below

PLATFORM_URL is the url which points to your CPD instance.

Tip: You can also use SDK to prepare the space for your work. Find more information in the Space Management sample notebook.

Action: Assign the space ID below

In [5]:

space_id = 'PASTE YOUR SPACE ID HERE'

To be able to interact with all resources available in Watson Machine Learning, set the space that you are using.

In [6]:

client.set.default_space(space_id)

Out[6]:

'SUCCESS'

Create an instance of COS client

Connect to the Cloud Object Storage instance for the by using the ibm_boto3 package (for detailed explanation how to do this see IBM Cloud Object Storage connection).

Action: Assign COS credentials below

In [7]:

cos_bucket_name = "PASTE YOUR COS BUCKET NAME HERE"
endpoint_url = "PASTE YOUR COS BUCKET ENDPOINT HERE"
access_key = "PASTE YOUR COS BUCKET ACCESS KEY HERE"
secret_access_key = "PASTE YOU COS BUCKET SECRET ACCESS KEY HERE"

Initialize the connection to the created and get the connection ID.

In [8]:

connection_details = client.connections.create(
    {
        "datasource_type": client.connections.get_datasource_type_uid_by_name(
            "bluemixcloudobjectstorage"
        ),
        "name": "Connection to COS for tests",
        "properties": {
            "bucket": cos_bucket_name,
            "access_key": access_key,
            "secret_key": secret_access_key,
            "iam_url": client.service_instance._href_definitions.get_iam_token_url(),
            "url": endpoint_url,
        },
    }
)

cos_connection_id = client.connections.get_id(
    connection_details
)

Out[8]:

Creating connections...
SUCCESS

Prepare data and connections for the Text Extraction service

The document, from which we are going to extract text, is located in the IBM Cloud Object Storage (COS). In this notebook, we will use the Granite Code Models paper as a source text document. The final results file, which will contain extracted text and necessary metadata, will be placed in the COS. So we will use the ibm_watsonx_ai.helpers.DataConnection and the ibm_watsonx_ai.helpers.S3Location class to create Python objects that will represent the references to the processed files. Reference to the final results will be used as an input for the AutoAI RAG experiment.

In [9]:

from ibm_watsonx_ai.helpers import DataConnection, S3Location

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

text_extraction_input_filename = "granite_code_models.pdf"
text_extraction_result_filename = "granite_code_models.md"

Create an input connection.

In [10]:

input_data_reference = DataConnection(
    connection_asset_id=cos_connection_id,
    location=S3Location(
        bucket=cos_bucket_name,
        path=text_extraction_input_filename,
    )
)
input_data_reference.set_client(client)

Download the document from the url and upload to the COS Bucket using created connection.

In [11]:

import requests
import tempfile

response = requests.get(data_url)
with tempfile.NamedTemporaryFile() as tmp:
    tmp.write(response.content)
    input_data_reference.write(tmp.name)

Output file connection.

In [12]:

result_data_reference = DataConnection(
    connection_asset_id=cos_connection_id,
    location=S3Location(
        bucket=cos_bucket_name,
        path=text_extraction_result_filename
    )
)
result_data_reference.set_client(client)

Process data using the Text Extraction service

Initialize the Text Extraction service endpoint.

In [13]:

from ibm_watsonx_ai.foundation_models.extractions import TextExtractionsV2

extraction = TextExtractionsV2(api_client=client)

Run a text extraction job for connections created in the previous step.

In [14]:

from ibm_watsonx_ai.metanames import TextExtractionsV2ParametersMetaNames
from ibm_watsonx_ai.foundation_models.extractions import TextExtractionsV2ResultFormats

response = extraction.run_job(
    document_reference=input_data_reference,
    results_reference=result_data_reference,
    parameters={
        TextExtractionsV2ParametersMetaNames.OCR_MODE: "enabled",
        TextExtractionsV2ParametersMetaNames.LANGUAGES: ["en"]
    },
    result_formats=TextExtractionsV2ResultFormats.MARKDOWN,
)

job_id = response['metadata']['id']

Wait for the job to be complete.

In [15]:

import json
import time

while True:
    job_details = extraction.get_job_details(job_id)
    status = job_details['entity']['results']['status']
    
    if status == "completed":
        print("Job completed successfully, details: {}".format(json.dumps(job_details, indent=2)))
        break
    
    if status == "failed":
        print("Job failed, details: {}. \n Try to run job again.".format(json.dumps(job_details, indent=2)))
        break

    time.sleep(10)

Out[15]:

Job completed successfully, details: {
  "entity": {
    "document_reference": {
      "connection": {
        "id": "9a6d9083-0657-41ec-a384-17327a90c3dd"
      },
      "location": {
        "bucket": "cf809fbf-8020-4c69-8435-7220aa26c522",
        "file_name": "granite_code_models.pdf"
      },
      "type": "connection_asset"
    },
    "parameters": {
      "create_embedded_images": "disabled",
      "languages": [
        "en"
      ],
      "mode": "standard",
      "ocr_mode": "enabled",
      "output_dpi": 72,
      "output_tokens_and_bbox": true,
      "requested_outputs": [
        "md"
      ]
    },
    "results": {
      "completed_at": "2025-05-15T10:15:35.666Z",
      "number_pages_processed": 28,
      "running_at": "2025-05-15T10:14:48.293Z",
      "status": "completed"
    },
    "results_reference": {
      "connection": {
        "id": "9a6d9083-0657-41ec-a384-17327a90c3dd"
      },
      "location": {
        "bucket": "cf809fbf-8020-4c69-8435-7220aa26c522",
        "file_name": "granite_code_models.md"
      },
      "type": "connection_asset"
    }
  },
  "metadata": {
    "created_at": "2025-05-15T10:14:45.638Z",
    "id": "d465b6b7-558e-45fc-88e7-018c7b446961",
    "modified_at": "2025-05-15T10:15:35.879Z",
    "space_id": "6f0b87a7-ae96-4467-b023-7fb5767d4798"
  }
}

Get the text extraction result.

In [16]:

from IPython.display import display, Markdown

result_data_reference.download(filename=text_extraction_result_filename)
with open(text_extraction_result_filename, 'r', encoding='utf-8') as file:
    # Display beginning of the result file
    display(Markdown((file.read()[:3000])))

Out[16]:

Granite Code Models: A Family of Open Foundation Models for Code Intelligence

Mayank Mishra⋆ Matt Stallone⋆ Gaoyuan Zhang⋆ Yikang Shen Aditya Prasad Adriana Meza Soria Michele Merler Parameswaran Selvam Saptha Surendran Shivdeep Singh Manish Sethi Xuan-Hong Dang Pengyuan Li Kun-Lung Wu Syed Zawad Andrew Coleman Matthew White Mark Lewis Raju Pavuluri Yan Koyfman Boris Lublinsky Maximilien de Bayser Ibrahim Abdelaziz Kinjal Basu Mayank Agarwal Yi Zhou Chris Johnson Aanchal Goyal Hima Patel Yousaf Shah Petros Zerfos Heiko Ludwig Asim Munawar Maxwell Crouse Pavan Kapanipathi Shweta Salaria Bob Calio Sophia Wen Seetharami Seelam Brian Belgodere Carlos Fonseca Amith Singhee Nirmit Desai David D. Cox Ruchir Puri† Rameswar Panda†

IBM Research

⋆Equal

Contribution

†Corresponding Authors [email protected], [email protected]

Abstract

Large Language Models (LLMs) trained on code are revolutionizing the software development process. Increasingly, code LLMs are being inte grated into software development environments to improve the produc tivity of human programmers, and LLM-based agents are beginning to show promise for handling complex tasks autonomously. Realizing the full potential of code LLMs requires a wide range of capabilities, including code generation, fixing bugs, explaining and documenting code, maintaining repositories, and more. In this work, we introduce the Granite series of decoder-only code models for code generative tasks, trained with code written in 116 programming languages. The Granite Code models family consists of models ranging in size from 3 to 34 billion parameters, suitable for applications ranging from complex application modernization tasks to on-device memory-constrained use cases. Evaluation on a comprehensive set of tasks demonstrates that Granite Code models consistently reaches state-of-the-art performance among available open-source code LLMs. The Granite Code model family was optimized for enterprise software devel opment workflows and performs well across a range of coding tasks (e.g. code generation, fixing and explanation), making it a versatile “all around” code model. We release all our Granite Code models under an Apache 2.0 license for both research and commercial use.

https://github.com/ibm-granite/granite-code-models

1 Introduction

Over the last several decades, software has been woven into the fabric of every aspect of our society. As demand for software development surges, it is more critical than ever to increase software development productivity, and LLMs provide promising path for augmenting human programmers. Prominent enterprise use cases for LLMs in software development productivity include code generation, code explanation, code fixing, unit test and documentation generation, application modernization, vulnerability detection, code translation, and more.

Recent years have seen rapid progress in LLM’s ability to generate and manipulate code, and a range of models with impressive coding a

Prepare data and connections for the AutoAI RAG experiment

Define a connection to COS Bucket and upload a json file to use for benchmarking.

Note: correct_answer_document_ids must refer to the document processed by text extraction service, not the initial document.

Benchmarking data.

In [17]:

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": [text_extraction_result_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": [text_extraction_result_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": [text_extraction_result_filename]
 },
]

Create a connection.

In [18]:

test_filename = "granite_code_models_benchmark.json"

test_data_reference = DataConnection(
    connection_asset_id=cos_connection_id,
    location=S3Location(bucket=cos_bucket_name, path=test_filename),
)
test_data_reference.set_client(client)

Upload benchmarking data to COS Bucket.

In [19]:

with tempfile.NamedTemporaryFile(mode="w") as tmp:
    json.dump(benchmarking_data, tmp, indent=4)
    tmp.flush()
    tmp_path = tmp.name
    test_data_reference.write(tmp_path)

Define the connections to the test data and input data for the AutoAI RAG experiment, using the output of the Text Extraction job as the input connection.

In [20]:

input_data_references = [result_data_reference]
test_data_references = [test_data_reference]

Run the AutoAI RAG experiment

Provide the input information for AutoAI RAG optimizer:

name - experiment name
description - experiment description
max_number_of_rag_patterns - maximum number of RAG patterns to create
optimization_metrics - target optimization metrics

In [21]:

from ibm_watsonx_ai.experiment import AutoAI

experiment = AutoAI(credentials, space_id=space_id)

rag_optimizer = experiment.rag_optimizer(
    name='AutoAI RAG - Text Extraction service experiment',
    description = "AutoAI RAG experiment on documents generated by text extraction service",
    max_number_of_rag_patterns=4,
    optimization_metrics=['answer_correctness']
)

Call the run() method to trigger the AutoAI RAG experiment. Choose one of two modes:

To use the interactive mode (synchronous job), specify background_mode=False
To use the background mode (asynchronous job), specify background_mode=True

In [22]:

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

Out[22]:

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

Running '89fff5da-8405-4aba-abd6-b21b66f8cd03'

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


pending...
running........................................................................................................................................................................................................................................................
completed
Training of '89fff5da-8405-4aba-abd6-b21b66f8cd03' finished successfully.

Compare and test of RAG Patterns

You can list the trained patterns and information on evaluation metrics in the form of a Pandas DataFrame by calling the summary() method. You can use the DataFrame to compare all discovered patterns and select the one you like for further testing.

In [23]:

summary = rag_optimizer.summary()
summary

Out[23]:

Get the selected pattern

Get the RAGPattern object from the RAG Optimizer experiment. By default, the RAGPattern of the best pattern is returned.

In [24]:

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

best_pattern = rag_optimizer.get_pattern()

Out[24]:

Best pattern is: Pattern1

Test the RAGPattern by querying it locally.

In [25]:

from ibm_watsonx_ai.deployments import RuntimeContext

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

In [37]:

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

context = RuntimeContext(
    api_client=client,
    request_payload_json={"messages": [{"role": "user", "content": question}]},
)
print(inference_service_function(context)["body"]["choices"][0]["message"]["content"])

Out[37]:

The models are trained using the causal language modeling objective and the Fill-In the-Middle (FIM) objective. The FIM objective is designed to predict inserted tokens given a context and subsequent text. The models are trained to work with both PSM (Prefix-Suffix-Middle) and SPM (Suffix-Prefix-Middle) modes, using relevant formatting control tokens. The overall loss is computed as a weighted combination of the two objectives, with α empirically set to 0.5 during training. The FIM objective is only used during pretraining, and is dropped during instruction finetuning.

Reference(s):
Document
4.2 Training Objective
For training of all our models, we use the causal language modeling objective and Fill-In the-Middle (FIM) (Bavarian et al., 2022) objective. The FIM objective is tasked to predict inserted tokens with the given context and subsequent text. We train our models to work with both PSM (Prefix-Suffix-Middle) and SPM (Suffix-Prefix-Middle) modes, with relevant formatting control tokens, same as StarCoder (Li et al., 2023a). The overall loss is computed as a weighted combination of the 2 objectives:
L = αLCLM + (1 − α)LF IM (1)
We emperically set α = 0.5 during training and find that this works well in practice leading to SOTA performance on both code completion and code infilling tasks. It should be noted that the FIM objective is only used during pretraining, however we drop it during instruction finetuning i.e we set α = 1.

Is there anything else you would like to know?

Deploy the RAGPattern

To deploy the RAGPattern, store the defined RAG function and then create a deployed asset.

In [27]:

deployment_details = best_pattern.inference_service.deploy(
    name="AutoAI RAG with Text Extraction service - test deployment",
    space_id=space_id,
    deploy_params={"tags": ["wx-autoai-rag"]}
)

Out[27]:

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

Synchronous deployment creation for id: '5f644c35-e6cf-43cb-ad04-7b0b13e67601' 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='8cabd997-4968-4675-a3e0-a48c4b6df6fb'
-----------------------------------------------------------------------------------------------

Test the deployed function

The RAG service is now deployed in our space. To test the solution, run the cell below. Questions have to be provided in the payload. Their format is provided below.

In [28]:

deployment_id = client.deployments.get_id(deployment_details)

payload = {
    "messages": [{"role": "user", "content": question}]
}
score_response = client.deployments.run_ai_service(deployment_id, payload)

In [29]:

from langchain_core.documents import Document
from IPython.display import display, Markdown
from ibm_watsonx_ai.foundation_models.extensions.rag.utils import verbose_search

reference_docs = score_response["choices"][0]["reference_documents"]
answer = score_response["choices"][0]["message"]["content"]
verbose_search(question, [Document(**d) for d in reference_docs])
display(Markdown(f"**Answer:** {answer}"))

Out[29]:

Question: What training objectives are used for the models?

Answer: The models are trained using the causal language modeling objective and the Fill-In the-Middle (FIM) objective. The FIM objective is designed to predict inserted tokens given a context and subsequent text. The models are trained to work with both PSM (Prefix-Suffix-Middle) and SPM (Suffix-Prefix-Middle) modes, using relevant formatting control tokens. The overall loss is computed as a weighted combination of the two objectives, with α empirically set to 0.5 during training. The FIM objective is only used during pretraining, and is dropped during instruction finetuning.

Reference(s): Document 4.2 Training Objective For training of all our models, we use the causal language modeling objective and Fill-In the-Middle (FIM) (Bavarian et al., 2022) objective. The FIM objective is tasked to predict inserted tokens with the given context and subsequent text. We train our models to work with both PSM (Prefix-Suffix-Middle) and SPM (Suffix-Prefix-Middle) modes, with relevant formatting control tokens, same as StarCoder (Li et al., 2023a). The overall loss is computed as a weighted combination of the 2 objectives: L = αLCLM + (1 − α)LF IM (1) We emperically set α = 0.5 during training and find that this works well in practice leading to SOTA performance on both code completion and code infilling tasks. It should be noted that the FIM objective is only used during pretraining, however we drop it during instruction finetuning i.e we set α = 1.

Is there anything else you would like to know?

Summary

You successfully completed this notebook!

You learned how to use AutoAI RAG with documents processed by the TextExtraction service.

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

Author:

Witold Nowogórski, Software Engineer at watsonx.ai.