Building interactive AI agents for ultra-fast machine learning tasks

The agent’s architecture is designed for modularity, scalability, and GPU acceleration. It consists of five core layers and one temporary data store that work together to transform natural language prompts into executables, data processing, and ML workflows. Figure 1 shows a high-level workflow showing how each layer interacts.

Let’s take a closer look at how the layers work together.

Layer 1: User Interface
The user interface was developed using a Streamlit-based conversational chatbot to allow users to interact with agents in plain English.

Layer 2: Agent Orchestrator
It is a central controller that works with all layers. It interprets user prompts, delegates execution to LLM to understand intent, invokes appropriate GPU-accelerated functionality from the tooling layer, and responds in natural language. Each orchestrator method is a lightweight wrapper around a GPU function. for example, _describe_data Inside the user query call basic_eda()meanwhile _optimize_ridge Inside the user query call optimize_ridge_regression().

Layer 3: LLM layer
The LLM layer acts as the agent’s inference engine and initializes the language model client to communicate with the Nemotron Nano 9B-v2 using the NVIDIA NIM API. This layer allows agents to interpret natural language inputs and transform them into structured executable actions through four key mechanisms: LLM models, retry strategies for resilient communication, function calls for structured tool calls, and function call windows.

• LLM model
  The architecture of the LLM layer is LLM-agnostic and works with any language model that supports function calls. For this application, we used the Nemotron Nano-9B-v2, which supports both function calls and advanced inference. Additionally, the model’s small size provides an optimal balance of efficiency and functionality, and it can be deployed on a single GPU for inference. It delivers up to 6x higher token generation throughput compared to other leading models in its size class, and its thought budget feature allows developers to control the number of “think” tokens used, reducing inference costs by up to 60%. This combination of outstanding performance and cost efficiency enables real-time, conversational workflows that are economically viable in production deployments.
• Retry strategies for resilient communication
  The LLM client implements an exponential backoff retry mechanism to handle temporary network failures and API rate limits, ensuring reliable communication even under adverse network conditions or high API loads.
• Function calls for structured tool calls
  Function calls bridge natural language and code execution by allowing LLM to translate user intent into structured tool calls in Agent Orchestrator. The agent defines the available tools using an OpenAI-compatible functional schema that specifies each tool’s name, purpose, parameters, and constraints.
• function call window
  Function calls transform LLM from a text generator to an inference engine capable of API orchestration. Nemotron Nano-9B-v2’s model provides a structured “API specification” of available tools to understand user intent, select the right function, extract the right type of parameters, and orchestrate multi-step data processing and ML operations. All of this is done through natural language, so you don’t need to understand API syntax or write any code.

  The complete function call flow shown in Figure 3 shows how natural language is translated into executable code. reference chat_agent.py and llm.py Scripts in GitHub code for the operations listed in Figure 3.

Layer 4: Memory layer
The memory layer (ExperimentStore) stores experiment metadata, including model configuration, performance metrics, and evaluation results such as accuracy and F1 score. This metadata is stored in a session-specific file in standard JSONL format and can be tracked and retrieved during a session using functions such as: get_recent_experiments() and show_history().

Layer 5: Temporary data storage
The temporary data storage layer contains session-specific output files (best_model.joblib and predictions.csv) are stored in your system’s temporary directory and user interface, and are available for immediate download and use. These files are automatically deleted when the agent shuts down.

Layer 6: Tool layer
The tools layer is the agent’s computational core and is responsible for performing data science functions such as data loading, exploratory data analysis (EDA), model training and evaluation, and hyperparameter optimization (HPO). The features selected for execution are based on the user’s query. Various optimization strategies are used, including:

1. Consistency and reproducibility
   Agents use various abstraction methods. scikit-learn (a popular open-source library) to ensure consistent data preprocessing and model training across training, test, and production environments. This design prevents common ML pitfalls such as data leakage and inconsistent preprocessing by automatically applying the exact same transformations (imputation values, scaling parameters, encoding mappings) learned during training to all inference data.
2. memory management
   Use memory optimization strategies to process large datasets. Float32 Conversion reduces memory usage, GPU memory management frees up active cache GPU memory, and dense output configurations are faster on the GPU compared to sparse formats.
3. Executing a function
   The tool execution agent uses the CUDA-X data science library, including: CUDF and cuML Achieve GPU-accelerated performance while maintaining the same syntax of pandas and scikit-learn. This zero-code-change speedup is achieved through Python’s module preloading mechanism, which allows developers to run existing CPU code on the GPU without refactoring. of cudf.pandas Accelerators replace pandas operations with equivalent GPUs. cuml.accel Automatically replace scikit-learn models with cuML’s GPU implementation.

The following command starts the Streamlit interface with GPU acceleration enabled for both the data processing and machine learning components.

python -m cudf.pandas -m cuml.accel -m streamlit run user_interface.py

Accelerate, modularize, and scale ML agents