image

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

Contents¶

This notebook contains the following parts:

  • Set up the environment
  • Prepare data and connections for the Text Extraction service
  • Process data using the Text Extraction service
  • Prepare data and connections for the AutoAI RAG experiment
  • Run the AutoAI RAG experiment
  • Compare and test RAG Patterns
  • Summary

Set up the environment¶

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

  • Create a watsonx.ai Runtime Service instance (a free plan is offered and information about how to create the instance can be found here).

Install and import the required modules and dependencies¶

In [ ]:
!pip install -U 'ibm-watsonx-ai[rag]>=1.1.24' | tail -n 1
!pip install -U "langchain_community>=0.3,<0.4" | tail -n 1

Define the watsonx.ai credentials¶

This cell defines the credentials required to work with the watsonx.ai service.

Action: Provide your IBM Cloud API key and the platform URL. For more information, see Managing user API keys.

In [22]:
import getpass

from ibm_watsonx_ai import Credentials

credentials = Credentials(
    url="https://us-south.ml.cloud.ibm.com",
    api_key=getpass.getpass("Please enter your watsonx.ai api key (hit enter): ")
)

Work with spaces¶

You need to create a space that will be used for your work. If you do not have a space, you can use Deployment Spaces Dashboard to create one.

  • Click New Deployment Space
  • Create an empty space
  • Select Cloud Object Storage
  • Select watsonx.ai Runtime instance and press Create
  • Go to Manage tab
  • Copy Space GUID and paste it below

Tip: You can also use SDK to prepare the space for your work. For more information, see the Space management notebook.

Action: Assign space ID below.

In [24]:
SPACE_ID = 'PASTE YOUR SPACE GUID HERE'

Create an instance of APIClient with authentication details¶

In [25]:
from ibm_watsonx_ai import APIClient

client = APIClient(credentials=credentials, space_id=SPACE_ID)

Create an instance of COS client¶

Connect to the default COS instance for the provided space by using the ibm_boto3 package.

In [26]:
import ibm_boto3

cos_credentials = client.spaces.get_details(space_id=SPACE_ID)['entity']['storage']['properties']

cos_client = ibm_boto3.client(
    service_name="s3",
    endpoint_url=cos_credentials["endpoint_url"],
    aws_access_key_id=cos_credentials["credentials"]["editor"]["access_key_id"],
    aws_secret_access_key=cos_credentials["credentials"]["editor"]["secret_access_key"],
)

Create a new bucket.

In [27]:
cos_bucket_name = "autoai-rag-with-text-extraction-experiment"

buckets_names = [bucket["Name"] for bucket in cos_client.list_buckets()["Buckets"]]
if not cos_bucket_name in buckets_names:
    cos_client.create_bucket(Bucket=cos_bucket_name)

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

In [28]:
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": cos_credentials["credentials"]["editor"]["access_key_id"],
            "secret_key": cos_credentials["credentials"]["editor"]["secret_access_key"],
            "iam_url": client.service_instance._href_definitions.get_iam_token_url(),
            "url": cos_credentials["endpoint_url"],
        },
    }
)

cos_connection_id = client.connections.get_id(
    connection_details
)

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 [29]:
from ibm_watsonx_ai.helpers import DataConnection, S3Location

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

te_input_filename = "granite_code_models_paper.pdf"
te_result_filename = "granite_code_models_paper.md"

Download and upload training data to the COS bucket. Then define a connection to the uploaded file.

In [30]:
import wget

wget.download(data_url, te_input_filename)
cos_client.upload_file(te_input_filename, cos_bucket_name, te_input_filename)
Out[30]:
'granite_code_models_paper (1).pdf'

Input file connection.

In [31]:
input_data_reference = DataConnection(
    connection_asset_id=cos_connection_id,
    location=S3Location(bucket=cos_bucket_name, path=te_input_filename),
)
input_data_reference.set_client(client)

Output file connection.

