Add trace analysis tools for Docent integration

[ekellbuch]

Summary

• Add upload_data.py to upload HuggingFace evaluation traces to Docent
• Add download_rubric.py to download rubric definitions from Docent collections
• Add README.md with documentation for the trace analysis pipeline

Test plan

• Test upload_data.py with a sample HuggingFace dataset
• Test download_rubric.py with an existing collection
• Verify README instructions are accurate

🤖 Generated with Claude Code

[gemini-code-assist]

Summary of Changes

Hello @ekellbuch, 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 introduces a suite of new tools aimed at enhancing the integration between HuggingFace evaluation datasets and the Docent platform. The primary goal is to enable seamless uploading of agent interaction traces for detailed analysis and to provide functionality for downloading existing rubric definitions. These tools streamline the process of evaluating agent performance and managing evaluation criteria within the Docent ecosystem.

Highlights

• New Trace Upload Utility: Introduced upload_data.py to facilitate the upload of HuggingFace evaluation traces to Docent, supporting both 'trajectory' and 'conversations' dataset formats. This script includes robust parsing for various message types, tool calls, and metadata, along with filtering and dry-run capabilities.
• New Rubric Download Utility: Added download_rubric.py, a Python script designed to fetch and save rubric definitions from specified Docent collections to local JSON files.
• Comprehensive Documentation: A new README.md has been added within the eval/trace_analysis directory, providing detailed instructions, examples, CLI options, and supported dataset formats for the entire trace analysis pipeline.
• Specialized LAION Trace Uploader: A dedicated script, eval/traces/trace_analysis.py, was added for uploading LAION thinking traces from HuggingFace to Docent, offering predefined model-to-HuggingFace dataset mappings.

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 introduces a set of useful tools for integrating with Docent by uploading trace data and managing rubrics. The new scripts are generally well-written and the README.md provides good documentation for the workflow. My main concern is the presence of two similar data upload scripts, eval/trace_analysis/upload_data.py and eval/traces/trace_analysis.py, which creates redundancy. I've suggested consolidating them to improve maintainability. I've also provided specific feedback to improve error handling, fix a documentation typo, and add batching support for uploads to handle large datasets more effectively.

[ekellbuch]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request introduces new tools for trace analysis with Docent, including scripts for uploading traces and downloading rubrics, along with documentation. The code is well-structured, but I've identified a few areas for improvement in terms of robustness, code clarity, and adherence to conventions. My comments focus on improving exception handling, simplifying complex logic, and fixing file formatting. Specifically, I've suggested changes to make exception handling more comprehensive, refactor brittle string parsing logic, improve readability of data extraction, and correct end-of-file formatting.

[gemini-code-assist]

Code Review

This pull request introduces a valuable set of tools for integrating with Docent to perform trace analysis. The upload_data.py script is particularly comprehensive in its ability to handle various data formats, and the download_rubric.py script is a useful addition for managing analysis definitions. The documentation in README.md is clear and helpful. My review focuses on enhancing the robustness, error handling, and maintainability of the new scripts to ensure they are as reliable as possible.

[richardzhuang0412]

lgtm