190 lines
7.5 KiB
Python
190 lines
7.5 KiB
Python
"""
|
||
homograph preprocessor
|
||
rewrites russian sentences containing stress homographs via local llm.
|
||
writes each sentence immediately — resume-safe if interrupted.
|
||
in --all mode: processes every sentence and reports discovered homographs.
|
||
|
||
usage:
|
||
python preprocess_homographs.py data/neutral-voice-dataset-list.txt
|
||
python preprocess_homographs.py data/*.txt --all
|
||
"""
|
||
|
||
import json
|
||
import re
|
||
import time
|
||
import argparse
|
||
import requests
|
||
from pathlib import Path
|
||
from collections import Counter
|
||
from tqdm import tqdm
|
||
|
||
# ── config ────────────────────────────────────────────────────────────────────
|
||
|
||
LLAMA_URL = "http://localhost:10000/completion"
|
||
SUFFIX = "-clean"
|
||
|
||
KNOWN_HOMOGRAPHS = {
|
||
"стоят", "стоит", "замок", "замки", "дома",
|
||
"мою", "моют", "пили", "пила", "вина",
|
||
"орган", "органы", "атлас", "полки", "полку",
|
||
"белки", "белка", "угол", "лечу", "лечит",
|
||
}
|
||
|
||
# ── prompts ───────────────────────────────────────────────────────────────────
|
||
|
||
SYSTEM = (
|
||
"Ты — редактор русских текстов. Ищешь омографы.\n\n"
|
||
|
||
"Омограф — ТОЛЬКО слово, у которого:\n"
|
||
"1) есть минимум 2 значения\n"
|
||
"2) значения различаются УДАРЕНИЕМ\n"
|
||
"3) написание одинаковое\n\n"
|
||
|
||
"Если ты не можешь явно назвать ОБА значения с разным ударением — это НЕ омограф.\n"
|
||
"Не угадывай. Если сомневаешься — не добавляй.\n\n"
|
||
|
||
"НЕ являются омографами:\n"
|
||
"- обычные глаголы (остыл, сделал, пошёл)\n"
|
||
"- обычные существительные (ведро, тишина)\n"
|
||
"- служебные слова (кто, это, в, и)\n"
|
||
"- слова без известной пары с другим ударением\n\n"
|
||
|
||
"Разрешённые ориентиры (примеры):\n"
|
||
"замок, стоят, плачу, мука, атлас, орган\n\n"
|
||
|
||
"Задача:\n"
|
||
"1) Найди омографы в исходном предложении\n"
|
||
"2) Перепиши предложение без них\n"
|
||
"3) Можно полностью перефразировать\n"
|
||
"4) В rewritten НЕ должно остаться омографов\n"
|
||
"5) Если их нет — верни текст без изменений\n\n"
|
||
|
||
"Перед добавлением слова в список задай себе вопрос: \"Какие ДВА значения и где ударение?\""
|
||
"Если ответа нет — не добавляй."
|
||
|
||
"Формат строго JSON:\n"
|
||
"{\"rewritten\": \"...\", \"homographs\": [\"...\"]}\n"
|
||
)
|
||
|
||
def build_prompt(sentence: str) -> str:
|
||
return (
|
||
f"<|begin_of_text|>"
|
||
f"<|start_header_id|>system<|end_header_id|>\n\n{SYSTEM}<|eot_id|>"
|
||
f"<|start_header_id|>user<|end_header_id|>\n\n{sentence}<|eot_id|>"
|
||
f"<|start_header_id|>assistant<|end_header_id|>\n\n"
|
||
)
|
||
|
||
# ── llm ───────────────────────────────────────────────────────────────────────
|
||
|
||
def call_llm(sentence: str) -> dict | None:
|
||
payload = {
|
||
"prompt": build_prompt(sentence),
|
||
"temperature": 0.2,
|
||
"top_p": 0.9,
|
||
"repeat_penalty": 1.1,
|
||
"enable_thinking": True,
|
||
"stop": ["<|eot_id|>", "<|end_of_text|>"],
|
||
"session": None,
|
||
"cache_prompt": False
|
||
}
|
||
try:
|
||
resp = requests.post(LLAMA_URL, json=payload, timeout=60)
|
||
resp.raise_for_status()
|
||
raw = resp.json().get("content", "").strip()
|
||
clean = re.sub(r"```json|```", "", raw).strip()
|
||
return json.loads(clean)
|
||
except Exception:
|
||
return None
|
||
|
||
# ── helpers ───────────────────────────────────────────────────────────────────
|
||
|
||
def has_known_homograph(sentence: str) -> bool:
|
||
words = re.sub(r"[^\w\s]", "", sentence.lower()).split()
|
||
return any(w in KNOWN_HOMOGRAPHS for w in words)
|
||
|
||
# ── main ──────────────────────────────────────────────────────────────────────
|
||
|
||
def process_file(path: Path, all_mode: bool):
|
||
lines = [l.strip() for l in path.read_text(encoding="utf-8").splitlines() if l.strip()]
|
||
out_path = path.parent / (path.stem + SUFFIX + path.suffix)
|
||
homo_out = path.parent / (path.stem + "-homographs.txt")
|
||
|
||
# resume support — count already written lines
|
||
done_count = 0
|
||
if out_path.exists():
|
||
done_count = sum(1 for l in out_path.read_text(encoding="utf-8").splitlines() if l.strip())
|
||
print(f" [*] resuming from line {done_count}/{len(lines)}")
|
||
|
||
todo = lines[done_count:]
|
||
if not todo:
|
||
print(f"[=] {path.name} already complete")
|
||
return
|
||
|
||
print(f"\n[>] {path.name} — {len(todo)} remaining ({done_count} done)")
|
||
|
||
changed = 0
|
||
failed = 0
|
||
discovered = Counter()
|
||
|
||
out_f = open(out_path, "a", encoding="utf-8")
|
||
homo_f = open(homo_out, "a", encoding="utf-8")
|
||
|
||
try:
|
||
for line in tqdm(todo):
|
||
if not all_mode and not has_known_homograph(line):
|
||
out_f.write(line + "\n")
|
||
out_f.flush()
|
||
continue
|
||
|
||
result = call_llm(line)
|
||
|
||
if result is None:
|
||
out_f.write(line + "\n")
|
||
out_f.flush()
|
||
failed += 1
|
||
continue
|
||
|
||
rewritten = result.get("rewritten", line).strip()
|
||
homographs = result.get("homographs", [])
|
||
|
||
for h in homographs:
|
||
if isinstance(h, str) and h.strip():
|
||
h = h.strip().lower()
|
||
discovered[h] += 1
|
||
homo_f.write(h + "\n")
|
||
homo_f.flush()
|
||
|
||
final = rewritten if (rewritten and rewritten != line) else line
|
||
if final != line:
|
||
changed += 1
|
||
|
||
out_f.write(final + "\n")
|
||
out_f.flush()
|
||
time.sleep(0.05)
|
||
|
||
finally:
|
||
out_f.close()
|
||
homo_f.close()
|
||
|
||
if discovered:
|
||
print(f"\n [*] {len(discovered)} unique homographs found → {homo_out}")
|
||
print(f" top 10: {', '.join(w for w, _ in discovered.most_common(10))}")
|
||
else:
|
||
print(f" [*] no new homographs discovered")
|
||
|
||
print(f" [✓] {changed} rewritten, {failed} failed → {out_path}")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Rewrite homograph sentences via local LLM")
|
||
parser.add_argument("files", nargs="+", help="txt files to process")
|
||
parser.add_argument("--all", action="store_true",
|
||
help="process ALL sentences and report discovered homographs")
|
||
args = parser.parse_args()
|
||
|
||
for f in args.files:
|
||
process_file(Path(f), all_mode=args.all)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |