Keyword extraction with kluster.ai API¶

One simple but powerful use case for AI models is keyword extraction, in which you feed the model a large dataset and ask it to provide a given number of keywords.

This tutorial runs through a notebook where you'll learn how to use the kluster.ai batch API to obtain keywords from a given dataset.

The example uses an extract from the AG News dataset. You can adapt this example by using the data and categories relevant to your use case. With this approach, you can effortlessly process datasets of any scale, big or small, to obtain keywords, all powered by a state-of-the-art language model.

Prerequisites¶

Before getting started, ensure you have the following:

• A kluster.ai account - sign up on the kluster.ai platform if you don't have one
• A kluster.ai API key - after signing in, go to the API Keys section and create a new key. For detailed instructions, check out the Get an API key guide

Setup¶

In this notebook, we'll use Python's getpass module to input the key safely. After execution, please provide your unique kluster.ai API key (ensure no spaces).

In [1]:

from getpass import getpass

api_key = getpass("Enter your kluster.ai API key: ")

from getpass import getpass api_key = getpass("Enter your kluster.ai API key: ")

Next, ensure you've installed OpenAI Python library:

In [2]:

pip install -q openai

pip install -q openai

Note: you may need to restart the kernel to use updated packages.

With the OpenAI Python library installed, we import the necessary dependencies for the tutorial:

In [3]:

from openai import OpenAI

import pandas as pd
import time
import json
import os
from IPython.display import clear_output, display

from openai import OpenAI import pandas as pd import time import json import os from IPython.display import clear_output, display

And then, initialize the client by pointing it to the kluster.ai endpoint, and passing your API key.

In [4]:

# Set up the client
client = OpenAI(
    base_url="https://api.kluster.ai/v1",
    api_key=api_key,
)

# Set up the client client = OpenAI( base_url="https://api.kluster.ai/v1", api_key=api_key, )

Get the data¶

Now that you've initialized an OpenAI-compatible client pointing to kluster.ai, we can discuss the data.

This notebook comes with a preloaded sample dataset based on the AG News dataset. It includes sections of news headlines and their leads, all set for processing. No additional setup is needed. Proceed to the next steps to begin working with this data.

In [6]:

df = pd.DataFrame({
    "text": [
        "Chorus Frog Found Croaking in Virginia - The Southern chorus frog has been found in southeastern Virginia, far outside its previously known range. The animal had never before been reported north of Beaufort County, N.C., about 125 miles to the south.",
        "Expedition to Probe Gulf of Mexico - Scientists will use advanced technology never before deployed beneath the sea as they try to discover new creatures, behaviors and phenomena in a 10-day expedition to the Gulf of Mexico's deepest reaches.",
        "Feds Accused of Exaggerating Fire ImpactP - The Forest Service exaggerated the effect of wildfires on California spotted owls in justifying a planned increase in logging in the Sierra Nevada, according to a longtime agency expert who worked on the plan.",
        "New Method May Predict Quakes Weeks Ahead - Swedish geologists may have found a way to predict earthquakes weeks before they happen by monitoring the amount of metals like zinc and copper in subsoil water near earthquake sites, scientists said Wednesday.",
        "Marine Expedition Finds New Species - Norwegian scientists who explored the deep waters of the Atlantic Ocean said Thursday their findings #151; including what appear to be new species of fish and squid #151; could be used to protect marine ecosystems worldwide."
    ]
})

df = pd.DataFrame({ "text": [ "Chorus Frog Found Croaking in Virginia - The Southern chorus frog has been found in southeastern Virginia, far outside its previously known range. The animal had never before been reported north of Beaufort County, N.C., about 125 miles to the south.", "Expedition to Probe Gulf of Mexico - Scientists will use advanced technology never before deployed beneath the sea as they try to discover new creatures, behaviors and phenomena in a 10-day expedition to the Gulf of Mexico's deepest reaches.", "Feds Accused of Exaggerating Fire ImpactP - The Forest Service exaggerated the effect of wildfires on California spotted owls in justifying a planned increase in logging in the Sierra Nevada, according to a longtime agency expert who worked on the plan.", "New Method May Predict Quakes Weeks Ahead - Swedish geologists may have found a way to predict earthquakes weeks before they happen by monitoring the amount of metals like zinc and copper in subsoil water near earthquake sites, scientists said Wednesday.", "Marine Expedition Finds New Species - Norwegian scientists who explored the deep waters of the Atlantic Ocean said Thursday their findings #151; including what appear to be new species of fish and squid #151; could be used to protect marine ecosystems worldwide." ] })

Perform batch inference¶

To execute the batch inference job, we'll take the following steps:

1. Create the batch job file - we'll generate a JSON lines file with the desired requests to be processed by the model
2. Upload the batch job file - once it is ready, we'll upload it to the kluster.ai platform using the API, where it will be processed. We'll receive a unique ID associated with our file
3. Start the batch job - after the file is uploaded, we'll initiate the job to process the uploaded data, using the file ID obtained before
4. Monitor job progress - (optional) track the status of the batch job to ensure it has been successfully completed
5. Retrieve results - once the job has completed execution, we can access and process the resultant data

This notebook is prepared for you to follow along. Run the cells below to watch it all come together.

Create the batch file¶

This example selects the klusterai/Meta-Llama-3.1-405B-Instruct-Turbo model. If you'd like to use a different model, feel free to change it by modifying the model field. In this notebook, you can also comment Llama 3.1 405B, and uncomment whatever model you want to try out.

Please refer to the Supported models section for a list of the models we support.

The following snippets prepare the JSONL file, where each line represents a different request. Note that each separate batch request can have its model. Also, we are using a temperature of 0.5, but feel free to change it and play around with the different outcomes (but we are only asking to respond with a single word, the genre).

In [7]:

# Prompt
SYSTEM_PROMPT = '''
    Extract up to 5 relevant keywords from the given text. Provide only the keywords between double quotes and separated by commas.
    '''

# Models
#model="deepseek-ai/DeepSeek-R1"
#model="deepseek-ai/DeepSeek-V3"
#model="klusterai/Meta-Llama-3.1-8B-Instruct-Turbo"
model="klusterai/Meta-Llama-3.1-405B-Instruct-Turbo"
#model="klusterai/Meta-Llama-3.3-70B-Instruct-Turbo"
#model="Qwen/Qwen2.5-VL-7B-Instruct"

# Ensure the directory exists
os.makedirs("keyword_extraction/data", exist_ok=True)

# Create the batch job file with the prompt and content
def create_batch_file(df):
    batch_list = []
    for index, row in df.iterrows():
        content = row['text']

        request = {
            "custom_id": f"keyword_extraction-{index}",
            "method": "POST",
            "url": "/v1/chat/completions",
            "body": {
                "model": model,
                "temperature": 0.5,
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": content}
                ],
            }
        }
        batch_list.append(request)
    return batch_list

# Save file
def save_batch_file(batch_list):
    filename = f"keyword_extraction/batch_job_request.jsonl"
    with open(filename, 'w') as file:
        for request in batch_list:
            file.write(json.dumps(request) + '\n')
    return filename

# Prompt SYSTEM_PROMPT = ''' Extract up to 5 relevant keywords from the given text. Provide only the keywords between double quotes and separated by commas. ''' # Models #model="deepseek-ai/DeepSeek-R1" #model="deepseek-ai/DeepSeek-V3" #model="klusterai/Meta-Llama-3.1-8B-Instruct-Turbo" model="klusterai/Meta-Llama-3.1-405B-Instruct-Turbo" #model="klusterai/Meta-Llama-3.3-70B-Instruct-Turbo" #model="Qwen/Qwen2.5-VL-7B-Instruct" # Ensure the directory exists os.makedirs("keyword_extraction/data", exist_ok=True) # Create the batch job file with the prompt and content def create_batch_file(df): batch_list = [] for index, row in df.iterrows(): content = row['text'] request = { "custom_id": f"keyword_extraction-{index}", "method": "POST", "url": "/v1/chat/completions", "body": { "model": model, "temperature": 0.5, "messages": [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": content} ], } } batch_list.append(request) return batch_list # Save file def save_batch_file(batch_list): filename = f"keyword_extraction/batch_job_request.jsonl" with open(filename, 'w') as file: for request in batch_list: file.write(json.dumps(request) + '\n') return filename

