| from torch.utils.data import Dataset |
| import logging |
| import gzip |
| from .. import SentenceTransformer |
| from ..readers import InputExample |
| from typing import List |
| import random |
|
|
|
|
| logger = logging.getLogger(__name__) |
|
|
| class ParallelSentencesDataset(Dataset): |
| """ |
| This dataset reader can be used to read-in parallel sentences, i.e., it reads in a file with tab-seperated sentences with the same |
| sentence in different languages. For example, the file can look like this (EN\tDE\tES): |
| hello world hallo welt hola mundo |
| second sentence zweiter satz segunda oración |
| |
| The sentence in the first column will be mapped to a sentence embedding using the given the embedder. For example, |
| embedder is a mono-lingual sentence embedding method for English. The sentences in the other languages will also be |
| mapped to this English sentence embedding. |
| |
| When getting a sample from the dataset, we get one sentence with the according sentence embedding for this sentence. |
| |
| teacher_model can be any class that implement an encode function. The encode function gets a list of sentences and |
| returns a list of sentence embeddings |
| """ |
|
|
| def __init__(self, student_model: SentenceTransformer, teacher_model: SentenceTransformer, batch_size: int = 8, use_embedding_cache: bool = True): |
| """ |
| Parallel sentences dataset reader to train student model given a teacher model |
| :param student_model: Student sentence embedding model that should be trained |
| :param teacher_model: Teacher model, that provides the sentence embeddings for the first column in the dataset file |
| """ |
| self.student_model = student_model |
| self.teacher_model = teacher_model |
| self.datasets = [] |
| self.datasets_iterator = [] |
| self.datasets_tokenized = [] |
| self.dataset_indices = [] |
| self.copy_dataset_indices = [] |
| self.cache = [] |
| self.batch_size = batch_size |
| self.use_embedding_cache = use_embedding_cache |
| self.embedding_cache = {} |
| self.num_sentences = 0 |
|
|
| def load_data(self, filepath: str, weight: int = 100, max_sentences: int = None, max_sentence_length: int = 128): |
| """ |
| Reads in a tab-seperated .txt/.csv/.tsv or .gz file. The different columns contain the different translations of the sentence in the first column |
| |
| :param filepath: Filepath to the file |
| :param weight: If more than one dataset is loaded with load_data: With which frequency should data be sampled from this dataset? |
| :param max_sentences: Max number of lines to be read from filepath |
| :param max_sentence_length: Skip the example if one of the sentences is has more characters than max_sentence_length |
| :param batch_size: Size for encoding parallel sentences |
| :return: |
| """ |
|
|
| logger.info("Load "+filepath) |
| parallel_sentences = [] |
|
|
| with gzip.open(filepath, 'rt', encoding='utf8') if filepath.endswith('.gz') else open(filepath, encoding='utf8') as fIn: |
| count = 0 |
| for line in fIn: |
| sentences = line.strip().split("\t") |
| if max_sentence_length is not None and max_sentence_length > 0 and max([len(sent) for sent in sentences]) > max_sentence_length: |
| continue |
|
|
| parallel_sentences.append(sentences) |
| count += 1 |
| if max_sentences is not None and max_sentences > 0 and count >= max_sentences: |
| break |
| self.add_dataset(parallel_sentences, weight=weight, max_sentences=max_sentences, max_sentence_length=max_sentence_length) |
|
|
|
|
| def add_dataset(self, parallel_sentences: List[List[str]], weight: int = 100, max_sentences: int = None, max_sentence_length: int = 128): |
| sentences_map = {} |
| for sentences in parallel_sentences: |
| if max_sentence_length is not None and max_sentence_length > 0 and max([len(sent) for sent in sentences]) > max_sentence_length: |
| continue |
|
|
| source_sentence = sentences[0] |
| if source_sentence not in sentences_map: |
| sentences_map[source_sentence] = set() |
|
|
| for sent in sentences: |
| sentences_map[source_sentence].add(sent) |
|
|
| if max_sentences is not None and max_sentences > 0 and len(sentences_map) >= max_sentences: |
| break |
|
|
| if len(sentences_map) == 0: |
| return |
|
|
| self.num_sentences += sum([len(sentences_map[sent]) for sent in sentences_map]) |
|
|
| dataset_id = len(self.datasets) |
| self.datasets.append(list(sentences_map.items())) |
| self.datasets_iterator.append(0) |
| self.dataset_indices.extend([dataset_id] * weight) |
|
|
|
|
| def generate_data(self): |
| source_sentences_list = [] |
| target_sentences_list = [] |
| for data_idx in self.dataset_indices: |
| src_sentence, trg_sentences = self.next_entry(data_idx) |
| source_sentences_list.append(src_sentence) |
| target_sentences_list.append(trg_sentences) |
|
|
| |
| src_embeddings = self.get_embeddings(source_sentences_list) |
|
|
| for src_embedding, trg_sentences in zip(src_embeddings, target_sentences_list): |
| for trg_sentence in trg_sentences: |
| self.cache.append(InputExample(texts=[trg_sentence], label=src_embedding)) |
|
|
| random.shuffle(self.cache) |
|
|
| def next_entry(self, data_idx): |
| source, target_sentences = self.datasets[data_idx][self.datasets_iterator[data_idx]] |
|
|
| self.datasets_iterator[data_idx] += 1 |
| if self.datasets_iterator[data_idx] >= len(self.datasets[data_idx]): |
| self.datasets_iterator[data_idx] = 0 |
| random.shuffle(self.datasets[data_idx]) |
|
|
| return source, target_sentences |
|
|
| def get_embeddings(self, sentences): |
| if not self.use_embedding_cache: |
| return self.teacher_model.encode(sentences, batch_size=self.batch_size, show_progress_bar=False, convert_to_numpy=True) |
|
|
| |
| new_sentences = [] |
| for sent in sentences: |
| if sent not in self.embedding_cache: |
| new_sentences.append(sent) |
|
|
| if len(new_sentences) > 0: |
| new_embeddings = self.teacher_model.encode(new_sentences, batch_size=self.batch_size, show_progress_bar=False, convert_to_numpy=True) |
| for sent, embedding in zip(new_sentences, new_embeddings): |
| self.embedding_cache[sent] = embedding |
|
|
| return [self.embedding_cache[sent] for sent in sentences] |
|
|
| def __len__(self): |
| return self.num_sentences |
|
|
| def __getitem__(self, idx): |
| if len(self.cache) == 0: |
| self.generate_data() |
|
|
| return self.cache.pop() |
|
|
