Initializing the SDK client

pip install google-genai

# with Google AI API key
GOOGLE_AI_API_KEY = os.getenv("GOOGLE_AI_API_KEY")

client = genai.Client(api_key=GOOGLE_AI_API_KEY)

response = client.models.generate_content(model="gemini-2.0-flash-exp", contents="How far is the moon from the earth?")

display(Markdown(response.text))
# The distance between the Earth and the Moon isn't constant, as the Moon travels in an elliptical orbit around our planet. Therefore, we talk about the distance in terms of averages:

# Average distance: Approximately 238,855 miles (384,400 kilometers).
# However, keep in mind:

# Perigee (Closest Approach): The Moon can get as close as about 225,623 miles (363,104 kilometers).
# Apogee (Farthest Distance): The Moon can be as far as about 252,088 miles (405,696 kilometers).
# So, while the average distance is a good figure to keep in mind, it's important to remember that the actual distance varies throughout the month.

print(response.usage_metadata)
# cached_content_token_count=None candidates_token_count=181 prompt_token_count=10 total_token_count=191

response = await client.aio.models.generate_content(
    model="gemini-2.0-flash-exp", contents="How far is the moon from the earth?", config=gen_config
)

# this is the path to your json credentials file (at least it is to mine, :))
GOOGLE_VERTEX_AI_CREDS = os.getenv("GOOGLE_VERTEX_AI_CREDS")

SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
creds = service_account.Credentials.from_service_account_file(GOOGLE_VERTEX_AI_CREDS, scopes=SCOPES)

client = genai.Client(vertexai=True, project="generative-playground", location="us-central1", credentials=creds)

response = client.models.generate_content(model="gemini-2.0-flash-exp", contents="How far is the moon from the earth?")

display(Markdown(response.text))
# The distance between the Earth and the Moon is not constant, as the Moon's orbit is elliptical (an oval shape), not a perfect circle. Here's a breakdown:

# Average Distance: The average distance is about 384,400 kilometers (238,900 miles). This is the number most commonly used.
# Perigee: This is the point in the Moon's orbit where it is closest to Earth. At perigee, the distance can be as close as 363,104 kilometers (225,623 miles).
# Apogee: This is the point in the Moon's orbit where it is furthest from Earth. At apogee, the distance can reach as far as 405,696 kilometers (252,088 miles).
# So, the answer is variable, but the average distance is approximately 384,400 kilometers (238,900 miles).

print(response.usage_metadata)
# cached_content_token_count=None candidates_token_count=221 prompt_token_count=9 total_token_count=230

for model in client.models.list():
    print(model)  # ain't returning nothing useful

Controlling Generation

You can influence the generation of content by providing a GenerateContentConfig object. With this object, you can provide a system instruction, safety settings, and generation parameters (e.g., temperature, top_p, etc.).

See also:

• Content Generation Parameters
• Safety Settings

safety_settings = [
    types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="BLOCK_NONE"),
    types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="BLOCK_NONE"),
    types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="BLOCK_NONE"),
    types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_NONE"),
]

gen_config = types.GenerateContentConfig(
    system_instruction="You are an expert in all things astronomy and you ONLY provide very concise answers",
    safety_settings=safety_settings,
    temperature=0,
    top_p=0.95,
    top_k=20,
    candidate_count=1,
    seed=5,
    max_output_tokens=100,
    stop_sequences=["STOP!"],
    presence_penalty=0.0,
    frequency_penalty=0.0,
)

response = client.models.generate_content(model="gemini-2.0-flash-exp", contents="How far is the moon from the earth?", config=gen_config)

With that in place we can see that our response is now much more concise and to the point.

response = client.models.generate_content(model="gemini-2.0-flash-exp", contents="How far is the moon from the earth?", config=gen_config)
print(response.text)
# About 384,400 kilometers.

print(response.usage_metadata)
# cached_content_token_count=None candidates_token_count=12 prompt_token_count=24 total_token_count=36

Tokenization Goodies

You can use the count_tokens method to calculate the number of input tokens before sending a request to the Gemini API. Super helpful to understanding the cost implications of your requests. Async operations are available for both count_tokens and compute_tokens as well as most other methods on your client instance via the aio property.

# count_tokens
response = client.models.count_tokens(model="gemini-2.0-flash-exp", contents="How far is the moon from the earth?")
print(response)
# total_tokens=9 cached_content_token_count=None

print("\n======\n")

# count_tokens async
response = await client.aio.models.count_tokens(model="gemini-2.0-flash-exp", contents="How far is the moon from the earth?")
print(response)
# total_tokens=9 cached_content_token_count=None

You can use compute_tokens to get the tokenized inputs (only available in Vertex AI). This can be really helpful in terms of understanding how the Gemini models tokenize your input and troubleshooing issues that might arise from generation.

# compute_tokens
response = client.models.compute_tokens(model="gemini-2.0-flash-exp", contents="How far is the moon from the earth?")
print(response)
# tokens_info=[TokensInfo(role='user', token_ids=['2299', '2166', '603', '573', '11883', '774', '573', '6683', '235336'], tokens=[b'How', b' far', b' is', b' the', b' moon', b' from', b' the', b' earth', b'?'])]


print("\n======\n")

# compute_tokens async
response = await client.aio.models.compute_tokens(model="gemini-2.0-flash-exp", contents="How far is the moon from the earth?")
print(response)
# tokens_info=[TokensInfo(role='user', token_ids=['2299', '2166', '603', '573', '11883', '774', '573', '6683', '235336'], tokens=[b'How', b' far', b' is', b' the', b' moon', b' from', b' the', b' earth', b'?'])]

Streaming Content

We’ve already seen how we can use the generate_content method to generate content but what if we want to stream the content as it is being generated?

The generate_content_stream method is a great way to stream content as it is generated. This is useful for things like chat interfaces or long-form content where you want to see the progress as it is being generated.

An async version is also available via the aio property.

for chunk in client.models.generate_content_stream(model="gemini-2.0-flash-exp", contents="How far is the moon from the earth?"):
    print(chunk.text, end="|", flush=True)

# The| distance between the Earth and the Moon is not constant, as the Moon's| orbit is elliptical. Here's a breakdown:

# * **Average Distance:**| The average distance between the Earth and the Moon is about **384,400 kilometers (238,900 miles)**. This| is the most commonly cited figure.

# * **Perigee:** This is the point in the Moon's orbit when it is closest to Earth. At| perigee, the distance can be as little as about **363,104 kilometers (225,623 miles)**.

# * **Apogee:** This is the point in the Moon's orbit| when it is farthest from Earth. At apogee, the distance can be as much as about **405,696 kilometers (252,088 miles)**.

# **Key Takeaway:** While we often| talk about the average distance, remember that the Moon's distance varies throughout its orbit, fluctuating between perigee and apogee.

# So, while the average is 384,400 kilometers, it's more accurate to say the distance ranges from about 363,000| km to 406,000 km.
|

I’ve included the pipe character so you can see the chunks as they are being generated.

Multi-Turn Chat

The chats module provides a way to interact with the Gemini API in a multi-turn chat interface.

The create method is used to create a new chat session. The send_message method is used to send a message to the chat session.

The history parameter is used to provide the chat history.

system_instruction = dedent("""
  You are an expert software developer and a helpful coding assistant.
  You are able to generate high-quality code in any programming language.
""").strip()

convo = []
chat = client.chats.create(
    model="gemini-2.0-flash-exp", config=types.GenerateContentConfig(system_instruction=system_instruction, temperature=0.5), history=convo
)

You can now use send_message to add a message to the chat.

response = chat.send_message("Write a simple python program that accepts a person's name and returns a greeting.")

Markdown(response.text)
# def greet_person(name):
#   """
#   Greets a person by name.

#   Args:
#     name: The name of the person to greet (string).

#   Returns:
#     A greeting string.
#   """
#   return f"Hello, {name}!"

# if __name__ == "__main__":
#   person_name = input("Please enter your name: ")
#   greeting = greet_person(person_name)
#   print(greeting)

# ... and an extended explanation of the code I'm omitting here ...

Currently we can look at, and even modify, the chat history by accessing the _curated_history property on the chat instance. That _ prefix tells me this part of the SDK is still under construction and likely to change.

chat._curated_history
# [Content(parts=[Part(video_metadata=None, code_execution_result=None, executable_code=None, file_data=None, function_call=None, function_response=None, inline_data=None, text="Write a simple python program that accepts a person's name and returns a greeting.")], role='user'),
#  Content(parts=[Part(video_metadata=None, code_execution_result=None, executable_code=None, file_data=None, function_call=None, function_response=None, inline_data=None, text='```python\ndef greet_person(name):\n  """\n  Greets a person by name.\n\n  Args:\n    name: The name of the person to greet (string).\n\n  Returns:\n    A greeting string.\n  """\n  return f"Hello, {name}!"\n\nif __name__ == "__main__":\n  person_name = input("Please enter your name: ")\n  greeting = greet_person(person_name)\n  print(greeting)\n```\n\n**Explanation:**\n\n1.  **`def greet_person(name):`**:\n    *   This line defines a function named `greet_person` that takes one argument, `name`.\n    *   The `name` argument will hold the person\'s name as a string.\n\n2.  **`"""..."""`**:\n    *   This is a docstring, which is a multiline string used to document what the function does. It\'s good practice to include docstrings to make your code more understandable.\n\n3.  **`return f"Hello, {name}!"`**:\n    *   This line uses an f-string (formatted string literal) to create the greeting.\n    *   `f"..."`  allows you to embed variables directly into the string by placing them inside curly braces `{}`.\n    *   The function returns the complete greeting string (e.g., "Hello, Alice!").\n\n4.  **`if __name__ == "__main__":`**:\n    *   This is a standard Python construct that ensures the code inside the `if` block only runs when the script is executed directly (not when it\'s imported as a module into another script).\n\n5.  **`person_name = input("Please enter your name: ")`**:\n    *   The `input()` function displays the message "Please enter your name: " to the user and waits for them to type something.\n    *   Whatever the user types is stored as a string in the `person_name` variable.\n\n6.  **`greeting = greet_person(person_name)`**:\n    *   This line calls the `greet_person` function, passing the `person_name` as an argument.\n    *   The function returns the greeting string, which is then stored in the `greeting` variable.\n\n7.  **`print(greeting)`**:\n    *   Finally, this line prints the greeting to the console.\n\n**How to run this code:**\n\n1.  Save the code in a file named, for example, `greeting.py`.\n2.  Open a terminal or command prompt.\n3.  Navigate to the directory where you saved the file.\n4.  Run the command `python greeting.py`.\n5.  The program will prompt you to enter your name. Type your name and press Enter.\n6.  The program will then print the greeting to the console.\n\n**Example Interaction:**\n\n```\nPlease enter your name: Bob\nHello, Bob!\n```\n')], role='model')

response = chat.send_message("Okay, write a unit test of the generated function.")

Markdown(response.text)
# import unittest
# from greeting import greet_person  # Assuming the previous code is in greeting.py

# class TestGreeting(unittest.TestCase):

#     def test_greet_with_valid_name(self):
#         self.assertEqual(greet_person("Alice"), "Hello, Alice!")
#         self.assertEqual(greet_person("Bob"), "Hello, Bob!")
#         self.assertEqual(greet_person("Charlie"), "Hello, Charlie!")

