Add zero-shot classification

[seanmor5]

I will add tests tomorrow when I get the multi-prompt case down, but right now it seems to be working fine:

Btw, is there a reason we hid documentation for TokenClassification ?

[jonatanklosko]

Awesome! I think ideally we should follow the same format as image/text classification, so %{predictions: [%{label: ..., score: ...}, ...]}. For now I wouldn't include prompt and label in the output, since the user has the input, and for multiple they can just zip.

Btw, is there a reason we hid documentation for TokenClassification?

We use defdelegate in the Bumblebee.Text module and there are all the docs :)

[jonatanklosko]

One thing to note is that here a single input becomes multiple inputs in the batch (each {prompt, hypothesis} pair). We could group them by adding an additional leading axis, however the number of hypothesis may be different for every input in the batch, so we can't really do that. Probably just documenting it is the way to go?

For Stable Diffusion if someone sets num_images_per_prompt: 2, the number of sub-inputs is fixed and we treat it as a single member of the batch (by adding a leading axis). I wonder if we should instead treat num_images_per_prompt: 2 as two inputs, so for batch size of 1 we would generate each image separately. This way if someone sets 4 images and batch size of 1 it would take more time, but don't blow up the memory, and they can set batch size to 4 to generate at once. This decoupling gives more control to the user.

@seanmor5 @josevalim thoughts?

[jonatanklosko]

FTR we decided to stick to the current approach, that is :batch_size always refers to the serving input. It may be inflated by options such as num_images_pre_prompt, or the number of labels in this case.

[jonatanklosko]

@seanmor5 I fixed the post processing to apply softmax per batch on the entailment label. I also adjusted the assertions based on this:

from transformers import pipeline

p = pipeline("zero-shot-classification", model="facebook/bart-large-mnli", candidate_labels=["cooking", "traveling", "dancing"])

p("one day I will see the world")

[jonatanklosko]