Refactor diarization#341
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @MahmoudAshraf97, I'm Gemini Code Assist1! 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 significantly refactors the speaker diarization module by introducing a new diarization Python package. The primary goal is to modularize the diarization process, making it easier to integrate and support additional diarization models in the future. The changes involve encapsulating the MSDD diarization logic into a dedicated class, streamlining the main processing scripts to use this new modular approach, and removing several outdated configuration files and helper scripts. This refactoring enhances code organization, maintainability, and extensibility for speaker diarization functionalities.
Highlights
- Modularization of Diarization Logic: The core diarization logic, previously scattered and reliant on external scripts, has been encapsulated within a new
diarizationPython package. This improves organization and maintainability. - Centralized MSDD Diarization: A dedicated
MSDDDiarizerclass has been introduced to centralize the setup and execution of the Multiscale Diarization Decoder (MSDD) model, abstracting away complex configuration and temporary file management. - Streamlined Main Scripts: The main
diarize.pyanddiarize_parallel.pyscripts have been streamlined. They now interact with the newMSDDDiarizerclass, significantly reducing their complexity and removing direct dependencies on NeMo internals and external helper scripts. - Improved Parallel Processing: The
diarize_parallel.pyscript now leverages Python'smultiprocessingmodule to run diarization in parallel, replacing the previoussubprocesscall tonemo_process.py. This offers better control and integration. - Cleanup of Obsolete Files: Redundant configuration files (
diar_infer_general.yaml,diar_infer_meeting.yaml) and thenemo_process.pyscript have been removed, simplifying the project structure and reducing overhead.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request is a solid refactoring that modularizes the diarization logic into a separate package, which will make it easier to add new models in the future. The switch from subprocess to multiprocessing for parallel execution is also a significant improvement. My review highlights a critical import error that needs to be fixed, a high-priority concern about dependency on private library internals, and a few medium-priority suggestions to enhance code structure and maintainability. Overall, these are great changes that improve the codebase.
| parser.add_argument( | ||
| "--diarizer", | ||
| default="msdd", | ||
| choices=["msdd"], | ||
| help="Choose the diarization model to use", | ||
| ) |
There was a problem hiding this comment.
The --diarizer argument is added but args.diarizer is never used to select the diarization model. The script unconditionally uses MSDDDiarizer.
This is inconsistent with diarize.py, which uses a conditional block to select and import the diarizer. To improve consistency and make it easier to add new diarization models in the future, consider adopting a similar pattern here. For example, you could select the diarizer class in the main block and pass it to the diarize_parallel function.
There was a problem hiding this comment.
Pull Request Overview
This PR refactors the diarization functionality by extracting speaker diarization logic from the main processing files into a dedicated diarization module. The refactor removes the standalone nemo_process.py script and configuration files in favor of a structured diarization package with the MSDD (Multi-Scale Diarization Decoder) implementation.
Key changes include:
- Creation of a new
diarizationmodule with MSDD-specific implementation - Integration of diarization as a library component rather than subprocess execution
- Consolidation of configuration management within the diarization module
Reviewed Changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| nemo_process.py | Removed standalone diarization script |
| nemo_msdd_configs/*.yaml | Removed external configuration files |
| helpers.py | Removed create_config function and related imports |
| diarize_parallel.py | Updated to use new diarization module with multiprocessing |
| diarize.py | Updated to use new diarization module directly |
| diarization/msdd/msdd.py | New MSDD diarization implementation |
| diarization/msdd/diar_infer_telephonic.yaml | New configuration file for MSDD |
| diarization/init.py | Module initialization with exports |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
You can also share your feedback on Copilot code review for a chance to win a $100 gift card. Take the survey.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1afe413664
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| args.device, | ||
| dtype=torch.float16 if args.device == "cuda" else torch.float32, | ||
| ) | ||
| nemo_process.join() |
There was a problem hiding this comment.
Consume queue payload before joining diarization process
This waits on nemo_process.join() before reading from results_queue, which can deadlock for larger diarization results (for example, long recordings with many RTTM segments): the child may block while flushing queue.put(...) to the pipe, while the parent is blocked waiting for child exit. Reading from the queue first (or using a different synchronization pattern) avoids hangs in real runs.
Useful? React with 👍 / 👎.
As mentioned in #239, This PR refactors the diarization module into a separate file to allow adding more models