#     def test_greet_with_empty_name(self):
#         self.assertEqual(greet_person(""), "Hello, !")

#     def test_greet_with_name_containing_spaces(self):
#         self.assertEqual(greet_person("John Doe"), "Hello, John Doe!")

#     def test_greet_with_name_containing_numbers(self):
#         self.assertEqual(greet_person("User123"), "Hello, User123!")

#     def test_greet_with_special_characters(self):
#         self.assertEqual(greet_person("!@#$%^"), "Hello, !@#$%^!")

# if __name__ == '__main__':
#     unittest.main()

# ... and an extended explanation of the code I'm omitting here ...

Pretty cool!

Conclusion

The SDK is still being fleshed out but I’m excited about the prospect of being able build with Gemini regardless of which API I’m using. Hopefully, some of the tips and tricks above will help you get started regardless of which API flavor you prefer.

In the next post, we’ll look at how to take advantage of Gemini’s multimodal capabilities including how to incorporate few-shot examples for both text only and multimodal tasks.

The Gemini Family

Figure 1: The Gemini Famly (google)

Performance

Internal Benchmarks

To better understand how the latest Gemini 2.0 Flash Experimental model stacks up against its predecessors, let’s examine the reported benchmark results.

The following table highlights the performance improvements compared to the Gemini 1.5 generation models.

Figure 2: Reported benchmarks (google)

The 1.5 models have worked really well for me across a number of projects and it looks like the 2.0 generation represents a substantial move in the right direction across the board! According to Google’s announcement post …

“Gemini 2.0 Flash builds on the success of 1.5 Flash, our most popular model yet for developers, with enhanced performance at similarly fast response times. Notably, 2.0 Flash even outperforms 1.5 Pro on key benchmarks, at twice the speed”

How Does It Compare to the Competition?x

Evaluating performance across large language models (LLMs) can be challenging due to the abundance of available benchmarks, not all of which provide consistent or reliable insights. Among these options, one resource that I’ve found consistently aligns with my own experiences is lmarena.ai (formerly LMSYS), which I will reference for this comparison.

Figure 2: Overall performance comparison across leading LLM models (courtesy of lmarena.ai)

Figure 3: Detailed breakdown of performance metrics across different categories (courtesy of lmarena.ai)

That’s alot of Geminis up there at the top!

Cost

As AI becomes an increasingly integral part of application infrastructure, the cost of using these models is becoming a critical consideration. Understanding how costs vary across providers like OpenAI, Anthropic, and others is essential for making informed decisions. So, how do the costs of using these models generally compare, and where does Gemini stand in this landscape?

Note: Input tokens refer to the tokens in the prompt sent to the model, while output tokens are those generated by the model in response.

Conclusion

If the future is indeed agentic, as I hypothesize, and cost-efficient, high-performing, general-use models play a central role in that vision, then Gemini emerges as a compelling choice. Based on the comparison of performance between Gemini models and those of other major API providers, combined with its cost-effectiveness, Gemini is undoubtedly worth considering as part of your AI strategy moving forward.

In the next blog in this series, we’ll explore the basics of using Gemini and dive into some of its standout features that set it apart from the competition using the new and improved unified SDK in 2.0. This unified SDK is one of the more exicting bits to developers like myself that have had to bounce back-and-forth between using the Gemini API and Vertex API to talk with Gemini. With 2.0, we finally have one SDK to rule them all (even if it still means we need to understand how the two different API options vary and when to use one over the other … something we’ll talk about in the next post as well).

What is a “Great Eval”?

I really liked the definition Ankur Goyal’s provided in his conference talk on doing LLM Eval For Text2SQL. He says that a “great eval” is comprised of three componnents:

1. Some data to act on
2. A task to perform with that data (e.g., “I want you to take some input and generate some output”)
3. One or more scoring functions that measure how well your task did the job

So as we iterate through our evaluation pipeline, that’s what we’ll do. At each step we’ll grab some data, define a task, and configure a number of scoring functions. It’s really that simple. That is an “eval.”

But before we get to that, what exaclty is an “L1” eval?

Level 1: Unit Tests (L1 Evals)

From the blog post mentioned above, we can infer some things about these particular evals.

1. Each should test a specific feature or scenario (e.g., the model called all the tools we expected).

2. They can also be general tests that aren’t specific to any feature or scenario (e.g., the model returns valid JSON).

3. You can use an LLM to synthetically generate inputs for these evals if necessary (see here for an exmaple)

4. You need to be able to track the results of these tests over time.

And remember this very important advice on interpreting the worthiness of your evals …

Lucky for us there is a really great tool that does all of this. It’s called Braintrust and it’s awesome! Before we get to it, we need to prepare a dataset to generate some predictions from.

Eval Preparation

Before we start eval’ng, we need to define our tools and a few methods for building a dataset we can pipe into Braintrust to get some initial feedback.

# For each task/tool, I include a number of prompts I imagine a user might use to implore the AI to perform
# that task.  When I build the actually examples, I randomly grab a single item for each task that should
# be called to build the final user prompt
example_translation_asks = [
    "If the text below is not in English, translate it into English; otherwise, return the text as is",
    "Translate the text below into English if it is not already in English; if it is, return it unchanged",
    "If the provided text is in a language other than English, translate it to English; if it's already in English, return it without changes",  # noqa: E501
    "Translate the following text to English if it is not written in English; otherwise, leave it unaltered",
    "If the text below is not in English, convert it to English; if it is in English, return it without modification",
]

example_summarization_asks = [
    "Provide a short summary of the text below and identify any overarching themes",
    "Summarize the document below and highlight any broad themes present",
    "Create a concise summary of the following document and identify any general themes",
    "Extract a brief 1-2 sentence summary from the text below and determine any broad themes",
    "Summarize the text below in 1-2 sentences and identify any major themes within it",
]
example_ner_asks = [
    "Extract the named entities from the author's statement below",
    "Determine the named entities in the document below",
    "Identify the specific named entities in the statement below",
    "Perform NER on the text below",
    "Do NER on the document below",
]
example_sentiment_asks = [
    "Determine the sentiment expressed in the document below",
    "Do sentiment analysis on the text below",
    "Identify and classify the sentiments expressed in the following document",
    "Perform sentiment analysis on the document below",
    "Extract the sentiment from the document below",
]
example_theme_asks = [
    "Identify the theme in the documents below and create an action plan based on them",
    "Determine the theme in the provided documents and develop an action plan accordingly",
    "Extract the theme in the documents below and formulate an action plan based on your findings",
    "Based on the documents provided, extract their theme and an action plan based on them",
    "What is a good theme that describes the text below? Also include an action plan based on the provided documents",
]


# Given a list of tools to call and whether our context is a single document or a colleciton of documents,
# this utility function does the work of coallesicing all the asks into a single human prompt.  Note how
# for "translation" tasks we sometimes ask for translation specifically and at other times we don't because
# I hope that ultimately, the model can just see the text is not in English and do the translation task
# without being specifically asked too.
def build_hypothetical_prompt(tool_names, is_doc: bool = False):
    # Vary the overall format of the prompt for different styles by which a user may ask for
    # different tasks
    if random.choice([True, False]):
        queries = []
        for tool_name in tool_names:
            if tool_name == "TranslationTask":
                # Vary this being asked explicity so model can learn to call this tool whenver
                # it detects a non-English document
                if random.choice([True, False]):
                    queries.append(random.choice(example_translation_asks))
            elif tool_name == "DocumentSummaryTask":
                queries.append(random.choice(example_summarization_asks))
            elif tool_name == "DocumentNERTask":
                queries.append(random.choice(example_ner_asks))
            elif tool_name == "DocumentSentimentTask":
                queries.append(random.choice(example_sentiment_asks))
            elif tool_name == "TopicSummaryTask":
                queries.append(random.choice(example_theme_asks))

        prompt = queries[0]
        for q in queries[1:]:
            prompt += f". {q}"
    else:
        tasks = []
        for tool_name in tool_names:
            if tool_name == "DocumentSummaryTask":
                tasks.append("summarization")
            elif tool_name == "DocumentNERTask":
                tasks.append("named entity recognition (ner)")
            elif tool_name == "DocumentSentimentTask":
                tasks.append("sentiment analysis")
            elif tool_name == "TopicSummaryTask":
                tasks.append("thematic analysis/action planning")

        random.shuffle(tasks)

        if is_doc:
            tasks.insert(0, "translation (if the document is not in English)")

        prompt = f"Tasks: {', '.join(tasks)}"

    return prompt.strip()


# When building examples, I use these weights to ensure some tasks are more frequently asked for than
# others. I'm doing this because it's my observation that this is what is more likely to be seen at
# inference time.
def get_tool_option_weights(tool_options: list[str]):
    weights = [
        0.1 if i == 1 else 0.9 / (len(tool_options) - 1)
        for i in range(1, len(tool_options) + 1)
    ]
    return weights


# This is the function that will be called to build a sampled dataset. It's designed to give the
# developer the ability to determine how many, and of what type, of data it should get from the
# document datasets created earlier. It also determines a random number of tools to call and
# includes that information in the dataset to be used for evals (e.g., testing that all expected
# tools were called, etc...)
def get_sample(
    docs_df: pd.DataFrame,
    topics_df: pd.DataFrame,
    n_docs: int = 5,
    n_docs_non_english: int = 2,
    n_chunks: int = 5,
    n_topics: int = 5,
    random_state: int | None = None,
):
    random.seed(random_state)

    # Add in some full and chunked documents for single document analysis
    df = docs_df.copy()
    test_df = df[df["AnswerLang"] == "Spanish"].sample(
        n=n_docs_non_english, random_state=random_state
    )
    test_df["AnswerText"] = test_df["AnswerText_NonEnglish"]
    test_df = pd.concat(
        [
            test_df,
            df[df["AnswerLang"] == "English"].sample(
                n=n_docs, random_state=random_state
            ),
        ]
    )
    test_df["_seq_id"] = -1
    test_df["_chunk_id"] = -1
    test_df = test_df.rename(columns={"AnswerText": "_text", "AnswerLang": "_lang"})
    test_df = test_df[["MLVerbatimId", "_seq_id", "_chunk_id", "_text", "_lang"]]

    chunk_df = df.sample(n=n_chunks, random_state=random_state)
    chunk_df = chunk_df.rename(columns={"_chunk": "_text", "AnswerLang": "_lang"})
    chunk_df["_lang"] = "English"
    chunk_df = chunk_df[["MLVerbatimId", "_seq_id", "_chunk_id", "_text", "_lang"]]

    test_df = pd.concat([test_df, chunk_df])
    test_df = test_df.drop_duplicates(subset=["_text"], keep="first")

    # Get a list as of dicts
    test_data_d = test_df.to_dict(orient="records")

    # Randomize the tools and the number of tools being called for variety
    tool_options = ["DocumentSummaryTask", "DocumentNERTask", "DocumentSentimentTask"]
    tool_option_weights = get_tool_option_weights(tool_options)
    for example in test_data_d:
        tools_to_call = []
        if example["_lang"] != "English":
            tools_to_call.append("TranslationTask")

        n_tools = random.choices(
            range(1, len(tool_options) + 1), weights=tool_option_weights, k=1
        )[0]
        tools_to_call += random.sample(tool_options, n_tools)

        if len(example["_text"]) >= 940:
            # For longer documents, vary the inclusion of 'TopicSummary'
            if random.choice([True, False]):
                tools_to_call.append("TopicSummaryTask")

        example["_n_tools"] = len(tools_to_call)
        example["_tools"] = tools_to_call
        example["ask"] = build_hypothetical_prompt(tools_to_call, is_doc=True)

    # Add in some summaries for related document analsysis
    df = topics_df.copy()
    test_df = df[df["pred_theme_id"] != -1].sample(
        n=n_topics, random_state=random_state
    )
    test_df = test_df.rename(columns={"_chunk": "_text"})
    test_df = test_df[["pred_theme_id", "_text"]]

    test_df = test_df.drop_duplicates(subset=["_text"], keep="first")

    # Get a list as of dicts
    topics_test_data_d = test_df.to_dict(orient="records")

    # Randomize the tools and the number of tools being called for variety
    for example in topics_test_data_d:
        tools_to_call = ["TopicSummaryTask"]
        # Randomly throw in the sentiment task 'DocumentSentimentTask'
        if random.choice([True, False]):
            tools_to_call.append("DocumentSentimentTask")

        example["_n_tools"] = len(tools_to_call)
        example["_tools"] = tools_to_call
        example["ask"] = build_hypothetical_prompt(tools_to_call, is_doc=False)

    return test_data_d + topics_test_data_d

