Model Training and Hyperparameter Tuning Using AI Assistant from Databricks

At UpTeam, we focus on productivity and innovation, especially when it comes to data. Data is the most essential asset for driving insights and making informed decisions. But it can also be a complicated process when companies are not using the right tools and the required expertise to make sense of it. 

By using Databricks, we improve our capabilities to process and analyze data efficiently. Recently, we had the opportunity to test and use the Databricks AI Assistant in a project for a leading company in research and healthcare. Here's how we leveraged the Assistant to optimize model training and hyperparameter tuning for getting refined results and predictions.

Seamless Integration and User-Friendly Experience

While our team has been working with Databricks for some time and also some team members have recently passed Data Engineer certifications, setting up the AI assistant was not a problem.

The AI Assistant from Databricks integrates pretty smoothly into existing workflows, providing AI-driven suggestions that feel like an extra team member. It quickly adapts to specific needs, enhancing daily operations significantly.

When working with Databricks, ingesting data into Delta Lake tables is a critical first step. Databricks Assistant simplifies this process with practical prompts so here are a few examples that are a reference on how we have used it:

1. Ingesting Data from an API (using Python):

Help me ingest patient data from this API into a Delta Lake table: https://healthdata.api/patients?status=active. Make sure to use PySpark, and be concise! If the Spark DataFrame columns have any spaces in them, make sure to remove them from the Spark DF.

2. Ingesting JSON Files from Cloud Storage 

I have JSON files in a healthcare data volume here: /Volumes/healthdata/default/patient_records.json. Write code to ingest this data into a Delta Lake table. Use SQL only, and be concise!

These prompts ensure the correct and efficient ingestion of data, enhancing workflow efficiency.

3. Transforming Data from Unstructured to Structured

Adhering to these data principles, Databricks Assistant helps extract structured data from unstructured formats using techniques like regular expressions and handling nested data structures.

Regular Expressions:

For example, parsing a "MedicalReport" column to extract report titles and dates:

Here is an example of the MedicalReport column in our dataset: 1. Annual Physical Examination (2021). The report title is between the number and the parentheses, and the date is within the parentheses. Write a function that extracts both the report title and the date from the MedicalReport column in the healthcare_raw DataFrame.

4. Nested Structs and Arrays:

When dealing with nested JSON data, Databricks Assistant can flatten these structures and extract the necessary metrics:

Write PySpark code to flatten the df and extract patient information, including medications and dosages.

This reduces the complexity and time required to handle nested data.

5. Refactoring, Debugging, and Optimization

Databricks Assistant can analyze and optimize poorly written code, improving readability and performance. For example, refactoring a Python function to calculate the total cost of treatments:

Rewrite this code in a more performant way, commented properly, and documented according to Python function documentation standards.

This prompt results in code that is more efficient and easier to maintain.

6. Transpiling Pandas to PySpark

The Assistant can convert inefficient Pandas code to PySpark, making it scalable for larger datasets:

Convert this Pandas code to PySpark for better performance and scalability.

While those are just a few prompts displayed, we have also used it when deciding to further diagnose errors and quickly identify and suggest corrections in the code, from typos to code optimizations. We have also used it for writing tests and enhancing the data pipeline reliability. Of course, it is a very useful tool to get help when needing to create and access documentation and Knowledge Base and help our customers navigate and efficiently utilize data.

Challenges and Solutions

Like any tool, there were initial challenges in integrating Databricks Assistant into our workflow. Understanding the best way to phrase prompts and utilize the Assistant’s features took some time. However, with a bit of practice and experimentation, we quickly overcame these challenges. We also found that occasionally, the Assistant's suggestions weren't perfect. This is where the human touch remains essential. By reviewing and tweaking the AI-generated code, we ensured that the final output met our standards. These experiences have taught us to see the Assistant as a valuable aid rather than a replacement for our skills and expertise.

In a data migration project, the AI Assistant efficiently transformed and loaded large datasets into Delta Lake tables, ensuring smooth and timely completion. In data transformation tasks, it handled nested JSON structures, accelerating data extraction and analysis. These capabilities have been particularly beneficial for projects involving complex data sources, delivering high-quality insights to clients.

We are excited about the potential enhancements Databricks might bring to the Assistant. Features that further improve context understanding and handle even more complex tasks could make it an indispensable tool. As Databricks continues to innovate, we believe the Assistant will play a key role in driving our productivity and success to new heights.

By implementing the AI Assistant from Databricks, our team proved how cutting-edge AI can transform model training and hyperparameter tuning, providing significant efficiency gains and high-quality results for a leading company in research and healthcare.

‍