PRLM/distilabel-intel-orca-dpo-pairs-balanced-subsets-translated

dataset_info:

• config_name: 1k-English
  features:
  • name: prompt
    dtype: string
  • name: completion
    dtype: string
  • name: score
    dtype: int64
    splits:
  • name: train
    num_bytes: 1649336
    num_examples: 1000
    download_size: 988517
    dataset_size: 1649336
• config_name: 1k-French
  features:
  • name: prompt
    dtype: string
  • name: completion
    dtype: string
  • name: score
    dtype: int64
    splits:
  • name: train
    num_bytes: 1622654
    num_examples: 1000
    download_size: 937396
    dataset_size: 1622654
• config_name: 1k-German
  features:
  • name: prompt
    dtype: string
  • name: completion
    dtype: string
  • name: score
    dtype: int64
    splits:
  • name: train
    num_bytes: 1540572
    num_examples: 1000
    download_size: 905410
    dataset_size: 1540572
• config_name: 1k-Hindi
  features:
  • name: prompt
    dtype: string
  • name: completion
    dtype: string
  • name: score
    dtype: int64
    splits:
  • name: train
    num_bytes: 3592006
    num_examples: 1000
    download_size: 1326017
    dataset_size: 3592006
• config_name: 1k-Italian
  features:
  • name: prompt
    dtype: string
  • name: completion
    dtype: string
  • name: score
    dtype: int64
    splits:
  • name: train
    num_bytes: 1521293
    num_examples: 1000
    download_size: 891313
    dataset_size: 1521293
• config_name: 1k-Portuguese
  features:
  • name: prompt
    dtype: string
  • name: completion
    dtype: string
  • name: score
    dtype: int64
    splits:
  • name: train
    num_bytes: 1504355
    num_examples: 1000
    download_size: 897872
    dataset_size: 1504355
• config_name: 1k-Spanish
  features:
  • name: prompt
    dtype: string
  • name: completion
    dtype: string
  • name: score
    dtype: int64
    splits:
  • name: train
    num_bytes: 1518093
    num_examples: 1000
    download_size: 881741
    dataset_size: 1518093
• config_name: 1k-Thai
  features:
  • name: prompt
    dtype: string
  • name: completion
    dtype: string
  • name: score
    dtype: int64
    splits:
  • name: train
    num_bytes: 3432191
    num_examples: 1000
    download_size: 1125403
    dataset_size: 3432191
    configs:
• config_name: 1k-English
  data_files:
  • split: train
    path: 1k-English/train-*
• config_name: 1k-French
  data_files:
  • split: train
    path: 1k-French/train-*
• config_name: 1k-German
  data_files:
  • split: train
    path: 1k-German/train-*
• config_name: 1k-Hindi
  data_files:
  • split: train
    path: 1k-Hindi/train-*
• config_name: 1k-Italian
  data_files:
  • split: train
    path: 1k-Italian/train-*
• config_name: 1k-Portuguese
  data_files:
  • split: train
    path: 1k-Portuguese/train-*
• config_name: 1k-Spanish
  data_files:
  • split: train
    path: 1k-Spanish/train-*
• config_name: 1k-Thai
  data_files:
  • split: train
    path: 1k-Thai/train-*

from concurrent.futures import ThreadPoolExecutor, as_completed
from datasets import load_dataset, Dataset
from typing import Optional
from threading import Lock
from tqdm import tqdm
import cohere
import time
import os


PROMPT_TEMPLATE = """Translate the following text from English to {target_language}.
Provide only the translated text without any additional explanation or formatting:
'{prompt}'"""

request_counter = 0
counter_lock = Lock()


def generate_w_cohere(prompt: str, model: str = "c4ai-aya-expanse-8b"):
    global request_counter
    with counter_lock:
        if request_counter >= 475:  # Rate limit for Cohere API is 500 requests per minute
            time.sleep(60)
            request_counter = 0
        request_counter += 1

    co = cohere.ClientV2(os.environ["COHERE_API_KEY"])
    response = co.chat(
        model=model,
        messages=[{"role": "user", "content": prompt}],
    )
    return response.message.content[0].text


def translate_w_cohere(
    dataset_path: str="PRLM/distilabel-intel-orca-dpo-pairs-balanced-subsets",
    subset: str="1k",
    split: str="train", 
    prompt_column: str="prompt", 
    response_column: str="completion",
    score_column: str="score",
):
    ds = load_dataset(dataset_path, subset, split=split).select_columns([prompt_column, response_column, score_column])

    for language in ["German", "Italian", "Portuguese", "Hindi", "Spanish", "Thai"]:  # "English", "French", 
        if language != "English":
            prompts = [PROMPT_TEMPLATE.format(target_language=language, prompt=prompt) for prompt in ds[prompt_column]]
            responses = [PROMPT_TEMPLATE.format(target_language=language, prompt=response) for response in ds[response_column]]
            
            prompt_translations = [None] * len(prompts)
            response_translations = [None] * len(responses)

            with ThreadPoolExecutor() as executor:
                prompt_futures = {executor.submit(generate_w_cohere, prompt): prompt for prompt in prompts}
                response_futures = {executor.submit(generate_w_cohere, response): response for response in responses}

                for future in tqdm(as_completed(prompt_futures), total=len(prompts), desc=f"Translating prompts to {language}", unit="prompt"):
                    index = prompts.index(prompt_futures[future])
                    prompt_translations[index] = future.result()

                for future in tqdm(as_completed(response_futures), total=len(responses), desc=f"Translating responses to {language}", unit="response"):
                    index = responses.index(response_futures[future])
                    response_translations[index] = future.result()
        else:
            prompt_translations = ds[prompt_column]
            response_translations = ds[response_column]

        translated_dataset = Dataset.from_dict({
            prompt_column: prompt_translations,
            response_column: response_translations,
            score_column: ds[score_column],
        })
        translated_dataset.push_to_hub(
            dataset_path+"-translated",
            subset+"-"+language,
            split="train", 
            private=True,
            token=os.environ["HF_TOKEN"],
        )


if __name__ == "__main__":
    import fire
    fire.Fire(translate_w_cohere)