Function collate_and_encode

sonar_encoding.py:48–87

Tokenize each text into sentences, batch the sentences through Sonar, and regroup them per text. Returns a list of dicts, each dict has: { "original_text": str, "sentences": [str, ...], "embeddings": [ [float, ..., float], ... ] # 1024-D }

(batch_texts, sonar_model)

def collate_and_encode(batch_texts, sonar_model):
    """
    Tokenize each text into sentences, batch the sentences through Sonar,
    and regroup them per text.
    Returns a list of dicts, each dict has:
      {
         "original_text": str,
         "sentences": [str, ...],
         "embeddings": [ [float, ..., float], ... ]  # 1024-D
      }
    """
    batch_sentences = []
    batch_map = [] 

    for i, txt in enumerate(batch_texts):
        sents = sent_tokenize(txt)
        sents = [s.strip() for s in sents if s.strip()]
        for j, s in enumerate(sents):
            batch_sentences.append(s)
            batch_map.append((i, j))

    if len(batch_sentences) > 0:
        with torch.no_grad():
            emb_tensors = sonar_model.predict(batch_sentences, source_lang="eng_Latn")
        all_embeddings = emb_tensors.cpu().tolist()
    else:
        all_embeddings = []

    results = [
        {"original_text": batch_texts[i], "sentences": [], "embeddings": []}
        for i in range(len(batch_texts))
    ]

    idx_in_sentences = 0
    for (text_i, _), emb_vec in zip(batch_map, all_embeddings):
        results[text_i]["sentences"].append(batch_sentences[idx_in_sentences])
        results[text_i]["embeddings"].append(emb_vec)
        idx_in_sentences += 1

    return results


def ddp_main():