Preserve dataset.features when using sample_dataset (#396)

* Update data.py

If a column has been defined as a ClassLabel then `sample_dataset` strips that information away and you lose names.

Test code

```python
 from datasets import load_dataset
import datasets
from sentence_transformers.losses import CosineSimilarityLoss

from setfit import SetFitModel, SetFitTrainer, sample_dataset


# Load a dataset from the Hugging Face Hub
dataset: datasets.DatasetDict = load_dataset("SetFit/sst5")
dataset = dataset.class_encode_column("label_text")

# Simulate the few-shot regime by sampling 8 examples per class
train_dataset: datasets.Dataset = sample_dataset(dataset["train"], label_column="label_text", num_samples=8)
eval_dataset: datasets.Dataset = dataset["validation"]

# Load a SetFit model from Hub
model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-mpnet-base-v2")

# Create trainer
trainer = SetFitTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss_class=CosineSimilarityLoss,
    metric="accuracy",
    batch_size=16,
    num_iterations=20, # The number of text pairs to generate for contrastive learning
    num_epochs=1, # The number of epochs to use for contrastive learning
    column_mapping={"text": "text", "label_text": "label"} # Map dataset columns to text/label expected by trainer
)

# Train and evaluate
trainer.train()
metrics = trainer.evaluate()

# Run inference
preds = model(["i loved the spiderman movie!", "pineapple on pizza is the worst 🤮"])

print(preds)
print(list(map(lambda x: train_dataset.features["label_text"].names[x], preds)))
```

* Preserve features when calling sample_dataset + tests

---------

Co-authored-by: Tom Aarsen <[email protected]>

Author: grofte

src/setfit/data.py

@@ -173,7 +173,7 @@ def sample_dataset(dataset: Dataset, label_column: str = "label", num_samples: i
     df = df.apply(lambda x: x.sample(min(num_samples, len(x))))
     df = df.reset_index(drop=True)
 
-    all_samples = Dataset.from_pandas(df)
+    all_samples = Dataset.from_pandas(df, features=dataset.features)
     return all_samples.shuffle(seed=seed)
 
 

tests/test_data.py

@@ -225,3 +225,11 @@ def test_correct_model_inputs(tokenizer_name):
     # Verify that the x_batch contains exactly those keys that the model requires
     x_batch, _ = next(iter(dataloader))
     assert set(x_batch.keys()) == set(tokenizer.model_input_names)
+
+
+def test_preserve_features() -> None:
+    dataset = load_dataset("SetFit/sst5", split="train[:100]")
+    label_column = "label_text"
+    dataset = dataset.class_encode_column(label_column)
+    train_dataset = sample_dataset(dataset, label_column=label_column, num_samples=8)
+    assert train_dataset.features[label_column] == dataset.features[label_column]