Let's run the functions we've defined before:

In [8]:

batch_list = create_batch_file(df)
filename = save_batch_file(batch_list)

batch_list = create_batch_file(df) filename = save_batch_file(batch_list)

Next, we can preview what that batch job file looks like:

In [9]:

!head -n 1 keyword_extraction/batch_job_request.jsonl

!head -n 1 keyword_extraction/batch_job_request.jsonl

{"custom_id": "keyword_extraction-0", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "klusterai/Meta-Llama-3.1-405B-Instruct-Turbo", "temperature": 0.5, "messages": [{"role": "system", "content": "\n    Extract up to 5 relevant keywords from the given text. Provide only the keywords between double quotes and separated by commas.\n    "}, {"role": "user", "content": "Chorus Frog Found Croaking in Virginia - The Southern chorus frog has been found in southeastern Virginia, far outside its previously known range. The animal had never before been reported north of Beaufort County, N.C., about 125 miles to the south."}]}}

Upload inference file to kluster.ai¶

Now that we’ve prepared our input file, it’s time to upload it to the kluster.ai platform. To do so, you can use the files.create endpoint of the client, where the purpose is set to batch. This will return the file ID, which we need to log for the next steps.

In [10]:

data_dir = 'keyword_extraction/batch_job_request.jsonl'

# Upload batch job request file
with open(data_dir, 'rb') as file:
    upload_response = client.files.create(
        file=file,
        purpose="batch"
    )

    # Print job ID
    file_id = upload_response.id
    print(f"File uploaded successfully. File ID: {file_id}")

data_dir = 'keyword_extraction/batch_job_request.jsonl' # Upload batch job request file with open(data_dir, 'rb') as file: upload_response = client.files.create( file=file, purpose="batch" ) # Print job ID file_id = upload_response.id print(f"File uploaded successfully. File ID: {file_id}")

File uploaded successfully. File ID: 67e40f05331c771bde3bca69

Start the batch job¶

Once the file has been successfully uploaded, we're ready to start (create) the batch job by providing the file ID we got in the previous step. To do so, we use the batches.create method, for which we need to set the endpoint to /v1/chat/completions. This will return the batch job details with the ID.

In [11]:

# Create batch job with completions endpoint
batch_job = client.batches.create(
    input_file_id=file_id,
    endpoint="/v1/chat/completions",
    completion_window="24h"
)

print("\nBatch job created:")
batch_dict = batch_job.model_dump()
print(json.dumps(batch_dict, indent=2))

# Create batch job with completions endpoint batch_job = client.batches.create( input_file_id=file_id, endpoint="/v1/chat/completions", completion_window="24h" ) print("\nBatch job created:") batch_dict = batch_job.model_dump() print(json.dumps(batch_dict, indent=2))

Batch job created:
{
  "id": "67e40f070997f511a77bf70b",
  "completion_window": "24h",
  "created_at": 1742999303,
  "endpoint": "/v1/chat/completions",
  "input_file_id": "67e40f05331c771bde3bca69",
  "object": "batch",
  "status": "pre_schedule",
  "cancelled_at": null,
  "cancelling_at": null,
  "completed_at": null,
  "error_file_id": null,
  "errors": [],
  "expired_at": null,
  "expires_at": 1743085703,
  "failed_at": null,
  "finalizing_at": null,
  "in_progress_at": null,
  "metadata": {},
  "output_file_id": null,
  "request_counts": {
    "completed": 0,
    "failed": 0,
    "total": 0
  }
}

Check job progress¶

Now that your batch job has been created, you can track its progress.

