264 lines
9.4 KiB
Python
264 lines
9.4 KiB
Python
"""Phase 6b — generate failure samples (tired/confused) and append to persona_train.jsonl.
|
|
|
|
Borrows CANONICAL_SYSTEM_PROMPT, failure prompts, and validation from gen_data.py.
|
|
Runs independently so we can iterate on failure prompting without touching the
|
|
persona generation pipeline.
|
|
|
|
Usage:
|
|
python append_failures.py
|
|
|
|
Environment:
|
|
INFERENCE_ROUTER_URL (default: http://inference.kvmx.ru)
|
|
"""
|
|
import json
|
|
import os
|
|
import sys
|
|
import random
|
|
import re
|
|
from itertools import cycle
|
|
from pathlib import Path
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
DATA = HERE / "data"
|
|
|
|
random.seed(42)
|
|
|
|
# steal the canonical system prompt and failure prompts from gen_data.py
|
|
sys.path.insert(0, str(HERE))
|
|
import gen_data # noqa: E402
|
|
|
|
CANONICAL_SYSTEM_PROMPT = gen_data.CANONICAL_SYSTEM_PROMPT
|
|
FAILURE_PROMPTS_TIRED = gen_data.FAILURE_PROMPTS_TIRED
|
|
FAILURE_PROMPTS_CONFUSED = gen_data.FAILURE_PROMPTS_CONFUSED
|
|
VALID_MOODS = gen_data.VALID_MOODS
|
|
validate_sample = gen_data.validate_sample
|
|
|
|
# ── router config ────────────────────────────────────────────────────────
|
|
ROUTER_URL = os.environ.get(
|
|
"INFERENCE_ROUTER_URL", "http://inference.kvmx.ru"
|
|
).rstrip("/")
|
|
|
|
# providers to cycle through — brief list of models that work well for RU
|
|
ROUTER_PAIRS = [
|
|
# format: (provider, model) — X-Provider header + bare model name
|
|
("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"),
|
|
]
|
|
|
|
|
|
# ── router call ──────────────────────────────────────────────────────────
|
|
def _chat(provider: str, model: str, msgs: list,
|
|
max_tokens: int = 200, temp: float = 0.7) -> 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"}
|
|
headers["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 _err.HTTPError as e:
|
|
_time.sleep(min(2 ** attempt, 8))
|
|
except Exception:
|
|
_time.sleep(1)
|
|
return None
|
|
|
|
|
|
def _parse_asst(content: str) -> dict | None:
|
|
try:
|
|
obj = json.loads(content)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
if not isinstance(obj, dict):
|
|
return None
|
|
if "response" not in obj or "mood" not in obj:
|
|
return None
|
|
if obj["mood"] not in VALID_MOODS:
|
|
return None
|
|
if 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},
|
|
]
|
|
|
|
|
|
# ── generation ───────────────────────────────────────────────────────────
|
|
def gen_failures(count: int) -> list[dict]:
|
|
"""Generate tired + confused samples, return list of validated dicts."""
|
|
samples = []
|
|
providers = cycle(ROUTER_PAIRS)
|
|
|
|
# tired: 2/3 of count
|
|
tired_target = count * 2 // 3
|
|
pool_tired = FAILURE_PROMPTS_TIRED * (tired_target // len(FAILURE_PROMPTS_TIRED) + 1)
|
|
random.shuffle(pool_tired)
|
|
|
|
print(f" tired: targeting {tired_target} samples...")
|
|
done = 0
|
|
for prompt_text in pool_tired[:tired_target]:
|
|
provider, model = next(providers)
|
|
msgs = [
|
|
{"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
|
|
{"role": "user", "content": (
|
|
f"{prompt_text}\n\n"
|
|
f"(Это вопрос, на который у тебя нет ответа. "
|
|
f"Не выдумывай. Честно скажи, что не знаешь или не можешь "
|
|
f"ответить. Используй mood: tired. Только tired, не neutral, "
|
|
f"не confused, не thinking — именно tired.)"
|
|
)},
|
|
]
|
|
content = _chat(provider, model, msgs, max_tokens=200, temp=0.7)
|
|
if content is None:
|
|
continue
|
|
asst_obj = _parse_asst(content)
|
|
if asst_obj is None or asst_obj["mood"] != "tired":
|
|
continue
|
|
|
|
asst_text = json.dumps(asst_obj, ensure_ascii=False)
|
|
training_msgs = flatten_msgs(prompt_text, asst_text)
|
|
err = validate_sample(training_msgs)
|
|
if err:
|
|
continue
|
|
samples.append({
|
|
"messages": training_msgs,
|
|
"mood": "tired",
|
|
"response": asst_obj["response"],
|
|
"source": "generated/failure",
|
|
})
|
|
done += 1
|
|
if done % 25 == 0:
|
|
print(f" ...{done}/{tired_target} tired valid")
|
|
|
|
print(f" tired: {done} valid")
|
|
|
|
# confused: 1/3 of count
|
|
confused_target = count // 3
|
|
pool_confused = FAILURE_PROMPTS_CONFUSED * (confused_target // len(FAILURE_PROMPTS_CONFUSED) + 1)
|
|
random.shuffle(pool_confused)
|
|
|
|
print(f" confused: targeting {confused_target} samples...")
|
|
done = 0
|
|
for prompt_text in pool_confused[:confused_target]:
|
|
provider, model = next(providers)
|
|
msgs = [
|
|
{"role": "system", "content": CANONICAL_SYSTEM_PROMPT},
|
|
{"role": "user", "content": (
|
|
f"{prompt_text}\n\n"
|
|
f"(Это неясный или слишком общий запрос. Ответь коротким "
|
|
f"уточняющим вопросом. Используй mood: confused. Только "
|
|
f"confused, не neutral, не tired, не thinking — именно confused.)"
|
|
)},
|
|
]
|
|
content = _chat(provider, model, msgs, max_tokens=200, temp=0.7)
|
|
if content is None:
|
|
continue
|
|
asst_obj = _parse_asst(content)
|
|
if asst_obj is None or asst_obj["mood"] != "confused":
|
|
continue
|
|
|
|
asst_text = json.dumps(asst_obj, ensure_ascii=False)
|
|
training_msgs = flatten_msgs(prompt_text, asst_text)
|
|
err = validate_sample(training_msgs)
|
|
if err:
|
|
continue
|
|
samples.append({
|
|
"messages": training_msgs,
|
|
"mood": "confused",
|
|
"response": asst_obj["response"],
|
|
"source": "generated/failure",
|
|
})
|
|
done += 1
|
|
if done % 25 == 0:
|
|
print(f" ...{done}/{confused_target} confused valid")
|
|
|
|
print(f" confused: {done} valid")
|
|
return samples
|
|
|
|
|
|
def main():
|
|
# target: 750 failure samples (15% of 5000)
|
|
target = 750
|
|
print(f"Generating up to {target} failure samples via router...")
|
|
print(f" Router: {ROUTER_URL}")
|
|
print(f" Models: {[f'{p}/{m}' for p,m in ROUTER_PAIRS]}")
|
|
|
|
samples = gen_failures(target)
|
|
print(f"\nGot {len(samples)} valid failure samples")
|
|
|
|
if not samples:
|
|
print("No samples generated — nothing to append.")
|
|
return
|
|
|
|
# dedup by response text against existing train data
|
|
train_path = DATA / "persona_train.jsonl"
|
|
if train_path.exists():
|
|
existing_responses: set[str] = set()
|
|
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_responses.add(asst_raw.get("response", "").strip().lower())
|
|
except (json.JSONDecodeError, KeyError, IndexError):
|
|
pass
|
|
before = len(samples)
|
|
samples = [s for s in samples if s["response"].strip().lower() not in existing_responses]
|
|
print(f"Dedup removed {before - len(samples)} (already in train file)")
|
|
|
|
# dedup within new samples
|
|
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
|
|
|
|
# append to persona_train.jsonl
|
|
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)} failure samples to {train_path}")
|
|
print(f"Mood distribution: tired={sum(1 for s in samples if s['mood']=='tired')}, "
|
|
f"confused={sum(1 for s in samples if s['mood']=='confused')}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|