TAO Mine Nearest Neighbors

Utilisez cette skill pour exécuter l'extraction par nearest-neighbor de TAO Data Services TMM. La skill consomme des parquets d'embedding et écrit un parquet source-sample extrait plus un résumé d'extraction. Elle ne calcule pas les embeddings ; les étapes en amont doivent d'abord produire les parquets d'embedding source et cible.

Le point d'entrée du conteneur est :

tmm nearest_neighbors -e /absolute/path/to/nearest_neighbors.yaml

TAO Data Services nécessite -e/--experiment_spec_file. Le script console tmm convertit ce YAML en arguments Hydra --config-path et --config-name en interne.

Inputs

L'utilisateur peut fournir soit un spec YAML nearest-neighbors existant, soit les champs nécessaires pour en générer un.

Champs spec requis :

Champ	Signification
`source_parquet`	Chemin absolu du parquet des embeddings candidat/source.
`target_parquet`	Chemin absolu du parquet des embeddings cible/requête.
`output_parquet`	Chemin absolu où TAO Data Services doit écrire les chemins de fichiers source extraits.

Champs optionnels courants :

Champ	Défaut	Signification
`topn`	`5`	Nombre d'échantillons source les plus proches à récupérer par échantillon cible.
`knn_metric`	`cosine`	L'un de `cosine`, `euclidean`, ou `manhattan`.
`source_embed_column_name`	`embedding`	Colonne d'embedding dans `source_parquet`.
`target_embed_column_name`	`embedding`	Colonne d'embedding dans `target_parquet`.
`filter_by_label`	`"false"`	Drapeau string. Quand `"true"`, TAO DS filtre les voisins en mettant en correspondance les colonnes `label` quand les deux parquets fournissent des labels.
`distance_threshold`	`-1.0`	Distance maximale à conserver. Un nombre négatif désactive le seuil.

Les deux parquets d'entrée doivent contenir une colonne filepath et une colonne d'embedding de type liste. Si filter_by_label est "true", les deux parquets doivent aussi contenir label.

Le template par défaut est assets/default_nearest_neighbors.yaml.

Quick Start

Exécutez à partir de la racine du repo tao-skill-bank. Résolvez l'image TAO Data Services épinglée depuis versions.yaml, vérifiez le spec, montez la racine d'exécution avec des chemins identiques hôte/conteneur, et transmettez les logs Docker.

SPEC=/absolute/path/to/nearest_neighbors.yaml
RUN_ROOT=/absolute/path/that/contains/specs/data/and/results
GPU_COUNT=1

python3 skills/data/tao-mine-nearest-neighbors/scripts/verify_nearest_neighbors_spec.py \
  --spec "$SPEC"

DS_IMAGE="$(scripts/resolve_versions_key.py images.tao_toolkit.data_services)"

docker run --rm --gpus "$GPU_COUNT" --shm-size=8g --network=host \
  -v "$RUN_ROOT:$RUN_ROOT" \
  -w "$RUN_ROOT" \
  "$DS_IMAGE" \
  tmm nearest_neighbors -e "$SPEC"

Utilisez au moins un GPU. Choisissez GPU_COUNT parmi le matériel disponible sur l'hôte ou la plateforme qui exécutera le conteneur. Si l'utilisateur ne connaît pas la bonne valeur, inspectez l'hôte avec nvidia-smi -L ou demandez quelle allocation GPU la run doit utiliser.

Ne passez pas --user $(id -u):$(id -g) au conteneur TAO data-services à moins d'avoir vérifié que l'image supporte cet UID. Certaines images TAO DS importent des paquets Python qui appellent getpass.getuser() au démarrage et échouent quand l'UID n'est pas présent dans /etc/passwd.

Generate A Spec

Si l'utilisateur fournit des chemins parquet source/cible/sortie au lieu d'un spec prêt, générez un spec à partir du template par défaut :

python3 skills/data/tao-mine-nearest-neighbors/scripts/prepare_nearest_neighbors_spec.py \
  --source-parquet /absolute/path/source_embeddings.parquet \
  --target-parquet /absolute/path/target_embeddings.parquet \
  --output-parquet /absolute/path/results/mined.parquet \
  --output-spec /absolute/path/specs/nearest_neighbors.yaml \
  --topn 5 \
  --knn-metric cosine \
  --filter-by-label false \
  --distance-threshold -1.0

Le YAML généré utilise des chemins absolus. Gardez le spec, les parquets d'entrée et le répertoire de sortie sous RUN_ROOT afin que les mêmes chemins se résolvent à l'intérieur du conteneur.

Preflight

Avant de lancer Docker :

Vérifiez Docker et l'accès au GPU :

docker info > /dev/null
nvidia-smi -L

Résolvez et téléchargez l'image data-services si nécessaire :

DS_IMAGE="$(scripts/resolve_versions_key.py images.tao_toolkit.data_services)"
docker image inspect "$DS_IMAGE" > /dev/null || docker pull "$DS_IMAGE"

Validez le spec :

python3 skills/data/tao-mine-nearest-neighbors/scripts/verify_nearest_neighbors_spec.py \
  --spec "$SPEC"

Confirmez que RUN_ROOT contient le spec, les deux parquets d'entrée et le répertoire de sortie. Montez RUN_ROOT au même chemin absolu à l'intérieur de Docker.

Outputs

La skill garantit les artefacts nommés par le spec :

Artefact	Localisation
parquet extrait	`output_parquet`
résumé d'extraction	`mining_summary.txt` à côté de `output_parquet`

La tâche nearest_neighbors actuelle de TAO Data Services écrit un parquet extrait avec des lignes filepath source uniques. Le fichier résumé rapporte les comptages d'extraction comme les requêtes traitées, les voisins considérés, les doublons supprimés et tout filtrage par label/distance.

Troubleshooting

The subtask nearest_neighbors requires -e/--experiment_spec_file : relancez avec tmm nearest_neighbors -e "$SPEC". Les remplacements Hydra seuls ne suffisent pas.

Parquet d'entrée introuvable à l'intérieur de Docker : le chemin YAML doit être visible à l'intérieur du conteneur. Utilisez un montage RUN_ROOT où les chemins hôte et conteneur sont identiques.

Le répertoire de sortie n'est pas accessible en écriture après la sortie de Docker : le conteneur TAO DS peut avoir écrit des fichiers en tant que root. Informez l'utilisateur, signalez les artefacts produits et demandez s'il faut réparer les permissions sur le répertoire de sortie avant de continuer.

Pas de GPU ou erreurs cuDF/cuML : l'extraction par nearest-neighbor nécessite au moins un GPU CUDA. Vérifiez nvidia-smi -L, le drapeau Docker --gpus et l'installation de NVIDIA container toolkit.