To monitor the job's progress, we can use the batches.retrieve method and pass the batch job ID. The response contains a status field that tells us if it is completed or not and the subsequent status of each job separately.

The following snippet checks the status every 10 seconds until the entire batch is completed:

In [ ]:

all_completed = False

# Loop to check status every 10 seconds
while not all_completed:
    all_completed = True
    output_lines = []

    updated_job = client.batches.retrieve(batch_job.id)

    if updated_job.status != "completed":
        all_completed = False
        completed = updated_job.request_counts.completed
        total = updated_job.request_counts.total
        output_lines.append(f"Job status: {updated_job.status} - Progress: {completed}/{total}")
    else:
        output_lines.append(f"Job completed!")

    # Clear the output and display updated status
    clear_output(wait=True)
    for line in output_lines:
        display(line)

    if not all_completed:
        time.sleep(10)

all_completed = False # Loop to check status every 10 seconds while not all_completed: all_completed = True output_lines = [] updated_job = client.batches.retrieve(batch_job.id) if updated_job.status != "completed": all_completed = False completed = updated_job.request_counts.completed total = updated_job.request_counts.total output_lines.append(f"Job status: {updated_job.status} - Progress: {completed}/{total}") else: output_lines.append(f"Job completed!") # Clear the output and display updated status clear_output(wait=True) for line in output_lines: display(line) if not all_completed: time.sleep(10)

'Job status: in_progress - Progress: 1/5'

Get the results¶

With the job completed, we'll retrieve the results and review the responses generated for each request. The results are parsed. To fetch the results from the platform, you must retrieve the output_file_id from the batch job and then use the files.content endpoint, providing that specific file ID. Note that the job status must be completed to retrieve the results!

In [ ]:

#Parse results as a JSON object
def parse_json_objects(data_string):
    if isinstance(data_string, bytes):
        data_string = data_string.decode('utf-8')

    json_strings = data_string.strip().split('\n')
    json_objects = []

    for json_str in json_strings:
        try:
            json_obj = json.loads(json_str)
            json_objects.append(json_obj)
        except json.JSONDecodeError as e:
            print(f"Error parsing JSON: {e}")

    return json_objects

# Retrieve results with job ID
job = client.batches.retrieve(batch_job.id)
result_file_id = job.output_file_id
result = client.files.content(result_file_id).content

# Parse JSON results
parsed_result = parse_json_objects(result)

# Extract and print only the content of each response
print("\nExtracted Responses:")
for item in parsed_result:
    try:
        content = item["response"]["body"]["choices"][0]["message"]["content"]
        print(content)
    except KeyError as e:
        print(f"Missing key in response: {e}")

#Parse results as a JSON object def parse_json_objects(data_string): if isinstance(data_string, bytes): data_string = data_string.decode('utf-8') json_strings = data_string.strip().split('\n') json_objects = [] for json_str in json_strings: try: json_obj = json.loads(json_str) json_objects.append(json_obj) except json.JSONDecodeError as e: print(f"Error parsing JSON: {e}") return json_objects # Retrieve results with job ID job = client.batches.retrieve(batch_job.id) result_file_id = job.output_file_id result = client.files.content(result_file_id).content # Parse JSON results parsed_result = parse_json_objects(result) # Extract and print only the content of each response print("\nExtracted Responses:") for item in parsed_result: try: content = item["response"]["body"]["choices"][0]["message"]["content"] print(content) except KeyError as e: print(f"Missing key in response: {e}")

Summary¶

This tutorial used the chat completion endpoint to perform a simple keyword extraction task with batch inference. This particular example extracted keywords from an AG News dataset.

To submit a batch job we've:

1. Created the JSONL file, where each line of the file represented a separate request
2. Submitted the file to the platform
3. Started the batch job, and monitored its progress
4. Once completed, we fetched the results

All of this using the OpenAI Python library and API, no changes needed!

Kluster.ai's batch API empowers you to scale your workflows seamlessly, making it an invaluable tool for processing extensive datasets. As the next steps, feel free to create your dataset or expand on this example. Good luck!