This notebook contains the steps and code to demonstrate the retrieval-augmented generation pattern in IBM watsonx.ai, using IBM Watson Discovery as the Search component.
Some familiarity with Python is helpful. This notebook uses Python 3.10.
The goal of this notebook is to demonstrate how to apply the retrieval-augmented generation pattern to a question-answering use case in watsonx.ai.
The website for an online seed catalog has many articles to help customers plan their garden and ultimately select which seeds to purchase. A new widget is being added to the website to answer customer questions based on the contents of the article the customer is viewing. Given a question related to a given article, answer the question based on the article.
This notebook contains the following parts:
The retrieval-augmented generation pattern involves three basic steps:
The term retrieval-augmented generation (RAG) was introduced in this paper: Retrieval-augmented generation for knowledge-intensive NLP tasks
"We build RAG models where the parametric memory is a pre-trained seq2seq transformer, and the non-parametric memory is a dense vector index of Wikipedia, accessed with a pre-trained neural retriever."
In that paper, the term "RAG models" refers to a specific implementation of a retriever (a specific query encoder and vector-based document search index) and a generator (a specific pre-trained, generative language model.)
However, the basic search-and-generate approach can be generalized to use different retriever components and foundation models.
In this notebook:
Before you use the sample code in this notebook, you must perform setup tasks.
The current project is the project in which you are running this notebook.
If an instance of Watson Machine Learning is not already associated with the current project, follow the instructions in this topic to do so: Adding associated services to a project.
Create an IBM Cloud API key by following these instructions: Creating an IBM Cloud API key
Then paste your new IBM Cloud API key in the code cell below.
g_cloud_apikey = ""
g_credentials = {
"url" : "https://us-south.ml.cloud.ibm.com",
"apikey" : g_cloud_apikey
}
You must perform the following setup tasks to use Watson Discovery as your Search component:
Create an instance of the IBM Watson Discovery service.
From the Manage page of your Discovery service instance in IBM Cloud, copy the API key and URL into the cell below.
g_discovery_apikey = ""
g_discovery_url = ""
!pip install ibm_watson
from ibm_watson import DiscoveryV2
from ibm_cloud_sdk_core.authenticators import IAMAuthenticator
authenticator = IAMAuthenticator( g_discovery_apikey )
g_discovery = DiscoveryV2( version= "2023-03-31", authenticator=authenticator )
g_discovery.set_service_url( g_discovery_url )
In Discovery, you organize your work in "Projects".
Create a Project in Discovery called "RAG with Discovery Search Project"
Note: You can also create a Project using the Discovery interface instead.
project_name = "RAG Search Project"
project_creation_result = g_discovery.create_project( name=project_name, type="document_retrieval" ).get_result()
g_discovery_project_id = project_creation_result["project_id"]
print( "Discovery project ID:\n" + g_discovery_project_id )
In Discovery, you assemble documents to search in "Collections".
Create a Collection in Discovery called "Discovery Search Collection"
Note: You can also create a Collection using the Discovery interface instead.
collection_name = "Gardening Articles Collection"
collection_creation_result = g_discovery.create_collection(project_id=g_discovery_project_id, name=collection_name, language="en" ).get_result()
g_collection_id = collection_creation_result["collection_id"]
print( "Discovery collection ID:\n" + g_collection_id )
In this notebook, the knowledge base is a collection of two articles.
(These articles were written as samples for watsonx.ai, they are not real articles published anywhere else. The authors and publication dates are fictional.)
article_01 = \
"Tomatoes are one of the most popular plants for vegetable gardens. Tip for success: If you select " \
"varieties that are resistant to disease and pests, growing tomatoes can be quite easy. For " \
"experienced gardeners looking for a challenge, there are endless heirloom and specialty varieties " \
"to cultivate. Tomato plants come in a range of sizes. There are varieties that stay very small, " \
"less than 12 inches, and grow well in a pot or hanging basket on a balcony or patio. Some grow " \
"into bushes that are a few feet high and wide, and can be grown is larger containers. Other " \
"varieties grow into huge bushes that are several feet wide and high in a planter or garden bed. " \
"Still other varieties grow as long vines, six feet or more, and love to climb trellises. Tomato " \
"plants do best in full sun. You need to water tomatoes deeply and often. Using mulch prevents " \
"soil-borne disease from splashing up onto the fruit when you water. Pruning suckers and even " \
"pinching the tips will encourage the plant to put all its energy into producing fruit."
article_02 = \
"Cucumbers are fun to grow for beginning gardeners and advanced gardeners alike. There are two " \
"types of cucumbers: slicing and pickling. Pickling cucumbers are smaller than slicing cucumbers. " \
"Cucumber plants come in two types: vining cucumbers, which are more common, and bush cucumbers. " \
"Vining cucumbers, which can grow to more than 5 feet tall, grow fast, yield lots of fruit, and you " \
"can train them up a trellis. Growing cucumbers up a trellis or fence can maximize garden space, " \
"keep fruit clean, and make it easier to harvest the fruit. Tropical plants, cucumbers are very " \
"sensitive to frost or cold weather. Cucumbers prefer full sun for 6 to 8 hours per day. Cucumbers " \
"need constant watering. Cucumbers can grow quickly and ripen in just 6 weeks. Harvest cucumbers " \
"every day or two because the more you harvest, the more the plant will produce. If any cucumber " \
"is left on the vine to fully mature, the plant will stop producing more cucumbers. You can extend " \
"the harvest season by planting cucumbers in batches, 2 weeks apart."
Create text files in the notebook working directory for each article:
with open ( "article_01.txt", "w" ) as file:
file.write( article_01 )
with open ( "article_02.txt", "w" ) as file:
file.write( article_02 )
See: Add a document in the Discovery V2 API.
knowledge_base = [
{
"file_name" : "article_01.txt",
"metadata" : { "title" : "Growing tomatoes",
"author" : "A. Rossi",
"published" : "2010" }
},
{
"file_name" : "article_02.txt",
"metadata" : { "title" : "Cucumbers for beginners",
"author" : "B. Melnyk",
"published" : "2018" }
}
]
import json
for article in knowledge_base:
file_name = article["file_name"]
metadata = article["metadata"]
with open( file_name, "rb" ) as f:
response_json = g_discovery.add_document( project_id=g_discovery_project_id,
collection_id=g_collection_id,
file=f,
filename=file_name,
metadata=json.dumps( metadata ),
file_content_type="text/plain" ).get_result()
print( file_name + "\n" + json.dumps( response_json, indent=3 ) + "\n" )
It takes a Discovery a few minutes to process the uploaded files before you can perform searches.
In the Discovery graphical interface, you can monitor the status of the file processing to see when the uploaded documents are ready for search.
Many articles that discuss retrieval-augmented generation assume the retrieval component uses a vector database.
However, to perform the general retrieval-augmented generation pattern, any search-and-retrieve method that can reliably return relevant content from the knowledge base will do.
In this notebook, the search component is a Watson Discovery search that returns one or the other of the two articles in the knowledge base, based on a natural language query match.
Define the natural language search in Discovery:
def search( question ):
response_json = g_discovery.query( project_id=g_discovery_project_id,
natural_language_query=question
).get_result()
#print( json.dumps( response_json, indent=3 ) )
results_arr = response_json["results"]
if( len( results_arr ) < 1 ):
return None
top_result = results_arr[0]
top_asset = { "title" : top_result["metadata"]["title"],
"author" : top_result["metadata"]["author"],
"published" : top_result["metadata"]["published"],
"text" : top_result["text"][0] }
return top_asset
search( "How tall do cucumbers grow?")
In this notebook, the task to be performed is a question-answering task.
There is no one, best prompt for any given task. However, models that have been instruction-tuned, such as bigscience/mt0-xxl-13b, google/flan-t5-xxl-11b, or google/flan-ul2-20b, can generally perform this task with the sample prompt below. Conservative decoding methods tend towards succinct answers.
In the prompt below, notice two string placeholders (marked with %s) that will be replaced at generation time:
prompt_template = """
Article:
###
%s
###
Answer the following question using only information from the article.
Answer in a complete sentence, with proper capitalization and punctuation.
If there is no good answer in the article, say "I don't know".
Question: %s
Answer:
"""
def augment( template_in, context_in, nlquery_in ):
return template_in % ( context_in, nlquery_in )
question = "How tall do cucumber plants grow?"
article_txt = article_02
augmented_prompt = augment( prompt_template, article_txt, question )
print( augmented_prompt )
You can prompt foundation models in watsonx.ai programmatically using the Python library.
See:
import os
from ibm_watson_machine_learning.foundation_models import Model
model_id = "google/flan-t5-xxl"
gen_parms = {
"DECODING_METHOD" : "greedy",
"MIN_NEW_TOKENS" : 1,
"MAX_NEW_TOKENS" : 50
}
g_wml_project_id = os.environ["PROJECT_ID"]
model = Model( model_id, g_credentials, gen_parms, g_wml_project_id )
def generate( model_in, augmented_prompt_in ):
generated_response = model_in.generate( augmented_prompt_in )
if ( "results" in generated_response ) \
and ( len( generated_response["results"] ) > 0 ) \
and ( "generated_text" in generated_response["results"][0] ):
return generated_response["results"][0]["generated_text"]
else:
print( "The model failed to generate an answer" )
print( "\nDebug info:\n" + json.dumps( generated_response, indent=3 ) )
return ""
output = generate( model, augmented_prompt )
print( output )
import re
def searchAndAnswer( model ):
question = input( "Type your question:\n")
if not re.match( r"\S+", question ):
print( "No question")
return
# Retrieve the relevant content
top_match = search( question )
if top_match is None:
print( "No good answer was found in the knowledge base" )
return;
article_text = top_match["text"]
# Augment a prompt with context
augmented_prompt = augment( prompt_template, article_text, question )
# Generate output
output = generate( model, augmented_prompt )
if not re.match( r"\S+", output ):
print( "The model failed to generate an answer")
print( "\nAnswer:\n" + output )
print( "\nSource: \"" + top_match["title"] + "\", " + top_match["author"] + " (" + top_match["published"] + ")" )
Test the solution by running the following cell multiple times.
*You will be prompted to enter a question each time.
searchAndAnswer( model )
You successfully completed this notebook!
You learned how to apply the general retrieval-augmented generation pattern with a Watson Discovery search component and a small knowledge base using watsonx.ai.
Check out our Documentation for more samples, tutorials, documentation, and how-tos.
Kevin MacDonald, Content Design - IBM Data and AI.
Copyright © 2023 IBM. This notebook and its source code are released under the terms of the MIT License.