In [32]:
result_data_reference = DataConnection(
    connection_asset_id=cos_connection_id, 
    location=S3Location(
        bucket=cos_bucket_name,
        path=te_result_filename
    )
)
result_data_reference.set_client(client)

Process data using the Text Extraction service¶

Initialize the Text Extraction service endpoint.

In [33]:
from ibm_watsonx_ai.foundation_models.extractions import TextExtractions

extraction = TextExtractions(
    credentials=credentials,
    space_id=SPACE_ID,
)

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

In [34]:
from ibm_watsonx_ai.metanames import TextExtractionsMetaNames

response = extraction.run_job(
    document_reference=input_data_reference,
    results_reference=result_data_reference,
    steps={
        TextExtractionsMetaNames.OCR: {
            "process_image": True,
            "languages_list": ["en"],
        },
        TextExtractionsMetaNames.TABLE_PROCESSING: {"enabled": True},
    },
    results_format="markdown",
)

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

Wait for the job to be complete.

In [35]:
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)

Job completed successfully, details: {
  "entity": {
    "document_reference": {
      "connection": {
        "id": "ba5a6803-15e9-4949-a88e-e4d5be84712d"
      },
      "location": {
        "bucket": "autoai-rag-with-text-extraction-experiment",
        "file_name": "granite_code_models_paper.pdf"
      },
      "type": "connection_asset"
    },
    "results": {
      "completed_at": "2025-01-09T15:18:43.682Z",
      "number_pages_processed": 28,
      "running_at": "2025-01-09T15:18:16.998Z",
      "status": "completed"
    },
    "results_reference": {
      "connection": {
        "id": "ba5a6803-15e9-4949-a88e-e4d5be84712d"
      },
      "location": {
        "bucket": "autoai-rag-with-text-extraction-experiment",
        "file_name": "granite_code_models_paper.md"
      },
      "type": "connection_asset"
    },
    "steps": {
      "ocr": {
        "languages_list": [
          "en"
        ]
      },
      "tables_processing": {
        "enabled": true
      }
    }
  },
  "metadata": {
    "created_at": "2025-01-09T15:18:15.034Z",
    "id": "e798ce5c-c1a8-48b9-b22d-e1a0b7be5e23",
    "modified_at": "2025-01-09T15:18:43.740Z",
    "space_id": "017a7971-1ddb-4dc0-a556-3f84229c4a07"
  }
}

Get the text extraction result.

In [36]:
from IPython.display import display, Markdown

cos_client.download_file(
    Bucket=cos_bucket_name,
    Key=te_result_filename,
    Filename=te_result_filename
)

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

†Corresponding Authors

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 abilities are available today. Models range in 43.5

  • 41.5

    • 37.4
    • 34.4
    • 31.2
    • 29.2

    - 29.0 * 33.2 * 18.9 + 19. - 21.3 - 15.1 * 15.0 - 29.6 - 26. - 14.2 O. - 7. * 8.9 * 2.9 + 20.

    • 15.0
    • 13.4
    • 13.5
    • 13.9
    • 12.4Generating Code

Explaining Code

Fixing Code

Average

O MistralAI/Mistral-7B W @ Meta/Llama-3-8B

@) Google/Gemma-7B

Meta/CodeLlama-7B @) Google/CodeGemma-7B @ IBM/Granite-8B-Code-Base W

@ BigCode/StarCoder2-7B

Apache 2.0 License

+ 49.6
  • 41.6
  • 36.9
  • 37.9
  • 19.8
  • 40.9
  • 38.4

29.1 |

  • 25.1

21

  • 41.9
  • 43.6
  • 5.8
  • 31.8
  • 25.3
  • 18.3
  • 13.7
  • 112.8

Generating Code

Explaining Code

Fixing Code

Average

O MistralAI/Mistral-7B-Instruct-v0.2

@ Meta/CodeLlama-7B-Instru

Prepare data and connections for the AutoAI RAG experiment¶

Upload a json file to use for benchmarking to COS and define a connection to this file.

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

In [37]:
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": [te_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": [te_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": [te_result_filename]
     },
]
In [38]:
import os

test_filename = "benchmark.json"

if not os.path.isfile(test_filename):
    with open(test_filename, "w") as json_file:
        json.dump(benchmarking_data, json_file, indent=4)

cos_client.upload_file(test_filename, cos_bucket_name, test_filename)

Test the data connection.

In [39]:
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)

test_data_references = [test_data_reference]

Use the reference to the Text Extraction job result as input for the AutoAI RAG experiment.

In [40]:
input_data_references = [result_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 [41]:
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=5,
    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 [51]:
rag_optimizer.run(
    input_data_references=input_data_references,
    test_data_references=test_data_references,
    background_mode=False
);


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

Running 'c8aa6019-c9c5-42da-a082-e2ea741e980e'

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


pending.....
running.........................
completed
Training of 'c8aa6019-c9c5-42da-a082-e2ea741e980e' 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 [52]:
summary = rag_optimizer.summary()
summary
Out[52]:
mean_answer_correctness mean_faithfulness mean_context_correctness chunking.chunk_size embeddings.model_id vector_store.distance_metric retrieval.method retrieval.number_of_chunks generation.model_id
Pattern_Name
Pattern2 0.7813 0.7882 1.0 1024 intfloat/multilingual-e5-large euclidean window 5 meta-llama/llama-3-70b-instruct
Pattern4 0.7566 0.8904 1.0 1024 intfloat/multilingual-e5-large cosine window 5 mistralai/mixtral-8x7b-instruct-v01
Pattern1 0.7207 0.7074 1.0 512 ibm/slate-125m-english-rtrvr euclidean window 5 mistralai/mixtral-8x7b-instruct-v01
Pattern5 0.7196 0.9690 1.0 1024 intfloat/multilingual-e5-large cosine window 3 ibm/granite-13b-chat-v2
Pattern3 0.6731 0.5211 1.0 1024 intfloat/multilingual-e5-large euclidean simple 5 meta-llama/llama-3-70b-instruct

Get the selected pattern¶

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

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

best_pattern = rag_optimizer.get_pattern()

Best pattern is: Pattern2
In [56]:
rag_optimizer.get_pattern_details(pattern_name=best_pattern_name)
Out[56]:
{'composition_steps': ['chunking',
  'embeddings',
  'vector_store',
  'retrieval',
  'generation'],
 'duration_seconds': 13,
 'location': {'evaluation_results': 'default_autoai_rag_out/c8aa6019-c9c5-42da-a082-e2ea741e980e/Pattern2/evaluation_results.json',
  'indexing_notebook': 'default_autoai_rag_out/c8aa6019-c9c5-42da-a082-e2ea741e980e/Pattern2/indexing_inference_notebook.ipynb',
  'inference_notebook': 'default_autoai_rag_out/c8aa6019-c9c5-42da-a082-e2ea741e980e/Pattern2/indexing_inference_notebook.ipynb'},
 'name': 'Pattern2',
 'settings': {'chunking': {'chunk_overlap': 256,
   'chunk_size': 1024,
   'method': 'recursive'},
  'embeddings': {'model_id': 'intfloat/multilingual-e5-large',
   'truncate_input_tokens': 512,
   'truncate_strategy': 'left'},
  'generation': {'context_template_text': '{document}',
   'model_id': 'meta-llama/llama-3-70b-instruct',
   'parameters': {'decoding_method': 'greedy',
    'max_new_tokens': 1000,
    'min_new_tokens': 1},
   'prompt_template_text': '<|begin_of_text|><|start_header_id|>system<|end_header_id|>\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don’t know the answer to a question, please don’t share false information.\n<|eot_id|><|start_header_id|>user<|end_header_id|>\n{reference_documents}\n[conversation]: {question}. Answer with no more than 150 words.  If you cannot base your answer on the given document, please state that you do not have an answer.<|eot_id|>\n<|start_header_id|>assistant<|end_header_id|>'},
  'retrieval': {'method': 'window', 'number_of_chunks': 5, 'window_size': 2},
  'vector_store': {'datasource_type': 'chroma',
   'distance_metric': 'euclidean',
   'index_name': 'autoai_rag_c8aa6019_20250109154049',
   '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'}}}}

