search

guillaume-be / rust-bert

Rust native ready-to-use NLP pipelines and transformer-based models (BERT, DistilBERT, GPT2,...)
https://docs.rs/crate/rust-bert
Apache License 2.0
2.64k stars 215 forks source link

Avoid panics in prediction code #297

Closed anna-hope closed 1 year ago

anna-hope commented 1 year ago

I am building a small application that serves the Zeroshot Model and allows getting predictions over HTTP via Axum: https://github.com/anna-hope/zeroshot-rs

I have structured it so that the model runs in a separate thread, per this example. The problem I'm now running into is that if the model prediction methods get invalid inputs (e.g. an empty label array), they panic and thus bring down that whole thread, without a straightforward way to recover. If that happens, it essentially renders the whole application inoperable.

Some examples:

https://github.com/guillaume-be/rust-bert/blob/5d2b107e9961f3f75b00d039afedd6192b5ed636/src/pipelines/zero_shot_classification.rs#L631

https://github.com/guillaume-be/rust-bert/blob/5d2b107e9961f3f75b00d039afedd6192b5ed636/src/pipelines/zero_shot_classification.rs#L742

https://github.com/guillaume-be/rust-bert/blob/5d2b107e9961f3f75b00d039afedd6192b5ed636/src/pipelines/zero_shot_classification.rs#L743

among other places.

I imagine similar patterns can be found in other pipelines, although I haven't checked.

Ideally, we should avoid using unwrap and panicking, and return Result instead. I understand that would cause a breaking change in the prediction API, but it would enable more robust and error-resilient uses of this library in applications such as mine. I understand I can validate the inputs in my application code, but more idiomatic and explicit error handling would lead to a better design overall.

Edit: better wording, expand explanation.

guillaume-be commented 1 year ago

Hello @anna-hope ,

Thank you for raising this. You are absolutely right - unwraps that may cause panic on user input should be replaced by Result for better error handling. Please note that even removing all unwrap from the library does not remove all risks of panic during runtime (especially those caused by tensor error (e.g. invalid shape) - these would however be implementation bug that need to be addressed rather than invalid user input). Some unwrap are also guarantees not to fail, and these are probably fine.

Would you like to submit a pull request to update the zero-shot classification (and/or other pipelines in the library)? Please let me know if you need any help.

Thanks!