model-training/llm/append_english.py

"""Phase 6c — generate English persona samples, append to persona_train.jsonl.

Reuses gen_data.py's CANONICAL_SYSTEM_PROMPT (which handles English input),
TOPIC_SEEDS_ENGLISH, and validation. Runs independently.

Usage:
    python append_english.py
"""
import json
import os
import sys
import random
from itertools import cycle
from pathlib import Path

HERE = Path(__file__).resolve().parent
DATA = HERE / "data"
random.seed(42)

sys.path.insert(0, str(HERE))
import gen_data

CANONICAL_SYSTEM_PROMPT = gen_data.CANONICAL_SYSTEM_PROMPT
TOPIC_SEEDS_ENGLISH = gen_data.TOPIC_SEEDS_ENGLISH
VALID_MOODS = gen_data.VALID_MOODS
validate_sample = gen_data.validate_sample

ROUTER_URL = os.environ.get(
    "INFERENCE_ROUTER_URL", "http://inference.kvmx.ru"
).rstrip("/")

ROUTER_PAIRS = [
    ("cerebras", "gemma-4-31b"),
    ("cerebras", "zai-glm-4.7"),
    ("github_models", "models/meta/llama-4"),
    ("github_models", "models/mistral/mistral-large-3"),
    ("mistral", "mistral-large-latest"),
    ("mistral", "codestral-latest"),
]


def _chat(provider: str, model: str, msgs: list,
          max_tokens: int = 300, temp: float = 0.8) -> str | None:
    import time as _time
    import urllib.request as _req
    import urllib.error as _err
    base = ROUTER_URL
    if not base.endswith("/v1"):
        base += "/v1"
    url = f"{base}/chat/completions"
    headers = {"Content-Type": "application/json", "X-Provider": provider}
    payload = json.dumps({
        "model": model, "messages": msgs,
        "max_tokens": max_tokens, "temperature": temp,
    }).encode()
    for attempt in range(3):
        try:
            r = _req.Request(url, data=payload, headers=headers)
            with _req.urlopen(r, timeout=30) as resp:
                body = json.loads(resp.read())
                content = (body["choices"][0]["message"]["content"] or "").strip()
            if content.startswith("```"):
                start = content.find("\n") + 1
                end = content.rfind("```")
                if end > start:
                    content = content[start:end].strip()
            return content
        except Exception:
            _time.sleep(min(2 ** attempt, 8))
    return None


def _parse_asst(content: str) -> dict | None:
    try:
        obj = json.loads(content)
    except json.JSONDecodeError:
        return None
    if not isinstance(obj, dict) or "response" not in obj or "mood" not in obj:
        return None
    if obj["mood"] not in VALID_MOODS or not isinstance(obj["response"], str):
        return None
    return obj


def flatten_msgs(user_text: str, asst_text: str) -> list[dict]:
    return [
        {"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
        {"role": "user", "content": user_text},
        {"role": "assistant", "content": asst_text},
    ]


def gen_english(target: int) -> list[dict]:
    """Generate English persona samples."""
    samples = []
    providers = cycle(ROUTER_PAIRS)
    pool = TOPIC_SEEDS_ENGLISH * (target // len(TOPIC_SEEDS_ENGLISH) + 1)
    random.shuffle(pool)

    print(f"  english: targeting {target} samples...")
    for idx, topic in enumerate(pool[:target]):
        provider, model = next(providers)
        msgs = [
            {"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
            {"role": "user", "content": topic},
        ]
        content = _chat(provider, model, msgs, max_tokens=300, temp=0.8)
        if content is None:
            continue
        asst_obj = _parse_asst(content)
        if asst_obj is None:
            continue
        # EN→EN bucket: reject majority-Cyrillic responses (language-mirroring failure).
        resp = asst_obj["response"]
        cyr = sum(1 for c in resp if "Ѐ" <= c <= "ӿ")
        lat = sum(1 for c in resp if c.isascii() and c.isalpha())
        if cyr > lat:
            continue

        asst_text = json.dumps(asst_obj, ensure_ascii=False)
        training_msgs = flatten_msgs(topic, asst_text)
        err = validate_sample(training_msgs)
        if err:
            continue
        samples.append({
            "messages": training_msgs,
            "mood": asst_obj["mood"],
            "response": asst_obj["response"],
            "source": "generated/english",
        })
        if len(samples) % 25 == 0:
            print(f"    ...{len(samples)}/{target} english valid")

    return samples


def main():
    target = 300
    print(f"Generating up to {target} English persona samples via router...")

    samples = gen_english(target)
    print(f"\nGot {len(samples)} valid English samples")

    if not samples:
        print("No samples generated.")
        return

    # dedup against existing train
    train_path = DATA / "persona_train.jsonl"
    existing: set[str] = set()
    if train_path.exists():
        with open(train_path) as f:
            for line in f:
                try:
                    d = json.loads(line)
                    msg = d.get("messages", [])
                    if len(msg) >= 3:
                        asst_raw = json.loads(msg[2]["content"])
                        existing.add(asst_raw.get("response", "").strip().lower())
                except Exception:
                    pass
    before = len(samples)
    samples = [s for s in samples if s["response"].strip().lower() not in existing]
    print(f"Dedup removed {before - len(samples)} (already in train file)")

    # dedup within
    seen: set[str] = set()
    deduped = []
    for s in samples:
        key = s["response"].strip().lower()
        if key not in seen:
            seen.add(key)
            deduped.append(s)
    samples = deduped

    with open(train_path, "a", encoding="utf-8") as f:
        for s in samples:
            f.write(json.dumps({"messages": s["messages"]}, ensure_ascii=False) + "\n")

    print(f"Appended {len(samples)} English samples to {train_path}")


if __name__ == "__main__":
    main()