remove unused field for chat_template.default for DPO training

[timofey]

Description

DPO trainer expects only prompt+chosen+rejected columns present in the final dataset. However, the chat_template.default implementation passes an extra column messages which causes validations to fail.

Motivation and Context

Fixes issue #2415

How has this been tested?

Tested using sample config:

rl: dpo

datasets:
  - path: "..../dataset.jsonl"
    chat_template: tokenizer_default
    type: chat_template.default
    field_messages: "messages"
    field_chosen: "chosen"
    field_rejected: "rejected"
    message_property_mappings:
      role: role
      content: content
    roles:
      user: [ "user" ]
      assistant: [ "assistant" ]
      system: [ "system" ]

then by running axolotl preprocess and axolotl train commands.

Without the fix, it fails with the following errors:

KeyError: "Invalid keys in the example: {'messages', 'prompt', 'rejected', 'chosen'}"

Screenshots (if appropriate)

Types of changes

Social Handles (Optional)

Summary by CodeRabbit

• New Features
  • The system now provides additional metadata indicating which columns should be removed during processing.
• Bug Fixes
  • Updated dataset tests to reflect correct feature presence and absence, improving data validation.

[coderabbitai]

Walkthrough

The default function in src/axolotl/prompt_strategies/dpo/chat_template.py was updated to return a tuple containing both the transformation function and a metadata dictionary specifying columns to remove, instead of returning only the transformation function. Corresponding test updates were made to unpack this tuple. Additionally, a dataset test was modified to check for the absence of the "conversation" feature and the presence of "chosen", "rejected", and "prompt" features.

Changes

File(s)	Change Summary
src/axolotl/prompt_strategies/dpo/chat_template.py	Modified default to return a tuple: (transform_fn, {"remove_columns": [field_messages]}).
tests/prompt_strategies/test_dpo_chat_templates.py	Updated tests to unpack default return into (transform_fn, _) instead of single variable.
tests/test_datasets.py	Changed dataset feature assertions: removed check for "conversation", added checks for "chosen", "rejected", and "prompt".

Sequence Diagram(s)

sequenceDiagram
    participant Caller
    participant ChatTemplate as chat_template.py
    Caller->>ChatTemplate: Call default()
    ChatTemplate-->>Caller: Return (transform_fn, {"remove_columns": [field_messages]})

Loading

Poem

In the code where columns roam,
A tuple now is heading home.
"Remove these fields!" it softly cries,
As transformation underlies.
No more errors, no more fright,
The data now is trimmed just right.
🐇✨

---

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2e7fac1 and cf14ab5.

📒 Files selected for processing (1)

• src/axolotl/prompt_strategies/dpo/chat_template.py (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• src/axolotl/prompt_strategies/dpo/chat_template.py

⏰ Context from checks skipped due to timeout of 90000ms (8)

• GitHub Check: PyTest (3.11, 2.7.0)
• GitHub Check: PyTest from Source Dist (3.11, 2.6.0)
• GitHub Check: PyTest (3.11, 2.6.0)
• GitHub Check: PyTest from Source Dist (3.11, 2.7.0)
• GitHub Check: PyTest from Source Dist (3.11, 2.5.1)
• GitHub Check: PyTest (3.11, 2.5.1)
• GitHub Check: pre-commit
• GitHub Check: pre-commit

✨ Finishing Touches

• 📝 Generate Docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[winglian]

The caller of the transform_fn is going to need to unpack the updated return (tuple)

[winglian]

I think I fixed the linting issues and tests that use the updated return values

[timofey]

oh thanks @winglian ! I'll take a stab at it tonight as well, looks like some of the tests are still failing

[timofey]

@winglian can you run the tests on the pipeline again please? Adjusted the failed test, since the "conversation" field is removed now

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

📢 Thoughts on this report? Let us know!