fix dataset formatting and mapping for Magistral reasoning

[rolandtannous]

Problem

reasoning_dataset.map(...)["conversations"] returns a list of conversations
But apply_chat_template() expects a single conversation (a list of message dicts), not a list of conversations.
the Jinja template tries to iterate over what it thinks are "messages", but they're actually entire conversations (lists). When it tries to access message['role'], it fails because a list doesn't have a 'role' key, causing it to fall through to the else clause that raises the exception:

 raise jinja2.exceptions.TemplateError(message)
jinja2.exceptions.TemplateError: Only user, system, assistant and tool roles are supported in the custom template made by Unsloth

Solution

aligned formats of formatting , mapping functions with what apply_chat_template expects.

def generate_conversation(example):
    problem  = example["problem"]
    solution = example["generated_solution"]
    conversation = [
        {"role" : "user",      "content" : problem},
        {"role" : "assistant", "content" : solution},
    ]
    return { "conversations": conversation }

reasoning_conversations = [
    tokenizer.apply_chat_template(conv["conversations"], tokenize=False)
    for conv in reasoning_dataset.map(generate_conversation)]

[gemini-code-assist]

Summary of Changes

Hello @rolandtannous, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical data formatting incompatibility within the Magistral reasoning notebook. The previous implementation incorrectly passed a list of conversations to apply_chat_template, which expects a single conversation. The changes ensure that the dataset mapping and template application functions are aligned, preventing a TemplateError and allowing the conversational data to be processed as intended.

Highlights

• Dataset Mapping Correction: The generate_conversation function was refactored to correctly process a single example and return a single conversation (a list of message dictionaries), aligning its output with the expectations of apply_chat_template.
• Template Application Alignment: The call to tokenizer.apply_chat_template was updated to iterate over individual conversations generated by the dataset mapping, resolving a jinja2.exceptions.TemplateError that occurred when a list of conversations was incorrectly passed to the function.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request correctly fixes the data format issue when calling tokenizer.apply_chat_template. The previous implementation passed a list of conversations, causing a TemplateError, and the new implementation correctly maps over each conversation individually. The changes in both the Jupyter notebook and the corresponding Python script are correct. I've added a couple of suggestions to improve code readability, mainly around formatting a list comprehension.

[gemini-code-assist]

The list comprehension is a bit difficult to read due to its formatting. The arguments to apply_chat_template are not clearly separated, and the for loop is on a separate line which is unusual for a list comprehension of this style. For better readability and adherence to common Python style (PEP 8), I suggest reformatting it.¹

A more readable version would be:

reasoning_conversations = [
    tokenizer.apply_chat_template(conv["conversations"], tokenize=False)
    for conv in reasoning_dataset.map(generate_conversation)
]

Style Guide References

Footnotes

1. PEP 8 provides guidelines on code layout, including how to format long lines and comprehensions, to improve readability. ↩

[gemini-code-assist]

This list comprehension is written across multiple lines, but line 131 is very long and contains multiple parts of the expression, which harms readability. According to PEP 8, lines should be limited to 79 or 99 characters for readability.¹ It's better to format long list comprehensions over multiple lines. This will make the code easier to read and maintain.

Suggested change

	reasoning_conversations = [
	tokenizer.apply_chat_template(conv["conversations"],tokenize = False,) for conv in reasoning_dataset.map(generate_conversation)
	]
	reasoning_conversations = [
	tokenizer.apply_chat_template(conv["conversations"], tokenize=False)
	for conv in reasoning_dataset.map(generate_conversation)
	]

Style Guide References

Footnotes

1. PEP 8 suggests limiting all lines to a maximum of 79 characters (or 99 for some projects) to improve readability. ↩