| import json |
| import random |
| from typing import List, Dict, Tuple, Iterator |
| import torch |
| from torch.utils.data import Dataset, Sampler |
| from transformers import PreTrainedTokenizerBase |
|
|
|
|
| class ConceptSynonymDataset(Dataset): |
| """ |
| Dataset for SapBERT metric learning. |
| Each item is a (synonym_text, label_idx) pair. |
| The concept_id is mapped to a contiguous integer label [0, N-1]. |
| """ |
|
|
| def __init__(self, jsonl_path: str): |
| self.samples: List[Tuple[str, int]] = [] |
| self.concept_to_label: Dict[str, int] = {} |
| self.label_to_concept: Dict[int, str] = {} |
| self.label_to_indices: Dict[int, List[int]] = {} |
| self.concept_names: Dict[str, str] = {} |
| self.vocabularies: Dict[str, str] = {} |
| self.concept_codes: Dict[str, str] = {} |
|
|
| self._load_data(jsonl_path) |
|
|
| def _load_data(self, jsonl_path: str): |
| label_counter = 0 |
| with open(jsonl_path, "r", encoding="utf-8") as f: |
| for line in f: |
| if not line.strip(): |
| continue |
| record = json.loads(line) |
| cid = str(record["concept_id"]) |
| synonyms = record.get("synonyms", []) |
| if len(synonyms) < 2: |
| continue |
|
|
| if cid not in self.concept_to_label: |
| label = label_counter |
| self.concept_to_label[cid] = label |
| self.label_to_concept[label] = cid |
| self.concept_names[cid] = record.get("concept_name", "") |
| self.vocabularies[cid] = record.get("vocabulary_id", "") |
| self.concept_codes[cid] = record.get("concept_code", "") |
| label_counter += 1 |
| else: |
| label = self.concept_to_label[cid] |
|
|
| if label not in self.label_to_indices: |
| self.label_to_indices[label] = [] |
|
|
| for syn in synonyms: |
| if syn and len(syn.strip()) > 0: |
| idx = len(self.samples) |
| self.samples.append((syn.strip(), label)) |
| self.label_to_indices[label].append(idx) |
|
|
| print( |
| f"Loaded {len(self.samples):,} terms across {len(self.concept_to_label):,} unique concepts from {jsonl_path}" |
| ) |
|
|
| def __len__(self) -> int: |
| return len(self.samples) |
|
|
| def __getitem__(self, idx: int) -> Tuple[str, int]: |
| return self.samples[idx] |
|
|
| @property |
| def num_classes(self) -> int: |
| return len(self.concept_to_label) |
|
|
|
|
| class MPerClassSampler(Sampler[int]): |
| """ |
| Samples M items per class in each batch. |
| Essential for metric learning (e.g., Multi-Similarity Loss) so that every |
| concept in a batch has at least M positive pairs. |
| """ |
|
|
| def __init__( |
| self, |
| dataset: ConceptSynonymDataset, |
| m: int = 2, |
| batch_size: int = 64, |
| length_before_new_iter: int = 100000, |
| ): |
| super().__init__() |
| self.dataset = dataset |
| self.m = m |
| self.batch_size = batch_size |
| self.classes_per_batch = max(1, batch_size // m) |
| self.labels = list(self.dataset.label_to_indices.keys()) |
| self.length = min(length_before_new_iter, len(self.dataset)) |
| self.total_batches = self.length // self.batch_size |
|
|
| def __len__(self) -> int: |
| return self.total_batches * self.batch_size |
|
|
| def __iter__(self) -> Iterator[int]: |
| batches_yielded = 0 |
| while batches_yielded < self.total_batches: |
| selected_classes = random.sample(self.labels, self.classes_per_batch) |
| batch = [] |
| for cls in selected_classes: |
| indices = self.dataset.label_to_indices[cls] |
| if len(indices) >= self.m: |
| batch.extend(random.sample(indices, self.m)) |
| else: |
| batch.extend(random.choices(indices, k=self.m)) |
|
|
| random.shuffle(batch) |
| for idx in batch: |
| yield idx |
| batches_yielded += 1 |
|
|
|
|
| def collate_fn(batch: List[Tuple[str, int]], tokenizer: PreTrainedTokenizerBase, max_length: int = 40): |
| """ |
| Collate function to dynamically tokenize batch texts and convert labels to tensors. |
| """ |
| texts = [item[0] for item in batch] |
| labels = torch.tensor([item[1] for item in batch], dtype=torch.long) |
|
|
| encoded = tokenizer( |
| texts, |
| padding=True, |
| truncation=True, |
| max_length=max_length, |
| return_tensors="pt", |
| ) |
|
|
| return { |
| "input_ids": encoded["input_ids"], |
| "attention_mask": encoded["attention_mask"], |
| "labels": labels, |
| } |
|
|
