SkyWater21/lv_go_emotions

from datasets import load_dataset
from deep_translator import GoogleTranslator
from deep_translator.exceptions import TranslationNotFound

original_dataset = load_dataset("go_emotions", name="simplified")
translator = GoogleTranslator(source="en", target="lv")

def translate_batch(batch):
    original_text = batch["text"]

    while True:
        try:
            translated_batch = translator.translate_batch(original_text)
            break
        except TranslationNotFound:
            # Translation can fail due to API limits, so we retry until it works
            print(f"Translation failed. Retrying...")
    
    # We fix untranslated entries (None values) by replacing them with the original text
    for i in range(len(translated_batch)):
        if not translated_batch[i]:
            translated_batch[i] = original_text[i]
            print(f"Replaced {original_text[i]} vs {translated_batch[i]}")
            
    batch["lv_text"] = translated_batch
    
    return batch

translated_dataset = original_dataset.map(
    translate_batch, batched=True, batch_size=500
)

0: admiration
1: amusement
2: anger
3: annoyance
4: approval
5: caring
6: confusion
7: curiosity
8: desire
9: disappointment
10: disapproval
11: disgust
12: embarrassment
13: excitement
14: fear
15: gratitude
16: grief
17: joy
18: love
19: nervousness
20: optimism
21: pride
22: realization
23: relief
24: remorse
25: sadness
26: surprise
27: neutral

0: anger
1: disgust
2: fear
3: joy
4: sadness
5: surprise
6: neutral