benchmark.py

import requests
import pandas as pd
import time
from datasets import load_dataset
import sacrebleu
import matplotlib.pyplot as plt
from comet import download_model, load_from_checkpoint

# =========================
# CONFIG
# =========================
API_KEY = "API_KEY" # Replace with your actual API key from the TranslatePlus dashboard
BASE_URL = "https://api.translateplus.io/v2/translate"

LANG_PAIRS = [
    ("eng_Latn", "fra_Latn"),  # French
    ("eng_Latn", "deu_Latn"),  # German
    ("eng_Latn", "spa_Latn"),  # Spanish
    ("eng_Latn", "ita_Latn"),  # Italian
    ("eng_Latn", "por_Latn"),  # Portuguese
    ("eng_Latn", "nld_Latn"),  # Dutch
    ("eng_Latn", "swe_Latn"),  # Swedish
    ("eng_Latn", "dan_Latn"),  # Danish
    ("eng_Latn", "fin_Latn"),  # Finnish
    ("eng_Latn", "pol_Latn"),  # Polish

    ("eng_Latn", "rus_Cyrl"),  # Russian
    ("eng_Latn", "tur_Latn"),  # Turkish
    ("eng_Latn", "arb_Arab"),  # Arabic
    ("eng_Latn", "hin_Deva"),  # Hindi
    ("eng_Latn", "ben_Beng"),  # Bengali
    ("eng_Latn", "urd_Arab"),  # Urdu

    ("eng_Latn", "zho_Hans"),  # Chinese Simplified
    ("eng_Latn", "jpn_Jpan"),  # Japanese
    ("eng_Latn", "kor_Hang"),  # Korean
    ("eng_Latn", "vie_Latn"),  # Vietnamese
]

SAMPLE_SIZE = 997
SLEEP_TIME = 0

# =========================
# LANGUAGE MAPPING
# =========================
def flores_to_api_lang(code):
    overrides = {
        "eng_Latn": "en",
        "urd_Arab": "ur",
        "fra_Latn": "fr",
        "deu_Latn": "de",
        "spa_Latn": "es",
        "ita_Latn": "it",
        "por_Latn": "pt",
        "rus_Cyrl": "ru",
        "hin_Deva": "hi",
        "arb_Arab": "ar",
        "zho_Hans": "zh-CN",
        "zho_Hant": "zh-TW",
        "ben_Beng": "bn",
        "pan_Guru": "pa",
        "tam_Taml": "ta",
        "tel_Telu": "te",
        "mar_Deva": "mr",
        "jpn_Jpan": "ja",
        "kor_Hang": "ko",
        "vie_Latn": "vi",
        "tha_Thai": "th",
        "tur_Latn": "tr",
        "pol_Latn": "pl",
        "nld_Latn": "nl",
        "swe_Latn": "sv",
        "dan_Latn": "da",
        "fin_Latn": "fi",
        "ell_Grek": "el",
        "ces_Latn": "cs",
        "ron_Latn": "ro",
        "hun_Latn": "hu",
        "ukr_Cyrl": "uk",
    }

    if code in overrides:
        return overrides[code]

    base = code.split("_")[0]
    fallback = {
        "eng": "en", "urd": "ur", "fra": "fr",
        "deu": "de", "spa": "es", "ita": "it",
        "por": "pt", "rus": "ru", "hin": "hi", "arb": "ar"
    }

    return fallback.get(base)

# =========================
# TRANSLATE FUNCTION
# =========================
def translate(text, source, target):
    try:
        response = requests.post(
            BASE_URL,
            headers={"X-API-KEY": API_KEY},
            json={
                "text": text,
                "source": source,
                "target": target
            },
            timeout=10
        )

        data = response.json()

        # ✅ correct parsing
        translation = data.get("translations", {}).get("translation")

        if not translation:
            raise ValueError(f"Invalid API response: {data}")

        return translation

    except Exception as e:
        print("Error:", e)
        return ""

# =========================
# LOAD DATASET ONCE
# =========================
print("Loading FLORES dataset...")
dataset = load_dataset("facebook/flores", "all")

# =========================
# LOAD COMET MODEL
# =========================
print("Loading COMET model...")
model_path = download_model("Unbabel/wmt22-comet-da")
comet_model = load_from_checkpoint(model_path)

# =========================
# MAIN LOOP
# =========================
summary = []

for src, tgt in LANG_PAIRS:
    print(f"\n=== {src} → {tgt} ===")

    source_lang = flores_to_api_lang(src)
    target_lang = flores_to_api_lang(tgt)

    if not source_lang or not target_lang:
        print(f"Skipping unsupported pair: {src}-{tgt}")
        continue

    data = dataset["dev"].select(range(SAMPLE_SIZE))

    source_col = f"sentence_{src}"
    target_col = f"sentence_{tgt}"

    results = []
    total_latency = 0

    for i, row in enumerate(data):
        source_text = row[source_col]
        reference_text = row[target_col]

        start = time.time()
        translated = translate(source_text, source_lang, target_lang)
        latency = time.time() - start

        total_latency += latency

        results.append({
            "source": source_text,
            "reference": reference_text,
            "hypothesis": translated,
            "latency": latency
        })

        print(f"{i+1}/{SAMPLE_SIZE}", end="\r")
        time.sleep(SLEEP_TIME)

    # =========================
    # METRICS
    # =========================
    refs = [r["reference"] for r in results]
    hyps = [r["hypothesis"] for r in results]
    srcs = [r["source"] for r in results]

    bleu = sacrebleu.corpus_bleu(hyps, [refs]).score

    comet_data = [
        {"src": s, "mt": h, "ref": r}
        for s, h, r in zip(srcs, hyps, refs)
    ]

    comet_score = comet_model.predict(
        comet_data,
        batch_size=8,
        gpus=0
    )["system_score"]

    avg_latency = total_latency / SAMPLE_SIZE

    print(f"\nBLEU: {bleu:.2f}")
    print(f"COMET: {comet_score:.4f}")
    print(f"Latency: {avg_latency:.3f}s")

    # SAVE
    df = pd.DataFrame(results)
    df.to_csv(f"results_{src}_{tgt}.csv", index=False)

    summary.append({
        "pair": f"{src}-{tgt}",
        "BLEU": bleu,
        "COMET": comet_score,
        "Latency": avg_latency
    })

# =========================
# SUMMARY + CHARTS
# =========================
summary_df = pd.DataFrame(summary)
summary_df.to_csv("benchmark_summary.csv", index=False)

# BLEU
plt.figure()
plt.bar(summary_df["pair"], summary_df["BLEU"])
plt.title("BLEU Scores")
plt.xticks(rotation=45)
plt.tight_layout()
plt.savefig("bleu_scores.png")

# COMET
plt.figure()
plt.bar(summary_df["pair"], summary_df["COMET"])
plt.title("COMET Scores")
plt.xticks(rotation=45)
plt.tight_layout()
plt.savefig("comet_scores.png")

# LATENCY
plt.figure()
plt.bar(summary_df["pair"], summary_df["Latency"])
plt.title("Latency (seconds)")
plt.xticks(rotation=45)
plt.tight_layout()
plt.savefig("latency.png")

print("\n✅ DONE")
print(summary_df)