Build Evals with Braintrust

To get started, log into braintrust, click the dropdown in the top-left (will have your username in there) and select “+ Create Project”. Give that project a name (e.g., ftcourse_project) and you’re good to go. You’ll need to create an API key and add it to your environmental variables to interact with that project in your code.

docs_df = pd.read_parquet(
    f"{DATA_DIR}/clean/{DATA_FILENAME}_sample_14k_chunked.parquet"
)
print(len(docs_df))
# 18955

topics_df = pd.read_parquet(
    f"{DATA_DIR}/clean/{DATA_FILENAME}_sample_14k_topics.parquet"
)
print(len(topics_df))
# 976

test_data = get_sample(
    docs_df,
    topics_df,
    n_docs=3,
    n_docs_non_english=2,
    n_chunks=5,
    n_topics=5,
    random_state=None,
)
print(len(test_data))
# 15

r = test_data[-1]
r

{'pred_theme_id': 159,
 '_text': array([<a list of strings>],dtype=object),
 '_n_tools': 1,
 '_tools': ['TopicSummaryTask'],
 'ask': 'Based on the documents provided, extract their theme and an action plan based on them'}

@braintrust.traced
def ask_ai(
    client,
    is_doc: bool,
    system_msg: str,
    human_msg: str,
    model: str = "gpt-4o",
    instructor_kwargs: dict = {},
) -> DocumentAnalysis | RelatedDocumentAnalysis:
    try:
        return client.chat.completions.create(
            model=model,
            response_model=DocumentAnalysis if is_doc else RelatedDocumentAnalysis,
            max_retries=3,
            messages=[
                {
                    "role": "system",
                    "content": system_msg,
                },
                {
                    "role": "user",
                    "content": human_msg,
                },
            ],
            **instructor_kwargs,
        )
    except Exception as e:
        print(e)

def convert_analysis_call_to_json(
    analysis_obj: DocumentAnalysis | RelatedDocumentAnalysis,
):
    j = [
        {"function_call": task.__class__.__name__, "args": task.model_dump()}
        for task in analysis_obj.tasks
    ]
    return j

@braintrust.traced
async def generate_tool_calls(
    item_d: dict,
    ask_attr: str = "ask",
    text_attr: str = "_text",
    tools_attr: str = "_tools",
    model_vender: str = "openai",
    model_name: str = "gpt-4o",
):
    is_doc = isinstance(item_d[text_attr], str)

    # Prep messages and defined required tools
    required_tools = item_d[tools_attr]
    ask = item_d[ask_attr]
    doc = item_d[text_attr] if is_doc else "\n\n".join(item_d[text_attr])

    if is_doc:
        system_msg = dedent("""\
            Execute each analysis task.
            Always translate any non-English documents into English before executing other tasks.""")
        human_msg = f"{ask}. Document: {doc}"
    else:
        system_msg = dedent("""\
            Execute each analysis task.""")
        human_msg = f"{ask}. Documents:\n{doc}"

    # Generate the response
    tool_calls = None
    error = None
    try:
        if model_vender == "openai":
            client = instructor.from_openai(OpenAI())
            results = ask_ai(
                client,
                is_doc,
                system_msg,
                human_msg,
                instructor_kwargs={"temperature": 0.0},
            )

        elif model_vender == "anthropic":
            client = instructor.from_anthropic(Anthropic())
            results = ask_ai(
                client,
                is_doc,
                system_msg,
                human_msg,
                model="claude-3-5-sonnet-20240620",
                instructor_kwargs={"max_tokens": 1024, "temperature": 0.0},
            )

        tool_calls = convert_analysis_call_to_json(results)
    except Exception as e:
        error = e

    return {
        "prompt": ask,
        "required_tools": required_tools,
        "tool_calls": tool_calls or [],
        "error": error,
    }

output = await generate_tool_calls(
    item_d=r
)  # , model_vender="anthropic", model_name="claude-3-5-sonnet-20240620")
output

{
    "prompt": "Based on the documents provided, extract their theme and an action plan based on them",
    "required_tools": ["TopicSummaryTask"],
    "tool_calls": [
        {
            "function_call": "TopicSummaryTask",
            "args": {
                "theme_name": "Campus Recreation and Fitness Programs",
                "action_plan": [
                    "Promote the variety of fitness classes available to all levels.",
                    "Highlight the benefits of free and low-cost classes to students.",
                    "Encourage participation in 'try it before you buy it' promotional weeks.",
                    "Emphasize the positive impact of fitness classes on physical and mental well-being.",
                    "Showcase the ease of registration and the supportive staff.",
                ],
            },
        }
    ],
    "error": None,
}

def overall_no_error(output):
    return output["error"] is None


def overall_is_valid_json(output):
    try:
        json.loads(json.dumps(output["tool_calls"]))
        return True
    except ValueError:
        return False


def overall_called_all_tools_acc(output):
    required_tools = output["required_tools"]

    tool_calls = output["tool_calls"]
    actual_tools_used = (
        [item["function_call"] for item in tool_calls] if tool_calls else []
    )

    correct_matches = len(set(required_tools).intersection(set(actual_tools_used)))
    total_possible_matches = len(required_tools)

    accuracy = correct_matches / total_possible_matches
    return accuracy

print(overall_no_error(output))
print(overall_is_valid_json(output))
print(overall_called_all_tools_acc(output))

True
True
1.0

await Eval(
    EVAL_PROJECT_NAME,
    experiment_name="it_000: Initial Dataset",
    data=[
        {
            "input": d,
            "metadata": {
                "model_vendor": "openai",
                "model_name": "gpt-4o",
                "is_single_doc": isinstance(d["_text"], str),
            },
        }
        for d in test_data
    ],
    task=partial(
        generate_tool_calls
    ),  # , model_vender="anthropic", model_name="claude-3-sonnet-20240229"),
    scores=[overall_no_error, overall_is_valid_json, overall_called_all_tools_acc],
)

=========================SUMMARY=========================
100.00% 'overall_called_all_tools_acc' score
100.00% 'overall_is_valid_json' score
100.00% 'overall_no_error' score

29.03s duration

See results for it_000: Initial Dataset at <link to experiment here>

Review Results

After your evals finish running, you can click on the summary link and/or just navigate to the “Experiments” tab of your project in BrainTrust to see the results. Out initial iteration above looks like this:

You can see for our initial experiment against 15 examples, our model performed very well with perfect scores across our three metrics. Clicking on our experiment name will allows us to look at how each example did.

From here we can look at specific examples by clicking on any one of them.

Here we can see both the input provided to the task and it’s output. Notice how each of the methods we decorated with the @braintrust.traced got included here in our trace. This allows us to look at each individually if we want too.

What I’m going to do is select all 15 examples, and add them to a dataset I named “golden dataset”. You can do that pretty easily as pictured below.

Navigating to the “Datasets” tab, I can click on my “golden dataset” and see the rows I just added. I decided to review each one of these examples and correct their outputs by updating the “Expected” section for each.

In the configuration sectoin of Braintrust I created a “reviewed-wg” tag I can apply to any rows in a dataset that I have personally reviewed/corrected. I can use this to know what examples need review and also for filtering out non-reviewed rows in this dataset when I use it in subsequent evals.

Takeaways

We’ve only just begun, but hopefully you can see that Braintrust provides a comprehensive platform for everything we need to build out our evaluation pipeline! It provides …

1. A mechanism to evaluate performance via scoring functions.

2. A mechanism to look at our data and see where things are working well or not so well overall and with specific examples via the Braintrust UI where we can look at overall runs and example-specific metrics.

3. The ability to make modifications to our system and measure how well they improved things or not via subsequent experiments and the ability to look at improvements and regressions on a case-by-case basis.

4. The ability to curate datasets that can be used for better evals and fine tuning via the ability to add the results of experiments to datasets that can be reviewed, corrected, and used for more complex evals and fine tunes

I’m a big fan of the platform and encourage y’all to watch Ankur’s talk from the workshop included below to learn more (I’ve already watched it 3x).

Next Steps

Initially, I was planning on showing y’all a couple more iterations in the eval pipeline before I realized this post is already long as hell. Braintrust is an awesome tool and I want to do it justice, and so my next blog post will cover two more experiments I ran before moving on to generate a dataset for fine tuning. In particular, I’ll cover the following:

• How to merge the “golden dataset” with more sampled data to build up our dataset along with better and more task specific scoring functions.
• How to really “look at the data” and review both improvements and regressions on a case-by-case basis.
• How to improve our prompts and tool definitions based on what we see and then see how things measure up by comparing results to previous runs.

Anyways, I know this was a lot but I hope it encourages you to adopt a “evals first” type of thinking and also to give Braintrust a go. If you made it this far, thanks for reading!

Tools

If you didn’t know this already, apparently “Pydantic is all you need”. There’s even this video by Jason Liu to prove it.

Let’s use it to build out our core tools that we’ll want to always make available for our users.

# Translation
class TranslationTask(BaseModel):
    english_translation: str
    source_language: str

# Document summarization
class DocumentSummaryTask(BaseModel):
    summary: str
    themes: list[str]

# NER
class NamedEntityType(str, Enum):
    PERSON = "PERSON"
    NORP = "NORP"
    FAC = "FAC"
    ORG = "ORG"
    GPE = "GPE"
    LOC = "LOC"
    PRODUCT = "PRODUCT"
    EVENT = "EVENT"
    WORK_OF_ART = "WORK_OF_ART"
    LAW = "LAW"
    LANGUAGE = "LANGUAGE"
    DATE = "DATE"
    TIME = "TIME"
    PERCENT = "PERCENT"
    MONEY = "MONEY"
    QUANTITY = "QUANTITY"
    ORDINAL = "ORDINAL"
    CARDINAL = "CARDINAL"
    OTHER = "OTHER"


