141 lines
5.3 KiB
Python
141 lines
5.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Скрипт загружает .txt файлы из lists/, разбивает по разделителю,
|
|
отправляет каждую запись в Ollama, собирает JSON в res/.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import pathlib
|
|
import shutil
|
|
import sys
|
|
import time
|
|
|
|
from ollama import Client
|
|
|
|
# --- Конфигурация ---
|
|
LISTS_DIR = pathlib.Path(__file__).parent / "lists"
|
|
RES_DIR = pathlib.Path(__file__).parent / "res"
|
|
DONE_DIR = pathlib.Path(__file__).parent / "done"
|
|
OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://localhost:11434")
|
|
MODEL = os.environ.get("MODEL", "mithunrajm06/gemma-3n-e4b:latest")
|
|
NUM_CTX = int(os.environ.get("NUM_CTX", "4096"))
|
|
NUM_PREDICT = int(os.environ.get("NUM_PREDICT", "512"))
|
|
RETRY_COUNT = 3
|
|
RETRY_DELAY = 5
|
|
LOG_STREAM = os.environ.get("LOG_STREAM", "").lower() in ("1", "true", "yes")
|
|
# показывать вход и выход по каждой записи (0/false — отключить)
|
|
SHOW_IO = os.environ.get("SHOW_IO", "1").lower() in ("1", "true", "yes")
|
|
|
|
# Разделители записей при разбиении файла (\n — перенос строки, \r\n — Windows)
|
|
RECORD_DELIMITERS = ("\n", "\r\n", "\r")
|
|
|
|
EXTRACT_PROMPT = """Extract brand, model, engine_code, date_range from the text.
|
|
If 1 record — return JSON object: {"brand":"","model":"","engine_code":"","date_range":""}
|
|
If 2+ records on the line — return JSON array: [{"brand":"",...},{"brand":"",...}]
|
|
Empty fields as "". No other text."""
|
|
|
|
|
|
def prepare_records(content: str) -> list[str]:
|
|
"""1) Подготовка списка записей из файла по разделителям."""
|
|
for sep in RECORD_DELIMITERS:
|
|
content = content.replace(sep, RECORD_DELIMITERS[0])
|
|
return [r.strip() for r in content.split(RECORD_DELIMITERS[0]) if r.strip()]
|
|
|
|
|
|
def _normalize_obj(obj: dict) -> dict:
|
|
return {
|
|
"brand": obj.get("brand", ""),
|
|
"model": obj.get("model", ""),
|
|
"engine_code": obj.get("engine_code", ""),
|
|
"date_range": obj.get("date_range", ""),
|
|
}
|
|
|
|
|
|
def extract_record(client: Client, record: str) -> list[dict]:
|
|
"""2) Отправка промпт + 1 запись в нейронку. 3) Получение ответа (1 или более объектов)."""
|
|
for attempt in range(1, RETRY_COUNT + 1):
|
|
try:
|
|
response = client.chat(
|
|
model=MODEL,
|
|
messages=[
|
|
{"role": "system", "content": EXTRACT_PROMPT},
|
|
{"role": "user", "content": record},
|
|
],
|
|
options={"num_ctx": NUM_CTX, "num_predict": NUM_PREDICT},
|
|
)
|
|
result = response.get("message", {}).get("content", "")
|
|
if "```" in result:
|
|
result = result.split("```")[1]
|
|
if result.strip().lower().startswith("json"):
|
|
result = result.strip()[4:].strip()
|
|
parsed = json.loads(result)
|
|
if isinstance(parsed, list):
|
|
return [_normalize_obj(o) for o in parsed]
|
|
return [_normalize_obj(parsed)]
|
|
except (json.JSONDecodeError, KeyError):
|
|
if attempt < RETRY_COUNT:
|
|
time.sleep(RETRY_DELAY)
|
|
else:
|
|
return [{"brand": "", "model": "", "engine_code": "", "date_range": record}]
|
|
except Exception:
|
|
if attempt < RETRY_COUNT:
|
|
time.sleep(RETRY_DELAY)
|
|
else:
|
|
raise
|
|
return [{"brand": "", "model": "", "engine_code": "", "date_range": record}]
|
|
|
|
|
|
def process_file(client: Client, filepath: pathlib.Path) -> bool:
|
|
"""Обработка файла: подготовка → запрос по записи → добавление в JSON."""
|
|
content = filepath.read_text(encoding="utf-8")
|
|
records = prepare_records(content)
|
|
out_path = RES_DIR / filepath.with_suffix(".json").name
|
|
|
|
if not records:
|
|
print(f" {filepath.name} — 0 записей")
|
|
return False
|
|
|
|
out_path.write_text("[]", encoding="utf-8")
|
|
total = len(records)
|
|
t0 = time.perf_counter()
|
|
data = []
|
|
|
|
for i, record in enumerate(records, 1):
|
|
print(f" [{i}/{total}] ", end="", flush=True)
|
|
objs = extract_record(client, record)
|
|
data.extend(objs)
|
|
out_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
if SHOW_IO:
|
|
inp = record[:60] + "…" if len(record) > 60 else record
|
|
out = json.dumps(objs, ensure_ascii=False)
|
|
print(f"вход: {inp}", flush=True)
|
|
print(f" выход: {out}", flush=True)
|
|
else:
|
|
print("✓", flush=True)
|
|
elapsed = time.perf_counter() - t0
|
|
print(f" {filepath.name} — {len(data)} объектов, {elapsed:.1f}s — OK")
|
|
return True
|
|
|
|
|
|
def main() -> None:
|
|
RES_DIR.mkdir(exist_ok=True)
|
|
DONE_DIR.mkdir(exist_ok=True)
|
|
|
|
client = Client(host=OLLAMA_HOST)
|
|
txt_files = sorted(LISTS_DIR.glob("*.txt"))
|
|
|
|
if not txt_files:
|
|
print(f"Нет .txt файлов в {LISTS_DIR}")
|
|
return
|
|
|
|
print(f"Файлов к обработке: {len(txt_files)}")
|
|
|
|
for filepath in txt_files:
|
|
if process_file(client, filepath):
|
|
shutil.move(str(filepath), DONE_DIR / filepath.name)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|