python/mt_engine.py

#!/usr/bin/env python3
"""Prompting + decode settings for Qwen3-4B translation."""
from __future__ import annotations

import re
import time

from openai import OpenAI

SYSTEM = (
    "You are a professional translator for invoices, notices, medical instructions, "
    "and news. Translate faithfully into the requested target language.\n"
    "Rules:\n"
    "- Output only the translation. No quotes, labels, romanization notes, or commentary.\n"
    "- Copy the listed preserve tokens exactly (IDs, emails, amounts, Latin brand names, acronyms).\n"
    "- If a preserve token is already Latin, keep it in Latin letters; do not transliterate it.\n"
    "- Do not replace a proper name with a similar-sounding different name.\n"
    "- receipt ≠ invoice. receipt → Arabic إيصال, French reçu, German Quittung, Spanish recibo. "
    "invoice → Arabic فاتورة, French facture, German Rechnung, Spanish factura, Italian fattura, "
    "Portuguese fatura, Turkish fatura, Greek τιμολόγιο, Russian счёт.\n"
    "- digital identity → Arabic هوية رقمية (not اعتماد).\n"
    "- In a banking context, Arabic حوّل / حول المبلغ إلى الحساب means transfer (wire) money, "
    "not convert a currency, unless the source is clearly about an exchange rate.\n"
    "- Keep technical acronyms GPU, OCR, CPU, CUDA, API in Latin letters when the source uses them.\n"
    "- Modern Standard Arabic; keep the definite article on city names (الرياض، الدوحة، أبوظبي).\n"
    "- Translate meaning, but never drop a number, date, or identifier."
)

_ID = re.compile(r"\b[A-Z]{2,}[A-Z0-9._/-]*\d[A-Z0-9._/-]*\b")
_EMAIL = re.compile(r"\b[\w.+-]+@[\w.-]+\.[A-Za-z]{2,}\b")
_ACRONYM = re.compile(
    r"\b(?:GPU|OCR|CPU|CUDA|API|AWB|LLC|USD|AED|EUR|VAT|TTC|IBAN|SWIFT|SKU|PO)\b"
)
_HYPHEN_NAME = re.compile(r"\b[A-Z][a-z]+-[A-Z][a-z]+\b")
_NAME_RUN = re.compile(
    r"\b[A-Z][a-z]+(?:\s+(?:and|&|of|de|del|al|the)\s+[A-Z][a-z]+)*(?:\s+[A-Z][a-z]+){1,3}\b"
)
_SKIP_NAMES = {
    "The Total", "The Patient", "Please Reset", "Please Send", "Please Pay",
    "On August", "Each Page", "This Stainless", "Dear Mr",
}


def copy_spans(source: str) -> list[str]:
    found: list[str] = []
    for rx in (_EMAIL, _ID, _ACRONYM, _HYPHEN_NAME):
        found.extend(rx.findall(source))
    for m in _NAME_RUN.finditer(source):
        span = m.group(0)
        if span in _SKIP_NAMES:
            continue
        if span.lower().startswith(("the ", "this ", "please ", "dear ")):
            continue
        found.append(span)
    out: list[str] = []
    for tok in found:
        if tok not in out:
            out.append(tok)
    return out[:12]


def term_hints(source: str, src_lang: str, tgt_lang: str) -> list[str]:
    hints: list[str] = []
    sl, tl = src_lang.lower(), tgt_lang.lower()
    low = source.lower()
    if tl == "arabic":
        if re.search(r"\breceipts?\b", low):
            hints.append("receipt → إيصال")
        if re.search(r"\binvoices?\b", low):
            hints.append("invoice → فاتورة")
        if "identity" in low:
            hints.append("digital identity → هوية رقمية")
        if re.search(r"\b(gpu|ocr|cuda)\b", low):
            hints.append("keep GPU / OCR / CUDA as Latin acronyms")
        if "cafe" in low or "café" in low:
            hints.append("Cafe in a brand name stays Latin (Harbor and Rye Cafe)")
        if "tablet" in low:
            hints.append("tablet → قرص")
        if "customs" in low:
            hints.append("customs officer → موظف الجمارك")
        if re.search(r"\bgoods\b", low):
            hints.append("goods → البضائع")
        if "finance" in low:
            hints.append("Finance department → قسم المالية")
        if "flight" in low:
            hints.append("flight → رحلة")
    if sl == "arabic" and tl == "english":
        if re.search(r"حو[ّ]?ل", source) and "حساب" in source:
            hints.append("حوّل المبلغ إلى الحساب = transfer (not convert)")
        if "إيصال" in source:
            hints.append("إيصال = receipt")
        if "فاتورة" in source:
            hints.append("فاتورة = invoice")
        if "هوية" in source:
            hints.append("هوية = identity")
        if "رحلة" in source and "مطار" in source:
            hints.append("الرحلة from an airport = flight (not trip)")
    if tl == "greek" and "invoice" in low:
        hints.append("invoice → τιμολόγιο")
    if tl == "russian" and "invoice" in low:
        hints.append("invoice → счёт")
    return hints


def prompt_for(case: dict) -> tuple[str, str]:
    src, tgt = case["src_lang"], case["tgt_lang"]
    preserve = list(case.get("preserve") or [])
    for span in copy_spans(case["source"]):
        if span not in preserve:
            preserve.append(span)
    hints = term_hints(case["source"], src, tgt)
    lines = [
        f"Translate from {src} to {tgt}.",
        "Output only the translation.",
    ]
    if preserve:
        lines.append(
            "Your translation MUST contain each of these strings verbatim: "
            + ", ".join(preserve)
            + "."
        )
    if hints:
        lines.append("Term hints: " + "; ".join(hints) + ".")
    lines.append("")
    lines.append(case["source"])
    return SYSTEM, "\n".join(lines)


def max_tokens_for(source: str) -> int:
    return min(512, max(96, 48 + 4 * max(1, len(source.split()))))


def strip_fences(text: str) -> str:
    text = (text or "").strip()
    text = re.sub(r"^```(?:\w+)?\s*", "", text)
    text = re.sub(r"\s*```$", "", text)
    text = re.sub(
        r"^(?:translation|arabic|english|french|german|spanish|italian|portuguese|"
        r"turkish|russian|greek)\s*:\s*",
        "",
        text,
        flags=re.I,
    )
    return text.strip().strip('"“”«»')


def translate_text(
    client: OpenAI,
    source: str,
    src_lang: str,
    tgt_lang: str,
    model: str = "Qwen/Qwen3-4B-Instruct-2507",
    preserve: list[str] | None = None,
) -> tuple[str, float]:
    return translate(
        client,
        model,
        {
            "source": source,
            "src_lang": src_lang,
            "tgt_lang": tgt_lang,
            "preserve": preserve or [],
        },
    )


def translate(client: OpenAI, model: str, case: dict) -> tuple[str, float]:
    system, user = prompt_for(case)
    messages = [
        {"role": "system", "content": system},
        {"role": "user", "content": user},
    ]
    kwargs = dict(
        model=model,
        temperature=0.0,
        top_p=1.0,
        max_tokens=max_tokens_for(case["source"]),
        messages=messages,
    )
    t0 = time.perf_counter()
    try:
        resp = client.chat.completions.create(**kwargs, seed=0)
    except TypeError:
        resp = client.chat.completions.create(**kwargs)
    dt = time.perf_counter() - t0
    return strip_fences(resp.choices[0].message.content or ""), dt