lhoestq:

final_output["labels"] = [label for label, score in zip(output['labels'], output['scores']) if score > threshold]
final_output["scores"] = [score for score in output['scores'] if score > threshold]



Thanks for the reply. However, when using this, I’m getting the following error:
2 updated_dataset_batched = unlabeled_data.map(zero_shot_classify_sequences_batched_labels_and_scores, 
----> 3                                              batched=True, batch_size=2)

~/(...)/lib/python3.6/site-packages/datasets/arrow_dataset.py in map(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint)
   1307                 fn_kwargs=fn_kwargs,
   1308                 new_fingerprint=new_fingerprint,
-> 1309                 update_data=update_data,
   1310             )
   1311         else:

~/(...)/lib/python3.6/site-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    202         }
    203         # apply actual function
--> 204         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    205         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    206         # re-apply format to the output

~/(...)/lib/python3.6/site-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    335             # Call actual function
    336 
--> 337             out = func(self, *args, **kwargs)
    338 
    339             # Update fingerprint of in-place transforms + update in-place history of transforms

~/projecten/datascouts/env_datascouts/lib/python3.6/site-packages/datasets/arrow_dataset.py in _map_single(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset, update_data)
   1580                     if update_data:
   1581                         batch = cast_to_python_objects(batch)
-> 1582                         writer.write_batch(batch)
   1583             if update_data:
   1584                 writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file

~/(...)/lib/python3.6/site-packages/datasets/arrow_writer.py in write_batch(self, batch_examples, writer_batch_size)
    274             typed_sequence = TypedSequence(batch_examples[col], type=col_type, try_type=col_try_type)
    275             typed_sequence_examples[col] = typed_sequence
--> 276         pa_table = pa.Table.from_pydict(typed_sequence_examples)
    277         self.write_table(pa_table, writer_batch_size)
    278 

~/(...)/lib/python3.6/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.from_pydict()

~/(...)/lib/python3.6/site-packages/pyarrow/array.pxi in pyarrow.lib.asarray()

~/(...)/lib/python3.6/site-packages/pyarrow/array.pxi in pyarrow.lib.array()

~/(...)/lib/python3.6/site-packages/pyarrow/array.pxi in pyarrow.lib._handle_arrow_array_protocol()

~/(...)/lib/python3.6/site-packages/datasets/arrow_writer.py in __arrow_array__(self, type)
     95                 out = pa.ExtensionArray.from_storage(type, pa.array(self.data, type.storage_dtype))
     96             else:
---> 97                 out = pa.array(self.data, type=type)
     98             if trying_type and out[0].as_py() != self.data[0]:
     99                 raise TypeError(

~/(...)/lib/python3.6/site-packages/pyarrow/array.pxi in pyarrow.lib.array()

~/(...)/lib/python3.6/site-packages/pyarrow/array.pxi in pyarrow.lib._sequence_to_array()

~/(...)/lib/python3.6/site-packages/pyarrow/error.pxi in pyarrow.lib.pyarrow_internal_check_status()

~/(...)/lib/python3.6/site-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowInvalid: Invalid null value

My function looks as follows:
# let's use a higher threshold
def zero_shot_classify_sequences_batched(examples, threshold=0.6):
    # first, send text through pipeline
    texts = examples['text']
    outputs = classifier(texts, candidate_labels, multi_class=True)
    # next, for each output:
    final_outputs = []
    for output in outputs:
        final_output = {}
        final_output["labels"] = [label for label, score in zip(output['labels'], output['scores']) if score > threshold]
        final_output["scores"] = [score for score in output['scores'] if score > threshold]
        final_outputs.append(final_output)
    
    assert len(final_outputs) == len(texts)
    # set final outputs
    examples['predicted_labels'] = final_outputs
    
    return examples