Test the RAGPattern by querying it locally.

In [57]:
questions = ["Which industry players are mentioned as IBM’s strategic partners?"]

payload = {
    client.deployments.ScoringMetaNames.INPUT_DATA: [
        {
            "values": questions,
            "access_token": client.service_instance._get_token()
        }
    ]
}

resp = best_pattern.inference_function()(payload)
In [58]:
print(resp["predictions"][0]["values"][0][0])


Based on the available information, IBM's strategic partners mentioned in the industry include:

* Salesforce: A leading customer relationship management (CRM) platform provider, with which IBM has a global strategic partnership to deliver joint solutions for artificial intelligence, blockchain, and the Internet of Things (IoT).
* Apple: A technology giant with which IBM has a partnership to develop enterprise mobility solutions, including iOS apps and IBM cloud services.
* Red Hat: An open-source software provider acquired by IBM, which has enabled the company to expand its cloud capabilities and offer a hybrid cloud platform.
* Box: A cloud content management platform provider with which IBM has a partnership to integrate its cloud storage and content management capabilities.

Please note that this information may not be exhaustive, and there might be other strategic partners not mentioned here.

Deploy the RAGPattern¶

Store the defined RAG function and create a deployed asset to deploy the RAGPattern.

In [59]:
deployment_details = best_pattern.deploy(
    name="AutoAI RAG deployment - ibm_watsonx_ai documentataion",
    space_id=SPACE_ID
)


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

Synchronous deployment creation for id: '9a0c710f-5e00-44b5-a22d-eec84e42fac9' 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='7c9a9ec8-c0f9-43e2-b2e0-7cac7eccc8ca'
-----------------------------------------------------------------------------------------------

Test the deployed function¶

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

In [60]:
deployment_id = client.deployments.get_id(deployment_details)

score_response = client.deployments.score(deployment_id, payload)
score_response
Out[60]:
{'predictions': [{'fields': ['answer', 'reference_documents'],
   'values': [["\n\nBased on the available information, IBM's strategic partners mentioned in the industry include:\n\n* Salesforce: A leading customer relationship management (CRM) platform provider, with which IBM has a global strategic partnership to deliver joint solutions for artificial intelligence, blockchain, and the Internet of Things (IoT).\n* Apple: A technology giant with which IBM has a partnership to develop enterprise mobility solutions, including iOS apps and IBM cloud services.\n* Red Hat: An open-source software provider acquired by IBM, which has enabled the company to expand its cloud capabilities and offer a hybrid cloud platform.\n* Box: A cloud content management platform provider with which IBM has a partnership to integrate its cloud storage and content management capabilities.\n\nPlease note that this information may not be exhaustive, and there might be other strategic partners not mentioned here.",
     []]]}]}
In [61]:
print(score_response["predictions"][0]["values"][0][0])


Based on the available information, IBM's strategic partners mentioned in the industry include:

* Salesforce: A leading customer relationship management (CRM) platform provider, with which IBM has a global strategic partnership to deliver joint solutions for artificial intelligence, blockchain, and the Internet of Things (IoT).
* Apple: A technology giant with which IBM has a partnership to develop enterprise mobility solutions, including iOS apps and IBM cloud services.
* Red Hat: An open-source software provider acquired by IBM, which has enabled the company to expand its cloud capabilities and offer a hybrid cloud platform.
* Box: A cloud content management platform provider with which IBM has a partnership to integrate its cloud storage and content management capabilities.

Please note that this information may not be exhaustive, and there might be other strategic partners not mentioned here.

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.

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