"Shape of query is incorrect, it has to be either a 1D array or 2D (1, N)"
#3323 
## Describe the bug
Passing a correctly shaped Numpy-Array to get_nearest_examples leads to the Exception

"Shape of query is incorrect, it has to be either a 1D array or 2D (1, N)"

Probably the reason for this is a wrongly converted assertion.

1.15.1:

`assert len(query.shape) == 1 or (len(query.shape) == 2 and query.shape[0] == 1)`

1.16.1:

```
 if len(query.shape) != 1 or (len(query.shape) == 2 and query.shape[0] != 1):
            raise ValueError("Shape of query is incorrect, it has to be either a 1D array or 2D (1, N)")
```

## Steps to reproduce the bug

follow the steps described here: https://huggingface.co/course/chapter5/6?fw=tf

```python
 question_embedding.shape # (1, 768)

 scores, samples = embeddings_dataset.get_nearest_examples(
    "embeddings", question_embedding, k=5 # Error
)

# "Shape of query is incorrect, it has to be either a 1D array or 2D (1, N)"
```

## Expected results
Should work without exception

## Actual results
Throws exception

## Environment info
- `datasets` version: 1.15.1
- Platform: Darwin-20.6.0-x86_64-i386-64bit
- Python version: 3.7.12
- PyArrow version: 6.0.
