How to prepare a training dataset for Sentence Transformers

Sentence Transformer losses read dataset columns as model inputs rather than inferring their meaning from friendly column names. A raw export can therefore run with the wrong input order or leak metadata into the trainer unless its schema is reduced before training.

The current trainer accepts Hugging Face Dataset and DatasetDict objects. Columns named label, labels, score, or scores become targets; every other column becomes an input in its existing order, which makes select_columns() the direct way to remove ticket metadata and preserve anchor-positive ordering.

A six-row support-search sample becomes separate training and evaluation splits. The training split exposes anchor and positive for MultipleNegativesRankingLoss, while the held-out evaluation split exposes sentence1, sentence2, and score; a second Python process reloads the saved Arrow data and prints both schemas.

1. Install Hugging Face Datasets in the Python environment used for Sentence Transformers training.

   $ python -m pip install --upgrade datasets

   The full Sentence Transformers training extra also installs datasets when the project still needs the trainer dependencies.
   Related: How to install Sentence Transformers with pip

2. Create prepare_training_dataset.py with the imports and raw support records.

   prepare_training_dataset.py

   from datasets import Dataset, DatasetDict, disable_progress_bars, load_from_disk
    
    
   disable_progress_bars()
    
   raw_rows = [
       {
           "ticket_id": "T-1001",
           "query": "reset an expired API token",
           "positive": "create a replacement access token",
           "score": 0.94,
           "queue": "identity",
       },
       {
           "ticket_id": "T-1002",
           "query": "restore a deleted project",
           "positive": "recover a project from deleted items",
           "score": 0.92,
           "queue": "projects",
       },
       {
           "ticket_id": "T-1003",
           "query": "export audit logs",
           "positive": "download the audit log CSV file",
           "score": 0.89,
           "queue": "compliance",
       },
       {
           "ticket_id": "T-1004",
           "query": "enable two factor authentication",
           "positive": "enroll an authenticator app",
           "score": 0.91,
           "queue": "identity",
       },
       {
           "ticket_id": "T-1005",
           "query": "change the billing contact",
           "positive": "update the primary billing contact",
           "score": 0.86,
           "queue": "billing",
       },
       {
           "ticket_id": "T-1006",
           "query": "invite a new support agent",
           "positive": "add another person to the support team",
           "score": 0.84,
           "queue": "identity",
       },
   ]
    
   raw_dataset = Dataset.from_list(raw_rows)

   The sample fields represent a project source record before preparation. Identifiers and routing metadata remain available for the leakage check but are absent from the model inputs.

3. Append the deterministic training and evaluation split logic below raw_dataset.

   split = raw_dataset.train_test_split(test_size=2, seed=42)
   assert set(split["train"]["ticket_id"]).isdisjoint(split["test"]["ticket_id"])
    
   ranking_train = split["train"].select_columns(["query", "positive"])
   ranking_train = ranking_train.rename_column("query", "anchor")
    
   similarity_eval = split["test"].select_columns(["query", "positive", "score"])
   similarity_eval = similarity_eval.rename_columns(
       {"query": "sentence1", "positive": "sentence2"}
   )
    
   prepared = DatasetDict({"train": ranking_train, "eval": similarity_eval})

   The split occurs before ticket_id and queue are removed, so the assertion can catch row leakage. select_columns() then fixes both the number and order of inputs seen by each loss or evaluator.

4. Append the schema assertions and persistence logic below prepared.

   assert prepared["train"].column_names == ["anchor", "positive"]
   assert prepared["eval"].column_names == ["sentence1", "sentence2", "score"]
   assert prepared["eval"].features["score"].dtype == "float64"
    
   prepared.save_to_disk("prepared/support-search")
   reloaded = load_from_disk("prepared/support-search")
    
   assert reloaded["train"].column_names == ["anchor", "positive"]
   assert reloaded["eval"].column_names == ["sentence1", "sentence2", "score"]
   assert len(reloaded["train"]) == 4
   assert len(reloaded["eval"]) == 2
    
   print("train columns:", reloaded["train"].column_names)
   print("evaluation columns:", reloaded["eval"].column_names)
   print("train rows:", len(reloaded["train"]))
   print("evaluation rows:", len(reloaded["eval"]))
   print("evaluation score type:", reloaded["eval"].features["score"].dtype)
   print("saved splits:", list(reloaded.keys()))

   The assertions stop the script before saving an incompatible schema or unexpected split size. A held-out evaluation split avoids sharing source records with training.

5. Run the completed preparation script.

   $ python prepare_training_dataset.py
   train columns: ['anchor', 'positive']
   evaluation columns: ['sentence1', 'sentence2', 'score']
   train rows: 4
   evaluation rows: 2
   evaluation score type: float64
   saved splits: ['train', 'eval']

6. Reload prepared/support-search in a separate Python process to confirm both saved splits.

   $ python -c 'from datasets import load_from_disk; data = load_from_disk("prepared/support-search"); print(data); print(data["train"][0])'
   DatasetDict({
       train: Dataset({
           features: ['anchor', 'positive'],
           num_rows: 4
       })
       eval: Dataset({
           features: ['sentence1', 'sentence2', 'score'],
           num_rows: 2
       })
   })
   {'anchor': 'invite a new support agent', 'positive': 'add another person to the support team'}