class NamedEntity(BaseModel):
    entity_type: NamedEntityType
    entity_mention: str 

class DocumentNERTask(BaseModel):
    named_entities: list[NamedEntity]

# Sentiment
class DocumentSentimentTask(BaseModel):
    positivity: int
    positive_statements: list[str]
    negative_statements: list[str]
    has_suggestions: bool 
    suggestions: list[str]
    feels_threatened: bool
    feels_threatened_examples: list[str]
    profanity: bool
    is_nonsense: bool

# Topic summarization
class TopicSummaryTask(BaseModel):
    theme_name: str
    action_plan: list[str]
```

class DocumentAnalysis(BaseModel):
    tasks: list[TranslationTask | DocumentSummaryTask | DocumentNERTask | DocumentSentimentTask | TopicSummaryTask]

class RelatedDocumentAnalysis(BaseModel):
    tasks: list[TopicSummaryTask | DocumentSentimentTask | TopicSummaryTask]

Single Document Analysis

df = pd.read_parquet(f"{DATA_DIR}/clean/{DATA_FILENAME}_sample_14k_chunked.parquet")

test_df = df[df["AnswerLang"] == "Spanish"].sample(n=2)
test_df["AnswerText"] = test_df["AnswerText_NonEnglish"]
test_df = pd.concat([test_df, df[df["AnswerLang"] == "English"].sample(n=3)])

def ask_ai(
    client,
    content: str,
    query: str | None = None,
    model: str = "gpt-4o",
    instructor_kwargs: dict = {},
) -> DocumentAnalysis:
    if not query:
        query = "Tasks: translation (if the document is not in English), summarization, ner, sentiment analysis."
    try:
        return client.chat.completions.create(
            model=model,
            response_model=DocumentAnalysis,
            max_retries=3,
            messages=[
                {
                    "role": "system",
                    "content": """
Execute each analysis task.
Always translate any non-English documents into English before executing other tasks.""",
                },
                {
                    "role": "user",
                    "content": f"{query}. Document: {content}",
                },
            ],
            **instructor_kwargs,
        )
    except Exception as e:
        print(e)

client = instructor.from_openai(OpenAI())

for r_idx, r in test_df.iterrows():
    print("==========")
    print(":: Document ::")
    print(r["AnswerText"])
    results = ask_ai(client, r["AnswerText"])

    print(":: Results ::")
    print(results.model_dump_json(indent=2))
    print("==========")

client = instructor.from_anthropic(Anthropic())

for r_idx, r in test_df.iterrows():
    print("==========")
    print(":: Document ::")
    print(r["AnswerText"])
    results = ask_ai(
        client, r["AnswerText"], model="claude-3-5-sonnet-20240620", instructor_kwargs={"max_tokens": 1024}
    )

    print(":: Results ::")
    print(results.model_dump_json(indent=2))
    print("==========")

client = instructor.patch(
    OpenAI(
        base_url="https://api.fireworks.ai/inference/v1",
        api_key=os.environ["FIREWORKS_API_KEY"],
    ),
)

for r_idx, r in test_df.iterrows():
    print("==========")
    print(":: Document ::")
    print(r["AnswerText"])
    results = ask_ai(client, r["AnswerText"], model="accounts/fireworks/models/firefunction-v2")

    print(":: Results ::")
    print(results.model_dump_json(indent=2))
    print("==========")

Related Document Analysis

df = pd.read_parquet(f"{DATA_DIR}/clean/{DATA_FILENAME}_sample_14k_topics.parquet")
test_df = df.sample(n=5)

def ask_ai(
    client,
    content: str,
    query: str | None = None,
    model: str = "gpt-4o",
    instructor_kwargs: dict = {},
) -> RelatedDocumentAnalysis:
    if not query:
        query = "Tasks: thematic analysis/action planning and sentiment analysis."
    try:
        return client.chat.completions.create(
            model=model,
            response_model=RelatedDocumentAnalysis,
            max_retries=3,
            messages=[
                {
                    "role": "system",
                    "content": dedent("""\
                        Execute each analysis task."""),
                },
                {
                    "role": "user",
                    "content": f"{query}. Documents:\n{content}",
                },
            ],
            **instructor_kwargs,
        )
    except Exception as e:
        print(e)

client = instructor.from_openai(OpenAI())

for r_idx, r in test_df.iterrows():
    print("==========")
    print(":: Document ::")
    docs = "\n\n".join(r["_chunk"])
    print(docs)

    results = ask_ai(client, docs)

    print(":: Results ::")
    print(results.model_dump_json(indent=2))
    print("==========")

client = instructor.from_anthropic(Anthropic())

for r_idx, r in test_df.iterrows():
    print("==========")
    print(":: Document ::")
    docs = "\n\n".join(r["_chunk"])
    print(docs)

    results = ask_ai(
        client, docs, model="claude-3-5-sonnet-20240620", instructor_kwargs={"max_tokens": 1024}
    )

    print(":: Results ::")
    print(results.model_dump_json(indent=2))
    print("==========")

client = instructor.patch(
    OpenAI(
        base_url="https://api.fireworks.ai/inference/v1",
        api_key=os.environ["FIREWORKS_API_KEY"],
    ),
)

for r_idx, r in test_df.iterrows():
    print("==========")
    print(":: Document ::")
    docs = "\n\n".join(r["_chunk"])
    print(docs)

    results = ask_ai(client, docs, model="accounts/fireworks/models/firefunction-v2")

    print(":: Results ::")
    print(results.model_dump_json(indent=2))
    print("==========")

Exprimenting with Replicate

Replicate is quickly becoming a favorite tool in my LLM arsenal. Hamel has a really nice example/tutorial for getting tool calling to work with the Llama3-70B Instruct model out of the box (not fine tuned). I stole some of his ideas and put together a simple workflow that allows me to incorporate testing with Llama3 variants here.

Below, I share running the single document analysis experiment against both the 70B and 8B variants.

doc_analysis_schema = DocumentAnalysis.model_json_schema()

# doc_analysis_schema

def format_tools():
    doc_analysis_schema = DocumentAnalysis.model_json_schema()

    result = ""
    for k, v in doc_analysis_schema["$defs"].items():
        result += f"<function name='{k}'>\n"
        result += json.dumps(v, indent=2)
        result += f"\n</function name='{k}'>\n"
    return result.strip()

PROMPT_TEMPLATE = """\
<|begin_of_text|><|start_header_id|>system<|end_header_id|>

{system_message}<|eot_id|><|start_header_id|>user<|end_header_id|>

<query>{prompt}</query><|eot_id|><|start_header_id|>assistant<|end_header_id|>

{prefill}"""

# print(PROMPT_TEMPLATE)

SYSTEM_PROMPT = """\
{system_persona}

You are given a set of tasks to perform and a document inside <query> tags and a set of possible functions inside <function-definitions> tags.
Calling these functions are optional. Carefully consider the question and determine if one or more functions can be used to answer the question. Place your thoughts and reasoning behind your decision in <function-thoughts> tags.
Below is a list of function definitions:
<function-definitions>
{tools}
</function-definitions>

For each function you want to call, execute the function and use the answer to provide values to each function parameter in a way that conforms to that function's schema. Include the function name and parameter values inside the <function-calls> tag.
Function calls MUST be in this format: <function-thoughts>Calling func1 would be helpful because of ...</function-thoughts><function-calls>[func1_name(params_name=params_value, params_name2=params_value2...), func2_name(params)]</function-calls>, WITHOUT any answer.
If the query is not in English, always translate it into English first and then proceed to call any other functions using the English translation.

If you do not wish to call any functions, say so in the <function-thoughts> tags followed by <function-calls>None</function-calls>"""

system_persona = """\
You are an expert NLP data scientist, skilled in machine translation, text summarization, NER, thematic analysis, strategic planning, sentiment analysis, and classification tasks."""

# print(SYSTEM_PROMPT)

tools = format_tools()  # .replace("\n", "")
system_message = SYSTEM_PROMPT.format(system_persona=system_persona, tools=tools)

# print(system_message)

for r_idx, r in test_df.iterrows():
    print("==========")
    print(f"Document:\n{r['AnswerText']}\n")

    prompt = f"Translate the document below into English if necessary. After that, perform the following tasks on the document below: summarize, perform named entity recognition, and sentiment analysis. Document:\n{r['AnswerText']}"

    input = {
        "max_tokenx": 1024,
        "temperature": 0,
        "top_p": 0.9,
        "top_k": 50,
        "presence_penalty": 0,
        "frequency_penalty": 0,
        "system_prompt": system_message,
        "prompt": f"<query>{prompt}</query>",
    }
    print("--- Example ---")
    print(f"prompt: {prompt}")
    print("")

    print("===== Llama3-70B-Instruct =====")
    output = replicate.run("meta/meta-llama-3-70b-instruct", input=input)
    print("".join(output))
    print("")

    print("===== Llama3-8B-Instruct =====")
    output = replicate.run("meta/meta-llama-3-8b-instruct", input=input)
    print("".join(output))
    print("")

Takeaways

After all this, I’m confident that my ideas are worthy to pursue. Based on the results above, it seems worthwhile to explore turning the following knobs to boost performance with these models:

1. Improved descriptions in the Pydantic classes

2. Improve the sytem prompt

3. Formulate varied human messages to simulate different ways the user may prompt the model

Can you think of any other adjustments that might be beneficial? If so, drop a comment below!

Next Steps

With some insights about what I can play with to improve my results, it’s time for some experimentation. But how can I know whether or not the changes I’m making present meaningful progress? How can I know specificially where the model is struggling and look at those examples to inform future experiments?

The answer: I need to set up an initial evaluation pipeline along with some scoring functions.

And that is exactly what we’ll do in this next post in this series … stay tuned.

Objective Definition

For this project I’ll be building a system where a user can ask an AI to perform various NLP tasks on one or more text documents. These documents represent survey comments sourced from several kinds of surveys delievered in higher education institutions (e.g., staff, faculty, student satisfaction and engagement surveys). There are core tasks that I need to perform like machine translation and sentiment analysis, but I’d also like to build a system whereby a user can provide it their own tools focused on NLP and the system figure out how to use them correctly.

Data Refinement

Based on my objectives, I need to chiefly be able to support two uses cases:

1. Tasks that need to be performed on a single document (e.g., translation, summary, NER, sentiment). Traditionally we have only done sentiment on chunks because it is the chunks that are associated to topics and we want to report sentiment by topic. This still needs to happen but it might also be interesting to predict sentiment across the full comment as well.

2. Tasks that need to be performed on a collection of related documents (.e.,g thematic analysis, topic sentiment).

In this section we ask, “How can I curate a set of survey comments and semantically related chunks that are varied enough to support these two use cases across all our survey clients?”

def init_document_preprocessing(df: pd.DataFrame, text_col: str):
    """Remove empty documents and clean up whitespace."""
    df = df.copy()

    # Remove all rows where `text_attr` is Nan
    df = df.dropna(subset=list(set([attr for attr in [text_col] if attr is not None])), how="all")

    # Remove all whitespace from `text_attr`, `language_attr`, `non_english_text_attr`
    for attr in list(set([text_col, "AnswerLang"])):
        if attr is not None:
            df[attr] = df[attr].apply(lambda v: str(v).strip())

    return df

df = init_document_preprocessing(raw_df, text_col="AnswerText")

def remove_duplicate_documents(df: pd.DataFrame, text_col: str):
    """Eliminate duplicate documents."""
    df = df.copy()
    unique_attrs = [text_col]

    # Remove duplicate rows (keeping the first instance)
    return df.drop_duplicates(subset=unique_attrs).copy()

df = remove_duplicate_documents(verbatims_df, text_col="AnswerText")

def remove_duplicate_sentences(text: str):
    """Remove duplicated sentences throughout an entire text sequence."""
    if text is None:
        text = ""

    text = re.sub(r"(?<=[.!?])\s*", " ", text)
    sentences = nltk.sent_tokenize(text)
    unique_sentences = list(dict.fromkeys(sentences))

    # Tokenize the text into sentences and remove duplicates by converting to set and back to list
    return " ".join(unique_sentences)


def remove_duplicate_sentences_in_documents(df: pd.DataFrame, text_col: str):
    df = df.copy()
    df[text_col] = df[text_col].apply(lambda v: remove_duplicate_sentences(v))
    return df

df = remove_duplicate_sentences_in_documents(verbatims_df, text_col="AnswerText")

remove_texts = [
    "none",
    "nothing",
    "not sure",
    "no suggestions",
    "no comment",
    "no comments",
    "nothing to add",
    "n/a",
    "none at this time",
    "no comments at all",
    "see previous comment",
    "not really",
]

def is_bad_value(s, remove_texts=[]):
    is_bad = False
    for rt in remove_texts:
        is_bad = bool(re.match(rf"^{rt.lower().strip()}[!.,;:?]*$", str(s).lower().strip()))
        if is_bad:
            break

    return is_bad


def remove_meaningless_documents(df, text_col: str):
    df = df.copy()
    df[text_col] = df[text_col].apply(lambda v: str(v) if v is not None and not is_bad_value(v, remove_texts) else "")
    return df


def remove_short_documents(df: pd.DataFrame, text_col: str, min_text_length: int = 5):
    df = df.copy()
    df = df[pd.notna(df[text_col]) & (df[text_col].str.len() >= min_text_length)]
    return df

df = remove_meaningless_documents(verbatims_df, text_col="AnswerText")
df = remove_short_documents(verbatims_df, text_col="AnswerText")

df.to_parquet(f"{DATA_DIR}/clean/documents_all.parquet")

filtered_df = df[(df["BenchmarkSurveyType"].str.startswith("CSS-"))]

# Calculate the number of samples for each BenchmarkSurveyType
sample_sizes = filtered_df["BenchmarkSurveyType"].value_counts(normalize=True) * 10000
sample_sizes = sample_sizes.round().astype(int)

# Sample records proportionally with a bias towards longer comments
sampled_df = pd.DataFrame()
for survey_type, size in sample_sizes.items():
    subset = filtered_df[filtered_df["BenchmarkSurveyType"] == survey_type]
    weights = subset["AnswerTextCharacterCount"] / subset["AnswerTextCharacterCount"].sum()
    sampled_records = subset.sample(n=size, weights=weights, random_state=42)
    sampled_df = pd.concat([sampled_df, sampled_records])

# Reset index if needed
sampled_df = sampled_df.reset_index(drop=True)

# This gets us down to 10k examples of customer satisfaction like surveys.

sampled_df = pd.concat([sampled_df, df[(df["BenchmarkSurveyType"].str.startswith("SAW"))]])

# Adding in these staff engagement surveys get us to a total of almost 14k examples

class SpanishTranslation(BaseModel):
    """The translation of a document from English to Spanish."""

    spanish_translation: str = Field(..., description="The English tranlsation")

def get_openai_translation_messages(domain: str = "survey comments"):
    system_msg = f"You are a world class translator. Translate the English {domain} below to Spanish. Properly escape strings."
    human_msg = "{input}"

    prompt_msgs = [SystemMessage(content=system_msg), HumanMessagePromptTemplate.from_template(human_msg)]
    return prompt_msgs

llm = ChatOpenAI(model="gpt-4")
messages = get_openai_translation_messages()

prompt = ChatPromptTemplate(messages=messages)
functions = [convert_to_openai_function(SpanishTranslation)]

chain = prompt | llm.bind(function_call={"name": SpanishTranslation.__name__}, functions=functions) | JsonOutputFunctionsParser()

spanish_sample_df = sampled_df.sample(250, random_state=9)

spanish_translations = []
for r_idx, r in spanish_sample_df.iterrows():
    rsp = chain.invoke({"input": str(r["AnswerText"])})
    spanish_translations.append(rsp["spanish_translation"])

sampled_df["AnswerText_NonEnglish"] = None
sampled_df["AnswerLang"] = "English"

spanish_sample_df["AnswerText_NonEnglish"] = spanish_translations
spanish_sample_df["AnswerLang"] = "Spanish"

# Set 'MLVerbatId' as the index for both DataFrames
sampled_df.set_index("MLVerbatimId", inplace=True)
spanish_sample_df.set_index("MLVerbatimId", inplace=True)

# Update df1 with values from df2
sampled_df.update(spanish_sample_df)

# Reset index if needed
sampled_df.reset_index(inplace=True)

df.to_parquet(f"{DATA_DIR}/clean/documents_sample_14k.parquet")

def chunk_paragraphs(text):
    paragraphs = re.split(
        "\n\n|\r\n|\\\\n\\\\n|\\\\r\\\\n|" + r"\\n\s*[-•*o]|\\n\s*\d+[.)]", text
    )  # text.split("\n\n")  # Split text by double line breaks to identify paragraphs

    return [chunk.strip() for chunk in paragraphs if len(chunk.strip()) > 4]

def chunk_texts(df: pd.DataFrame):
    """Chunk the paragraphs keeping any bullet points alongside their context."""
    df = df.copy()

    df.insert(1, "_seq", df["AnswerText"].apply(lambda v: chunk_paragraphs(v)))

    df = df.explode("_seq").reset_index(drop=True)
    df["_seq"] = df["_seq"].str.strip()
    df.insert(2, "_seq_id", df.groupby(["MLVerbatimId"]).cumcount())
    df["_seq_length"] = df["_seq"].str.len()

    return df

def remove_bullet_points(text):
    return re.sub(r"(^|\n)\s*[-•*o]\s*|\n\s*\d+[.)]\s*", " ", text).strip()

chunked_df = chunk_texts(sampled_df)
chunked_df["_seq"] = chunked_df["_seq"].apply(remove_bullet_points)

encoder = HuggingFaceEncoder(name="thenlper/gte-large")

chunker = StatisticalChunker(
    encoder=encoder,
    threshold_adjustment=0.01,
    dynamic_threshold=True,
    window_size=5,  # 5,
    min_split_tokens=23,  # 100,
    max_split_tokens=300,  # 500
    split_tokens_tolerance=10,
    plot_chunks=False,
    enable_statistics=False,
)

def get_semantic_chunks(txt: str, min_chars_to_chunk: int = 90):
    try:
        if len(txt.strip()) < min_chars_to_chunk:
            return [txt.strip()]
        chunks = chunker(docs=[txt])
        return [chunk.content.strip() for chunk in chunks[0]]
    except Exception as e:
        return [txt.strip()]


seqs = chunked_df["_seq"].values.tolist()
chunked_docs = []
for seq in tqdm(seqs):
    chunked_docs.append(get_semantic_chunks(seq))

chunked_df.insert(1, "_chunk", chunked_docs)
chunked_df = chunked_df.explode("_chunk").reset_index(drop=True)
chunked_df["_chunk"] = chunked_df["_chunk"].str.strip()
chunked_df.insert(2, "_chunk_id", chunked_df.groupby(["MLVerbatimId", "_seq_id"]).cumcount())
chunked_df["_chunk_length"] = chunked_df["_chunk"].str.len()

chunked_df.to_parquet(f"{DATA_DIR}/clean/documents_sample_14k_chunked.parquet")

topics_df = chunked_df.copy()

cols = ["pred_theme_id", "pred_orig_theme_name"]

topics_df["pred_theme_id"] = topics_df["pred_theme_id"].astype(int)
topics_df["_chunk"] = topics_df["_chunk"].astype(str)

summaries_df = topics_df.groupby(by=cols)["_chunk"].agg(list).reset_index()

summaries_df.to_parquet(f"{DATA_DIR}/clean/documents_sample_14k_topics.parquet")

Next Steps

With some high quality context, we’ll move on to getting a “vibe check” for how likely what we want to build is possible by seeing what kinds of structured outputs we get when using the big dogs. In particular we’ll be running some tests with OpenAI, Anthropic, Fireworks, and Replicate to develop a good intutition of how well to expect things might work with both closed and open source models.

Again, I’m using Hamel’s ftcourse repo as a general guide for building this project out so make sure to check it out as y’all start your own journeys. The topic in this blog post isn’t necessarily covered in any of his notebooks so consider this the 00 notebook that is more or less implied in the course.

Also, I welcome any ideas on improving anything and everything presented above. Especially if you notice any really egregious and glaring errors in my thinking or workflow, I’d defintely appreciate your thoughts. Either way, thanks for reading to the end :)

v0: Using Enum

This is the initial Enum and pydantic classes I started with. It works pretty damn well with OpenAI’s GPT-4o but fails spectacularly when using any of the Anthopic models.

class EntityGroup(str, Enum):
    """A named entity type."""

    PERSON = "PERSON"
    ORGANIZATION = "ORGANIZATION"
    LOCATION = "LOCATION"
    DATE = "DATE"
    TIME = "TIME"
    PERCENT = "PERCENT"
    MONEY = "MONEY"
    QUANTITY = "QUANTITY"
    ORDINAL = "ORDINAL"
    CARDINAL = "CARDINAL"
    EMAIL = "EMAIL"
    PHONE_NUMBER = "PHONE_NUMBER"
    CREDIT_CARD_NUMBER = "CREDIT_CARD_NUMBER"
    SSN = "SSN"


class NamedEntity(BaseModel):
    entity_group: EntityGroup
    word: str


class DocumentNER(BaseModel):
    named_entities: list[NamedEntity]

When using the Anthropic models, I would consistently see it trying to set entity_group to a string rather than a proper enum value from the EntityGroup enum.

After iterating through a number of prompt and class/field description modifications, I decided to give up and replace my Enum with a Literal. And guess what, everything worked great across all model vendors.

I also decided to lookup the named entities used in Spacy and use those names in my Enum as it makes sense to me that perhaps these libraries might have been included in the training of these LLMs and so maybe will help it do a better job of finding the entities I care about.

v1: Using Literal

Using the Literal type fixed everything and works great across all models! Here’s what it looks like:

class NamedEntity(BaseModel):
    """A named entity found in a document."""

    entity_type: Literal[
        "PERSON",
        "NORP",
        "FAC",
        "ORG",
        "GPE",
        "LOC",
        "PRODUCT",
        "EVENT",
        "WORK_OF_ART",
        "LAW",
        "LANGUAGE",
        "DATE",
        "TIME",
        "PERCENT",
        "MONEY",
        "QUANTITY",
        "ORDINAL",
        "CARDINAL",
        "OTHER",
    ]
    entity: str


class DocumentNERTask(BaseModel):
    named_entities: list[NamedEntity]

This works great … but I really wanted to use an Enum for the reasons listed at the top of this post. And as I’m the kinda guy who enjoys fighting with CUDA installs on his local DL rig, I decided to give it a go after taking a few hours off to enjoy the Euros and Copa America tourneys (also Germany should have won; that was a handball but nah, I’m not angry, nope, not bent at all).

v2: Using Enum Revisted

Here’s the TL;DR version of the code. This version is working fabulously across all APIs and I have yet to encounter a single exception involving Instructor being unable to assign a valid value from the Enum.

class NamedEntityType(str, Enum):
    """Valid types of named entities to extract."""

    PERSON = "PERSON"
    NORP = "NORP"
    FAC = "FAC"
    ORG = "ORG"
    GPE = "GPE"
    LOC = "LOC"
    PRODUCT = "PRODUCT"
    EVENT = "EVENT"
    WORK_OF_ART = "WORK_OF_ART"
    LAW = "LAW"
    LANGUAGE = "LANGUAGE"
    DATE = "DATE"
    TIME = "TIME"
    PERCENT = "PERCENT"
    MONEY = "MONEY"
    QUANTITY = "QUANTITY"
    ORDINAL = "ORDINAL"
    CARDINAL = "CARDINAL"
    OTHER = "OTHER"


class NamedEntity(BaseModel):
    """A named entity result."""

    def convert_str_to_named_entity_type(v: str | NamedEntityType) -> NamedEntityType:
        """Ensure entity type is a valid enum."""
        if isinstance(v, NamedEntityType):
            return v
        else:
            try:
                return NamedEntityType(v)
            except ValueError:
                return NamedEntityType.OTHER

    entity_type: Annotated[str, BeforeValidator(convert_str_to_named_entity_type)]
    entity_mention: str = Field(..., description="The named entity recognized.")


class DocumentNERTask(BaseModel):
    named_entities: list[NamedEntity]

Besides the return of the Enum, the most noticeable change involves the inclusion of a BeforeValidator that ensures the value is assigned to a valid enum as defined in NamedEntity. In cases where it wants to add an entity to the list of named_entities that isn’t defined in the NamedEntityType enum or is named differently (e.g., “ORGANIZATION” vs. “ORG”), it will assign them to OTHER.

With this in place, I now have a solution that is:

1. More resiliant

2. Can be used in debugging named entity recogintion (e.g, I can explore what named entities might be missing from the Enum or getting named differently by looking at those that get associated with the OTHER value)

3. I can use that same beautiful Enum across all parts of my application

v2.0.1: Using Enum and fuzzywuzzy

A suggestion from a Twitter user inspired me to enhance our approach by implementing similarity-based matching rather than relying on exact matches. To make it so, I installed the fuzzywuzzy library and made the necessary modifications to increase the likelihood of delivering high-quality results.

class NamedEntityType(str, Enum):
    """Valid types of named entities to extract."""

    PERSON = "PERSON"
    NORP = "NORP"
    FAC = "FAC"
    ORG = "ORG"
    GPE = "GPE"
    LOC = "LOC"
    PRODUCT = "PRODUCT"
    EVENT = "EVENT"
    WORK_OF_ART = "WORK_OF_ART"
    LAW = "LAW"
    LANGUAGE = "LANGUAGE"
    DATE = "DATE"
    TIME = "TIME"
    PERCENT = "PERCENT"
    MONEY = "MONEY"
    QUANTITY = "QUANTITY"
    ORDINAL = "ORDINAL"
    CARDINAL = "CARDINAL"
    OTHER = "OTHER"


class NamedEntity(BaseModel):
    """A named entity result."""

    def convert_str_to_named_entity_type(v: str | NamedEntityType) -> NamedEntityType:
        """Ensure entity type is a valid enum."""
        if isinstance(v, NamedEntityType):
            return v
        else:
            try:
                match, score = fuzzy_process.extractOne(v.upper(), [e.value for e in list(NamedEntityType)])
                return NamedEntityType(match) if score >= 60 else NamedEntityType.OTHER
            except ValueError:
                return NamedEntityType.OTHER

    entity_type: Annotated[str, BeforeValidator(convert_str_to_named_entity_type)]
    entity_mention: str = Field(..., description="The named entity recognized.")


class DocumentNERTask(BaseModel):
    named_entities: list[NamedEntity]

This improves those cases where, for example, the LLM wants to define the entity type as “ORGANIZATION” but it is defined in the Enum as “ORG”.

Another option potentially worth exploring is to use the llm_validator function to make a call out to the LLM when exceptions happen and prompt it to coerce the value into something in the Enum. This could hike up your costs a bit but I imagine using a cheap model like GPT-3.5-Turbo could do the job just fine, and would likely you give an addtional robustness in quality results.

Conclusion

That’s it.

If you found this helpful and/or have suggestions on how to improve the use of Enums in Instructor, lmk in the comments below or on X. Until then, time to enjoy some football and see if Brazil can make it into the semis.

How to take this course

Here I want to share 4 tips to getting the most of this course.

Key Takeaways

I’ve learned a lot to say the least. Below are some of the key takeways I’ve taken away from the course that I hope will aid your navigation and success with it as well.

Next Steps

Build something.

Following Hamel’s course project repo as a guide, I’m going to explore improving a ML pipeline I built for work that allows users to upload a collection of survey comments, interview answers, etc… and have a variety of NLP analysis performed on the documents collectively and individually (eg., translation, sentiment, summarization, NER, and thematic analysis). The data and responses are styled in a way familiar to faculty, staff, and students in higher ed.

The current system makes a call out to gpt-4 for every document and for every collection of related documents to perform each task. It can get pricey and can take a long time depending on the number of documents being processed. I did this because I couldn’t get reliable results from asking the LLM to perform a number of tasks, whether the document/s where short or long.

For this course project, I’m going to start from the beginning and try to improve things by:

1. Curating a dataset from both generated data and local traces I have from the existing pipeline

2. Follow and “eval first” development structure, so I can evaluate the proprietary LLMs (e.g., gpt-4, gpt-4o, gpt-3.5-turbo, etc..) I’m using right now.

3. Use tool calling and more complex pydantic models to reduce the number of calls I need to make and lower costs.

4. Finetune a variety of smaller models on a curated datset to see if I can improve the quality of my predictions while also lowering costs and latency even more.

If folks have any suggestions or ideas to improve this approach, please let me know in the comments or on twitter. I’d especially love to hear from people who have gone down a similar path and what worked (or didn’t) for them.

Tabular Modeling

What is it?

“Tabular modeling takes data in the form of a table (like a spreadsheet or CSV). The objective is to predict the value of one column based on the values in the other columns.” Tabular data is also called “structured data” while “unstructured data” represents things like text, images, audio, etc…

Why is it important?

Though it is reported that 80-90% of data is unstructured (think images, text, audio), ironically, it appears that the vast majority of “real world” machine learning is concerened with tabular/structured data.

If you are just starting out with Data Science you should know that still the vast majority of DS problems in the industry concern structured/tabular data. This is what you should focus on in order to make a professional inroad.

1/3

— Bojan Tunguz (@tunguz) February 23, 2022

What are they?

1. For structured data, ensembles of decision trees (e.g., random forests and gradient boosting machines like XGBoost).

2. For unstructured data, multilayered neural networks learned with SGD.

See pages 282-284 for more discussion on the pros/cons of decision trees and neural networks for tabular data.

Categorical Embeddings

Continuous v. Categorical

Continuous variables “contain a real numbers that be fed into a model directly and are meaningful in and out of themselves. Examples include”age” and “price”.

Categorical variables “contain a number of discrete levels, such as ‘movie ID,’ for which addition and multiplication don’t have any meaning (even if they’re stored as numbers). Other examples include dates, columns indicating”sex”, “gender”, “department”, etc…

How do we represent “categorical” data?

We learned this in chapter 8, we represent such data via entity embeddings.

Because “an embedding layer is exactly equivalent to placing an ordinary linear layer after every one-hot-encoded input layer … the embedding transforms the categorical variables into inputs that are both continuous and meaningful.”

In other words …

Imports

from kaggle import api
from dtreeviz.trees import *
from fastai.tabular.all import *
from IPython.display import Image, display_svg, SVG
from pandas.api.types import is_string_dtype, is_numeric_dtype, is_categorical_dtype
from sklearn.ensemble import RandomForestRegressor
from sklearn.tree import DecisionTreeRegressor, export_graphviz

pd.options.display.max_rows = 20
pd.options.display.max_columns = 8

Data preparation

path = Path("bluebook")
path

Path('bluebook')

if not path.exists():
    path.mkdir()
    api.competition_download_cli("bluebook-for-bulldozers", path=path)
    file_extract(path / "bluebook-for-bulldozers.zip")

path.ls(file_type="text")

Downloading bluebook-for-bulldozers.zip to bluebook

100%|██████████| 48.4M/48.4M [00:00<00:00, 112MB/s] 

(#7) [Path('bluebook/Machine_Appendix.csv'),Path('bluebook/Test.csv'),Path('bluebook/ValidSolution.csv'),Path('bluebook/median_benchmark.csv'),Path('bluebook/TrainAndValid.csv'),Path('bluebook/Valid.csv'),Path('bluebook/random_forest_benchmark_test.csv')]

train_df = pd.read_csv(path / "TrainAndValid.csv", low_memory=False)
test_df = pd.read_csv(path / "Test.csv", low_memory=False)

train_df.columns

Index(['SalesID', 'SalePrice', 'MachineID', 'ModelID', 'datasource',
       'auctioneerID', 'YearMade', 'MachineHoursCurrentMeter', 'UsageBand',
       'saledate', 'fiModelDesc', 'fiBaseModel', 'fiSecondaryDesc',
       'fiModelSeries', 'fiModelDescriptor', 'ProductSize',
       'fiProductClassDesc', 'state', 'ProductGroup', 'ProductGroupDesc',
       'Drive_System', 'Enclosure', 'Forks', 'Pad_Type', 'Ride_Control',
       'Stick', 'Transmission', 'Turbocharged', 'Blade_Extension',
       'Blade_Width', 'Enclosure_Type', 'Engine_Horsepower', 'Hydraulics',
       'Pushblock', 'Ripper', 'Scarifier', 'Tip_Control', 'Tire_Size',
       'Coupler', 'Coupler_System', 'Grouser_Tracks', 'Hydraulics_Flow',
       'Track_Type', 'Undercarriage_Pad_Width', 'Stick_Length', 'Thumb',
       'Pattern_Changer', 'Grouser_Type', 'Backhoe_Mounting', 'Blade_Type',
       'Travel_Controls', 'Differential_Type', 'Steering_Controls'],
      dtype='object')

train_df.describe().T

advanced_describe(train_df)

train_df.ProductSize.unique()

array([nan, 'Medium', 'Small', 'Large / Medium', 'Mini', 'Large',
       'Compact'], dtype=object)

sizes = ["Large", "Large / Medium", "Medium", "Small", "Mini", "Compact"]
train_df.ProductSize = train_df.ProductSize.astype("category")
train_df.ProductSize = train_df.ProductSize.cat.set_categories(
    sizes, ordered=True
)  # note: "inplace=True" is depreciated as of 1.30
train_df.ProductSize.unique()

[NaN, 'Medium', 'Small', 'Large / Medium', 'Mini', 'Large', 'Compact']
Categories (6, object): ['Large' < 'Large / Medium' < 'Medium' < 'Small' < 'Mini' < 'Compact']

dep_var = "SalePrice"
train_df[dep_var] = np.log(train_df[dep_var])

train_df = add_datepart(train_df, "saledate")
test_df = add_datepart(test_df, "saledate")

[col for col in train_df.columns if col.startswith("sale")]

['saleYear',
 'saleMonth',
 'saleWeek',
 'saleDay',
 'saleDayofweek',
 'saleDayofyear',
 'saleIs_month_end',
 'saleIs_month_start',
 'saleIs_quarter_end',
 'saleIs_quarter_start',
 'saleIs_year_end',
 'saleIs_year_start',
 'saleElapsed']

procs = [Categorify, FillMissing]

Creating our TabularPandas

cont, cat = cont_cat_split(train_df, 1, dep_var=dep_var)

cond = (train_df.saleYear < 2011) | (train_df.saleMonth < 10)
train_idxs = np.where(cond)[0]
valid_idxs = np.where(~cond)[0]

splits = (list(train_idxs), list(valid_idxs))

to = TabularPandas(train_df, procs, cat, cont, y_names=dep_var, splits=splits)
type(to)

fastai.tabular.core.TabularPandas

len(to.train), len(to.valid)

(404710, 7988)

to.show(3)

to.items.head(3)

to.classes["ProductSize"]

['#na#', 'Large', 'Large / Medium', 'Medium', 'Small', 'Mini', 'Compact']

# save to filesystem
save_pickle(path / "to.pkl", to)

# load from filesystem
to = load_pickle(path / "to.pkl")

Approach 1: Decision Trees

“A decision tree asks a series of binary (yes or no) questions about the data. After each question, the data at that part of the tree is split between a Yes and a No branch…. After one or more questions, either a prediction can be made on the basis of all previous answers or another question is required.”

“… for regression, we take the target mean of the items in the group”

train_xs, train_y = to.train.xs, to.train.y
valid_xs, valid_y = to.valid.xs, to.valid.y

m = DecisionTreeRegressor(max_leaf_nodes=4)
m.fit(train_xs, train_y)

DecisionTreeRegressor(max_leaf_nodes=4)

draw_tree(m, train_xs, size=7, leaves_parallel=True, precision=2)

samp_idx = np.random.permutation(len(train_y))[:500]

dtreeviz(
    m,
    train_xs.iloc[samp_idx],
    train_y.iloc[samp_idx],
    train_xs.columns,
    dep_var,
    fontname="DejaVu Sans",
    scale=1.6,
    label_fontsize=10,
    orientation="LR",
)

/usr/local/lib/python3.7/dist-packages/sklearn/base.py:451: UserWarning: X does not have valid feature names, but DecisionTreeRegressor was fitted with feature names
  "X does not have valid feature names, but"

train_xs.loc[train_xs["YearMade"] < 1900, "YearMade"] = 1950
valid_xs.loc[valid_xs["YearMade"] < 1900, "YearMade"] = 1950

samp_idx = np.random.permutation(len(train_y))[:500]

dtreeviz(
    m,
    train_xs.iloc[samp_idx],
    train_y.iloc[samp_idx],
    train_xs.columns,
    dep_var,
    fontname="DejaVu Sans",
    scale=1.6,
    label_fontsize=10,
    orientation="LR",
)

/usr/local/lib/python3.7/dist-packages/sklearn/base.py:451: UserWarning: X does not have valid feature names, but DecisionTreeRegressor was fitted with feature names
  "X does not have valid feature names, but"

m = DecisionTreeRegressor()
m.fit(train_xs, train_y)

DecisionTreeRegressor()

def r_mse(preds, targs):
    return round(math.sqrt(((preds - targs) ** 2).mean()), 6)


def m_rmse(m, xs, y):
    return r_mse(m.predict(xs), y)

m_rmse(m, train_xs, train_y)

0.0

m_rmse(m, valid_xs, valid_y)

0.333111

m.get_n_leaves(), len(train_xs)

(324559, 404710)

m = DecisionTreeRegressor(min_samples_leaf=25)
m.fit(to.train.xs, to.train.y)

m_rmse(m, to.train.xs, to.train.y), m_rmse(m, to.valid.xs, to.valid.y)

(0.211677, 0.268059)

Approach 2: Random Forests

A random forest “is a model that averages the predictions of a large number of decision trees, which are generated by randomly varying various parameters that specify what data is used to train the tree (what columns and rows are included in each tree) and other tree parameters”

Why does it work so well?

Because it utlizes bagging.

What is “bagging”?

From the “Bagging Predictors” paper … “Bagging predictors is a method for generating multiple versions of a predictor and using these to get an aggregated predictor…. The multiple versions are formed by making bootstrap replicates (a randomly chosen subset of rows) of the learning set.”

This means that we can improve the performance of a model by training it multiple times with a different random subset of the data each time, and then averaging the predictions.

See p.298 for more details on how bagging works.

def fit_rf(
    xs,
    y,
    n_estimators=40,
    max_samples=200_000,
    max_features=0.5,
    min_samples_leaf=5,
    **kwargs,
):
    return RandomForestRegressor(
        n_jobs=-1,  # Tells sklearn to use all our CPUs to build the trees in parallel
        n_estimators=n_estimators,  # The number of trees
        max_samples=max_samples,  # The number of rows to sample for training ea. tree
        max_features=max_features,  # The number of columns to sample at each split
        min_samples_leaf=min_samples_leaf,  # Stop when all leaf nodes have at least this number of samples
        oob_score=True,
    ).fit(xs, y)

m = fit_rf(train_xs, train_y)

m_rmse(m, train_xs, train_y), m_rmse(m, valid_xs, valid_y)

(0.170771, 0.232215)

tree_preds = np.stack(
    [t.predict(valid_xs.values) for t in m.estimators_]
)  # added .values (see: https://stackoverflow.com/a/69378867)

r_mse(tree_preds.mean(0), valid_y)

0.232215

plt.plot([r_mse(tree_preds[: i + 1].mean(0), valid_y) for i in range(40)])

r_mse(m.oob_prediction_, train_y)

0.210601

preds = np.stack([t.predict(valid_xs.values) for t in m.estimators_])
preds.shape  # => (# of trees, # of predictions)

(40, 7988)

preds_std = preds.std(0)  # get rid of first dimension (the trees)

# std for first 5 rows of validation set
preds_std[:5]

array([0.25911739, 0.08550421, 0.11939131, 0.29835108, 0.16490343])

def rf_feature_importance(m, df):
    return pd.DataFrame(
        {"cols": df.columns, "imp": m.feature_importances_}
    ).sort_values("imp", ascending=False)


def plot_fi(fi_df):
    return fi_df.plot("cols", "imp", "barh", figsize=(12, 7), legend=False)

fi_df = rf_feature_importance(m, train_xs)

# Let's look at the 10 most important features
fi_df[:10]

plot_fi(fi_df[:30])

cols_to_keep = fi_df[fi_df.imp > 0.005].cols
len(cols_to_keep)

21

train_xs_keep = train_xs[cols_to_keep]
valid_xs_keep = valid_xs[cols_to_keep]

m = fit_rf(train_xs_keep, train_y)

m_rmse(m, train_xs_keep, train_y), m_rmse(m, valid_xs_keep, valid_y)

(0.181366, 0.231985)

len(train_xs.columns), len(train_xs_keep.columns)

(66, 21)

plot_fi(rf_feature_importance(m, train_xs_keep[:30]))

cluster_columns(train_xs_keep)

def get_oob(df, y):
    m = RandomForestRegressor(
        n_estimators=10,
        min_samples_leaf=15,
        max_samples=50_000,
        max_features=0.5,
        n_jobs=-1,
        oob_score=True,
    )
    m.fit(df, y)
    return m.oob_score_

get_oob(train_xs_keep, train_y)

0.8723047876078355

redundant_cols = [
    "saleYear",
    "saleElapsed",
    "ProductGroupDesc",
    "ProductGroup",
    "fiModelDesc",
    "fiBaseModel",
    "Hydraulics_Flow",
    "Grouser_Tracks",
    "Coupler_System",
]

{c: get_oob(train_xs_keep.drop(c, axis=1), train_y) for c in redundant_cols}

{'Coupler_System': 0.8729675628149272,
 'Grouser_Tracks': 0.8735142839769211,
 'Hydraulics_Flow': 0.8730586366090629,
 'ProductGroup': 0.8731366740450576,
 'ProductGroupDesc': 0.8728712200071638,
 'fiBaseModel': 0.87126582392193,
 'fiModelDesc': 0.8714879359004835,
 'saleElapsed': 0.8681954238416791,
 'saleYear': 0.8713304844511609}

cols_to_drop = ["saleYear", "ProductGroupDesc", "fiBaseModel", "Grouser_Tracks"]

get_oob(train_xs_keep.drop(cols_to_drop, axis=1), train_y)

0.8695950962336326

train_xs_final = train_xs_keep.drop(cols_to_drop, axis=1)
valid_xs_final = valid_xs_keep.drop(cols_to_drop, axis=1)

# save to filesystem
save_pickle(path / "train_xs_final.pkl", train_xs_final)
save_pickle(path / "valid_xs_final.pkl", valid_xs_final)

train_xs_final = load_pickle(path / "train_xs_final.pkl")
valid_xs_final = load_pickle(path / "valid_xs_final.pkl")

m = fit_rf(train_xs_final.values, train_y)
m_rmse(m, train_xs_final.values, train_y), m_rmse(m, valid_xs_final.values, valid_y)

(0.183291, 0.232867)

p = valid_xs_final["ProductSize"].value_counts().sort_index().plot.barh()
c = to.classes["ProductSize"]
plt.yticks(range(len(c)), c)

([<matplotlib.axis.YTick at 0x7f6564704d10>,
  <matplotlib.axis.YTick at 0x7f6562ce8350>,
  <matplotlib.axis.YTick at 0x7f6562b34d10>,
  <matplotlib.axis.YTick at 0x7f65645b61d0>,
  <matplotlib.axis.YTick at 0x7f656473b310>,
  <matplotlib.axis.YTick at 0x7f656473bd90>,
  <matplotlib.axis.YTick at 0x7f656473b250>],
 [Text(0, 0, '#na#'),
  Text(0, 0, 'Large'),
  Text(0, 0, 'Large / Medium'),
  Text(0, 0, 'Medium'),
  Text(0, 0, 'Small'),
  Text(0, 0, 'Mini'),
  Text(0, 0, 'Compact')])

x = valid_xs_final["YearMade"].hist()

from sklearn.inspection import PartialDependenceDisplay

fig, ax = plt.subplots(figsize=(12, 4))
PartialDependenceDisplay.from_estimator(
    m, valid_xs_final, ["YearMade", "ProductSize"], grid_resolution=20, ax=ax
)

from treeinterpreter import treeinterpreter

row = valid_xs_final.iloc[:5]

pred, bias, contributions = treeinterpreter.predict(m, row.values)

pred[0], bias[0], contributions[0].sum()

(array([10.00298786]), 10.104457702350853, -0.10146984568626016)

import waterfall_chart

waterfall_chart.plot(
    valid_xs_final.columns,
    contributions[0],
    threshold=0.08,
    rotation_value=45,
    formatting="{:,.3f}",
)

combined_df = pd.concat([train_xs_final, valid_xs_final])
is_valid = np.array([0] * len(train_xs_final) + [1] * len(valid_xs_final))

ood_m = fit_rf(combined_df, is_valid)

rf_feature_importance(ood_m, combined_df)[:6]

m = fit_rf(train_xs_final, train_y)
print("orig:", m_rmse(m, valid_xs_final, valid_y))

orig: 0.231632

for c in ("SalesID", "saleElapsed", "MachineID"):
    m = fit_rf(train_xs_final.drop(c, axis=1), train_y)
    print(c, m_rmse(m, valid_xs_final.drop(c, axis=1), valid_y))

SalesID 0.231139
saleElapsed 0.235471
MachineID 0.229936

time_vars = ["SalesID", "MachineID"]
train_xs_final_time = train_xs_final.drop(time_vars, axis=1)
valid_xs_final_time = valid_xs_final.drop(time_vars, axis=1)

m = fit_rf(train_xs_final_time, train_y)
m_rmse(m, valid_xs_final_time, valid_y)

0.228524

train_xs["saleYear"].hist()

recent_sales = train_xs["saleYear"] > 2004

train_xs_recent = train_xs_final_time[recent_sales]
train_y = train_y[recent_sales]

m = fit_rf(train_xs_recent, train_y)
m_rmse(m, train_xs_recent, train_y), m_rmse(m, valid_xs_final_time, valid_y)

(0.177533, 0.229519)

Approach 3: Neural Networks

nn_df = pd.read_csv(path / "TrainAndValid.csv", low_memory=False)

sizes = ["Large", "Large / Medium", "Medium", "Small", "Mini", "Compact"]
nn_df.ProductSize = nn_df.ProductSize.astype("category")
nn_df.ProductSize = nn_df.ProductSize.cat.set_categories(
    sizes, ordered=True
)  # note: "inplace=True" is depreciated as of 1.30

nn_df[dep_var] = np.log(nn_df[dep_var])

nn_df = add_datepart(nn_df, "saledate")

nn_df_final = nn_df[list(train_xs_final_time.columns) + [dep_var]]

fastai uses the max_card to determine what columns should be treated as categoricals … “Embedding sizes larger that 10,000 should generaly be used only after you’ve tested whether there are better ways to group the variable”

nn_cont, nn_cat = cont_cat_split(nn_df_final, max_card=9000, dep_var=dep_var)

Therfore, “saleElapsed” needs to be made into a continous variable … which it alread is

nn_cont

['saleElapsed']

How to find the cardinality of each categorical?

nn_df_final[nn_cat].nunique()

YearMade                73
ProductSize              6
Coupler_System           2
fiProductClassDesc      74
Hydraulics_Flow          3
ModelID               5281
fiSecondaryDesc        177
fiModelDesc           5059
Enclosure                6
Hydraulics              12
ProductGroup             6
fiModelDescriptor      140
Tire_Size               17
Drive_System             4
dtype: int64

“… two variables pertaining to the ‘model’ of the equipment, both with similar very high cardinalities, suggests that they may contain similar, redundant information.”

cols_to_drop = ["fiModelDescriptor"]
train_xs_recent2 = train_xs_recent.drop(cols_to_drop, axis=1)
valid_xs_recent2 = valid_xs_final_time.drop(cols_to_drop, axis=1)

m2 = fit_rf(train_xs_recent2, train_y)
m_rmse(m2, train_xs_recent2, train_y), m_rmse(m2, valid_xs_recent2, valid_y)

(0.176776, 0.230194)

nn_cat.remove("fiModelDescriptor")

print(nn_cont)
print(nn_cat)

['saleElapsed']
['YearMade', 'ProductSize', 'Coupler_System', 'fiProductClassDesc', 'Hydraulics_Flow', 'ModelID', 'fiSecondaryDesc', 'fiModelDesc', 'Enclosure', 'Hydraulics', 'ProductGroup', 'Tire_Size', 'Drive_System']

nn_procs = [Categorify, FillMissing, Normalize]

nn_to = TabularPandas(
    nn_df_final, nn_procs, nn_cat, nn_cont, splits=splits, y_names=dep_var
)

dls = nn_to.dataloaders(bs=1024)

y = nn_to.train.y
y.min(), y.max()

(8.465899467468262, 11.863582611083984)

from fastai.tabular.all import *

learn = tabular_learner(
    dls, y_range=(8, 12), layers=[500, 250], n_out=1, loss_func=F.mse_loss
)

learn.lr_find()

SuggestedLRs(valley=0.00015848931798245758)

learn.fit_one_cycle(5, 1e-2)

preds, targs = learn.get_preds()
r_mse(preds, targs)

0.223821

learn.save("nn")

Path('models/nn.pth')

Approach 4: Boosting

Here “we add models instead of averaging them” (as compared to ensembling). See p.323-24 for all the details, but here’s how it works in general:

1. Train small model that underfits
2. Get predictions for model from #1
3. Get the “residuals” (‘the error for each point in the training set’) by subtracting the predictions from the targets.
4. Go back to step 1, “but instead of using the original targets, use the residuals as the targets for training.
5. Continue stesp 1-4 until you need to stop (e.g., max number of trees, you start overfitting, etc…)

Steps 3-4 are the bread and butter of things.

Some go as far as saying that such models are all you need :)

XGBoost Is All You Need

Deep Neural Networks and Tabular Data: A Surveyhttps://t.co/Z2KsHP3fvp pic.twitter.com/uh5NLS1fVP

— Bojan Tunguz (@tunguz) March 30, 2022

How do you make predictions with an ensemble of boosted trees?

Answer: By calculating the predictions from each tree and then adding them all together.

XGBoost and sklearn’s HistGradientBooting models are legit ML models to try …

Here are my top used ML algos for tabular data problems:

1. XGBoost
2. HistGradientBoosting
3. Logistic regression/ Ridge regression
4. LightGBM
5. MLP
6. A blend of 1. And 5.

— Bojan Tunguz (@tunguz) April 22, 2022

Other things to try

Ensembling (pp.322-23)

This is a form of “bagging” where averaging the predictions of models trained using different algorithms, each of which it is reasonable to expect would produce differnt kinds of errors, would likely product better predictions.

Combining embeddings with other models (pp.324-25)

“… shows that you can get much of the performance improvement of a neural network without having to use a neural network as inference time” by using the embeddings as inputs (which are just array lookups) to small decision tree ensemble!

In summary

“This will give you a strong baseline … can then use that model for feature selection and partial dependence analysis to get a better understanding of your data.”

Pros/Cons of each approach:

Random forests: Pros:

• Easiest to train
• Resiliant to hyperparameter choices
• Require little preprocessing
• Fast to train
• Should not overfit if you have enough trees Cons:
• Can be less accurate (esp. if extrapolation is required “such as predicting future time periods.”

Gradient boosting machines (GBMs): Pros:

• Often more accurate than random forests Cons:
• Sensitive to hyperparmeter choices
• Can overfit

Neural Networks: Pros:

• Can extrapolate well
• Can provide good results. Cons:
• Take longest time to train
• Require extra preprocessing (e.g., normalization)

Resources

1. https://book.fast.ai - The book’s website; it’s updated regularly with new content and recommendations from everything to GPUs to use, how to run things locally and on the cloud, etc…

2. “Hands-On Gradient Boosting with XGBoost and scikit-learn” - Unread personally, but on my bookshelf and recommended by those in the know.

3. tsai - A time series/sequences focused library built on PyTorch and fastai by Ignacio Oguiza

Collaborative Filtering

What is it?

Think recommender systems which “look at which products the current user has used or liked, find other users who have used or liked similar products, and then recommend other products that those users have used or liked.”

The key to making collaborative filtering and tabular models, is the idea of latent factors.

What are “latent factors” and what is the problem they solve?

Remember that models can only work with numbers, and while something like “price” can be used to accurately reflect the value of a house, how do we represent numerically concepts like the day of week, the make/model of a car, or the job function of an employee?

The answer is with latent factors.

In a nutshell, latent factors are numbers associated to a thing (e.g., day of week, model of car, job function, etc…) that are learnt during model training. At the end this process, we have numbers that provide a representation of the thing we can use and explore in a variety of ways. These factors are called “latent” because we don’t know what they are beforehand.

If we had something like this …

… how could we predict what users would rate movies they have yet to see? Let’s take a look.

from fastai.collab import *
from fastai.tabular.all import *

path = untar_data(URLs.ML_100k)

ratings_df = pd.read_csv(path/"u.data", delimiter="\t", header=None, names=["user", "movie", "rating", "timestamp"])
ratings_df.head()

How do we numerically represent user 196 and movie 242? With latent factors we don’t have to know, we can have such a representation learnt using SGD.

How do we set this up?

1. “… randomly initialized some parameters [which] will be a set of latent factors for each user and movie.”

2. “… to calculate our predictions [take] the dot product of each movie with each user.

3. “… to calculate our loss … let’s pick mean squared error for now, since that is one reasonable way to represent the accuracy of a prediction”

With this in place, “we can optimize our parameters (the latent factors) using stochastic gradient descent, such as to minimize the loss.” In a picture, it looks like this …

movies_df = pd.read_csv(path/"u.item", delimiter="|", header=None, names=["movie", "title"], usecols=(0,1), encoding="latin-1")
movies_df.head()

ratings_df = ratings_df.merge(movies_df)
ratings_df.head()

dls = CollabDataLoaders.from_df(ratings_df, item_name="title", user_name="user", rating_name="rating")
dls.show_batch()

So how do we create these latent factors for our users and movies?

“We can represent our movie and user latent factor tables as simple matrices” that we can index into. But as looking up in an index is not something our models know how to do, we need to use a special PyTorch layer that will do this for us (and more efficiently than using a one-hot-encoded, OHE, vector to do the same).

And that layer is called an embedding. It “indexes into a vector using an integer, but has its derivative calcuated in such a way that it is identical to what it would have been if it had done a matric multiplication with a one-hot-encoded vector.”

Collaborative Filtering: From Scratch (dot product)

A dot product approach

n_users = len(dls.classes["user"])
n_movies = len(dls.classes["title"])
n_factors = 5

print(n_users, n_movies, n_factors)

944 1665 5

class DotProduct(Module):
  def __init__(self, n_users, n_movies, n_factors):
      super().__init__()
      self.users_emb = Embedding(n_users, n_factors)
      self.movies_emb = Embedding(n_movies, n_factors)

  def forward(self, inp):
    users = self.users_emb(inp[:,0])
    movies = self.movies_emb(inp[:,1])
    return (users * movies).sum(dim=1)