112 lines
4.2 KiB
Python
112 lines
4.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Скрипт загружает .txt файлы из lists/, отправляет каждый в локальную Ollama
|
|
с системным промптом и сохраняет ответ в res/. Успешные файлы переносятся в done/.
|
|
При падении можно перезапустить — необработанные остаются в lists/.
|
|
"""
|
|
|
|
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", "gemma3n:e4b")
|
|
RETRY_COUNT = 3
|
|
RETRY_DELAY = 5
|
|
LOG_STREAM = os.environ.get("LOG_STREAM", "").lower() in ("1", "true", "yes")
|
|
|
|
SYSTEM_PROMPT = """You are a data processing assistant.
|
|
The input text is a list of tab-separated lines: brand, model, engine_code, date_range.
|
|
Convert to a JSON array of objects with fields: brand, model, engine_code, date_range.
|
|
The number of elements in JSON MUST match the number of lines in the list.
|
|
Keep empty fields as empty string "".
|
|
Output ONLY the JSON array. No explanations, no markdown, no other text before or after."""
|
|
|
|
|
|
def process_file(client: Client, filepath: pathlib.Path) -> bool:
|
|
"""Обрабатывает один файл. Возвращает True при успехе."""
|
|
content = filepath.read_text(encoding="utf-8")
|
|
out_path = RES_DIR / filepath.with_suffix(".json").name
|
|
|
|
for attempt in range(1, RETRY_COUNT + 1):
|
|
try:
|
|
t0 = time.perf_counter()
|
|
print(f" {filepath.name} — ", end="", flush=True)
|
|
|
|
result = ""
|
|
stream = client.chat(
|
|
model=MODEL,
|
|
messages=[
|
|
{"role": "system", "content": SYSTEM_PROMPT},
|
|
{"role": "user", "content": f"Convert to JSON. Output only the JSON array:\n\n{content}"},
|
|
],
|
|
stream=True,
|
|
)
|
|
for chunk in stream:
|
|
part = chunk.get("message", {}).get("content", "") or ""
|
|
result += part
|
|
if part:
|
|
if LOG_STREAM:
|
|
sys.stdout.write(part)
|
|
else:
|
|
sys.stdout.write(".")
|
|
sys.stdout.flush()
|
|
|
|
elapsed = time.perf_counter() - t0
|
|
try:
|
|
data = json.loads(result)
|
|
out_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
except json.JSONDecodeError:
|
|
out_path.write_text(result, encoding="utf-8")
|
|
|
|
print(f"\n" if LOG_STREAM else " ", end="")
|
|
print(f" {len(result)} chars, {elapsed:.1f}s — OK")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f" {filepath.name} — попытка {attempt}/{RETRY_COUNT} — ошибка: {e}")
|
|
if attempt < RETRY_COUNT:
|
|
time.sleep(RETRY_DELAY)
|
|
else:
|
|
print(f" {filepath.name} — пропущен после {RETRY_COUNT} попыток")
|
|
return False
|
|
|
|
return False
|
|
|
|
|
|
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:
|
|
out_path = RES_DIR / filepath.with_suffix(".json").name
|
|
if out_path.exists():
|
|
print(f" {filepath.name} — уже в res/, перенос в done/")
|
|
shutil.move(str(filepath), DONE_DIR / filepath.name)
|
|
continue
|
|
if process_file(client, filepath):
|
|
shutil.move(str(filepath), DONE_DIR / filepath.name)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|