returning model raw outputs as well - #8
Conversation
Reviewer's guide (collapsed on small PRs)Reviewer's GuideThis PR updates the entity prediction APIs to also expose the raw model outputs while preserving the previous return shape for single-text predictions, and refactors the internal call to keep the model output object intact before indexing into its first element. Sequence diagram for updated entity prediction and raw outputssequenceDiagram
actor Client
participant GlinerModel
participant UnderlyingModel
Client->>GlinerModel: predict_entities(text, labels, flat_ner, threshold, multi_label)
GlinerModel->>GlinerModel: batch_predict_entities([text], labels, flat_ner, threshold, multi_label)
GlinerModel->>UnderlyingModel: model(model_input)
UnderlyingModel-->>GlinerModel: model_output_raw
GlinerModel->>GlinerModel: model_output = model_output_raw[0]
GlinerModel->>GlinerModel: postprocess to all_entities
GlinerModel-->>GlinerModel: return all_entities, model_output_raw
GlinerModel->>GlinerModel: all_entities[0]
GlinerModel-->>Client: entities_for_single_text
Class diagram for updated prediction methods exposing raw outputsclassDiagram
class GlinerModel {
+predict_entities(text, labels, flat_ner, threshold, multi_label) entities
+batch_predict_entities(texts, labels, flat_ner, threshold, multi_label) (all_entities, model_output_raw)
+prepare_model_inputs(texts, labels) (model_input, raw_batch)
}
class UnderlyingModel {
+__call__(model_input) model_output_raw
}
GlinerModel --> UnderlyingModel : uses for inference
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Summary of ChangesHello @arthrod, 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 enhances the GLiNER model's prediction capabilities by exposing the raw outputs from the underlying model. This allows users to access more detailed information beyond just the extracted entities, which can be valuable for debugging, advanced analysis, or custom post-processing. The primary Highlights
Changelog
Using Gemini Code AssistThe 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
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 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
|
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts (beta)
Comment |
There was a problem hiding this comment.
Hey - I've left some high level feedback:
- Changing
batch_predict_entitiesto return(all_entities, model_output_raw)instead of justall_entitiesis a breaking API change; consider either adding a new method for the extended return value or keeping the old return shape and exposing the raw outputs via an optional flag or attribute to avoid breaking existing callers. - In
predict_entities,model_outputis assigned but never used and the method still only returnsall_entities[0]; either return the raw outputs there as well or remove the unused variable to keep the intent clear. - Update the docstrings (or function comments) for
predict_entitiesandbatch_predict_entitiesto describe the new return values so callers understand that the batch method now returns both entities and raw model outputs.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- Changing `batch_predict_entities` to return `(all_entities, model_output_raw)` instead of just `all_entities` is a breaking API change; consider either adding a new method for the extended return value or keeping the old return shape and exposing the raw outputs via an optional flag or attribute to avoid breaking existing callers.
- In `predict_entities`, `model_output` is assigned but never used and the method still only returns `all_entities[0]`; either return the raw outputs there as well or remove the unused variable to keep the intent clear.
- Update the docstrings (or function comments) for `predict_entities` and `batch_predict_entities` to describe the new return values so callers understand that the batch method now returns both entities and raw model outputs.Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b6df740f6f
ℹ️ 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".
| all_entities.append(entities) | ||
|
|
||
| return all_entities | ||
| return all_entities, model_output_raw |
There was a problem hiding this comment.
Preserve batch_predict_entities return contract
Changing batch_predict_entities to always return (all_entities, model_output_raw) breaks existing callers that were written against the previous contract (entities = model.batch_predict_entities(...) returning only entity lists). In downstream inference pipelines this can immediately raise type errors or silently corrupt post-processing when code iterates/serializes the returned value, so this should be made backward-compatible (for example via an opt-in flag for raw outputs) rather than a hard return-type change.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
gliner/model.py (1)
138-174:⚠️ Potential issue | 🟠 MajorBreaking API change: update docstring and consider backward compatibility.
Changing the return type from
ListtoTuple[List, Any]is a breaking change. Existing callers using:entities = model.batch_predict_entities(texts, labels)will now receive a tuple and likely break.
Two issues to address:
- Update the docstring to document the new return type:
""" Predict entities for a batch of texts. texts: List of texts | List[str] labels: List of labels | List[str] ... + + Returns: + Tuple[List[List[Dict]], Any]: A tuple containing (all_entities, model_output_raw) """
- Consider backward compatibility via an optional parameter:
def batch_predict_entities(self, texts, labels, flat_ner=True, threshold=0.5, multi_label=False, return_model_output=False): ... if return_model_output: return all_entities, model_output_raw return all_entities🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@gliner/model.py` around lines 138 - 174, The batch_predict_entities function now returns a tuple (all_entities, model_output_raw) which is a breaking API change; update the function docstring to document the new return value and add a backward-compatible optional flag (e.g., return_model_output=False) to preserve the original behavior. Modify the signature of batch_predict_entities to accept return_model_output=False, and at the end return (all_entities, model_output_raw) only when return_model_output is True, otherwise return just all_entities; reference the existing symbols model_output_raw, all_entities, and the function name batch_predict_entities when making these changes.
🧹 Nitpick comments (1)
gliner/model.py (1)
132-135: Unused variable and API inconsistency.The
model_outputvariable is unpacked but never used. Per Python convention, prefix it with an underscore to indicate it's intentionally discarded:- all_entities, model_output = self.batch_predict_entities( + all_entities, _model_output = self.batch_predict_entities( [text], labels, flat_ner=flat_ner, threshold=threshold, multi_label=multi_label )Additionally, should
predict_entitiesalso return raw outputs for API consistency withbatch_predict_entities? If backward compatibility is the goal, consider adding an optionalreturn_raw_output=Falseparameter instead.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@gliner/model.py` around lines 132 - 135, The local variable model_output in predict_entities is unused — rename it to _model_output when unpacking the result of batch_predict_entities to signal it's intentionally discarded (i.e., replace model_output with _model_output in predict_entities). If you want API consistency with batch_predict_entities, add an optional parameter return_raw_output: bool = False to predict_entities' signature and, when True, return a tuple (all_entities[0], _model_output) otherwise keep the current behavior of returning all_entities[0]; ensure the call to batch_predict_entities preserves the existing parameters (labels, flat_ner, threshold, multi_label) and maintain backward compatibility by defaulting return_raw_output to False.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@gliner/model.py`:
- Around line 138-174: The batch_predict_entities function now returns a tuple
(all_entities, model_output_raw) which is a breaking API change; update the
function docstring to document the new return value and add a
backward-compatible optional flag (e.g., return_model_output=False) to preserve
the original behavior. Modify the signature of batch_predict_entities to accept
return_model_output=False, and at the end return (all_entities,
model_output_raw) only when return_model_output is True, otherwise return just
all_entities; reference the existing symbols model_output_raw, all_entities, and
the function name batch_predict_entities when making these changes.
---
Nitpick comments:
In `@gliner/model.py`:
- Around line 132-135: The local variable model_output in predict_entities is
unused — rename it to _model_output when unpacking the result of
batch_predict_entities to signal it's intentionally discarded (i.e., replace
model_output with _model_output in predict_entities). If you want API
consistency with batch_predict_entities, add an optional parameter
return_raw_output: bool = False to predict_entities' signature and, when True,
return a tuple (all_entities[0], _model_output) otherwise keep the current
behavior of returning all_entities[0]; ensure the call to batch_predict_entities
preserves the existing parameters (labels, flat_ner, threshold, multi_label) and
maintain backward compatibility by defaulting return_raw_output to False.
There was a problem hiding this comment.
Code Review
This pull request modifies the prediction helpers to return raw model outputs, updating batch_predict_entities and predict_entities accordingly. It also resolves a critical issue where the original model response was an empty, invalid JSON, by ensuring a minimal valid JSON structure is always returned. The changes are clear and achieve the stated goal, with a minor suggestion for improving code clarity by handling an unused variable.
| all_entities, model_output = self.batch_predict_entities( | ||
| [text], labels, flat_ner=flat_ner, threshold=threshold, multi_label=multi_label | ||
| )[0] | ||
| ) |
There was a problem hiding this comment.
The model_output variable is assigned but never used. To signal that it's intentionally being ignored, you can use an underscore _ as the variable name. This improves code clarity and adheres to common Python conventions for unused variables.
| all_entities, model_output = self.batch_predict_entities( | |
| [text], labels, flat_ner=flat_ner, threshold=threshold, multi_label=multi_label | |
| )[0] | |
| ) | |
| all_entities, _ = self.batch_predict_entities( | |
| [text], labels, flat_ner=flat_ner, threshold=threshold, multi_label=multi_label | |
| ) |
Summary by Sourcery
Return both extracted entities and raw model outputs from prediction helpers.
New Features:
Enhancements: