diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..d5f8487 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,12 @@ +.git +.venv +__pycache__/ +*.pyc +runtime/ +debug/ +images/ +input.* +audio.* +video.* +output.* +*.json diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..39db138 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,36 @@ +# syntax=docker/dockerfile:1 + +ARG GPU=nvidia + +FROM nvidia/cuda:13.0.2-cudnn-runtime-ubuntu24.04 AS base-nvidia +FROM rocm/dev-ubuntu-24.04:7.2.1-complete AS base-amd + +FROM base-${GPU} AS app + +ENV DEBIAN_FRONTEND=noninteractive \ + PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + XDG_CACHE_HOME=/tmp/sumka-cache + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + ca-certificates curl ffmpeg fonts-dejavu-core \ + python3 python3-pip python3-venv \ + && curl -fsSL -o /tmp/chrome.deb \ + https://dl.google.com/linux/direct/google-chrome-stable_current_amd64.deb \ + && apt-get install -y /tmp/chrome.deb \ + && rm -rf /var/lib/apt/lists/* /tmp/chrome.deb + +WORKDIR /app +COPY requirements.txt requirements-nvidia.txt requirements-amd.txt ./ +ARG GPU +RUN python3 -m venv /opt/venv \ + && /opt/venv/bin/pip install --no-cache-dir --upgrade pip \ + && /opt/venv/bin/pip install --no-cache-dir -r requirements-${GPU}.txt + +ENV PATH=/opt/venv/bin:$PATH + +COPY . ./ +WORKDIR /work + +ENTRYPOINT ["python3", "/app/main.py"] diff --git a/README.md b/README.md index 175307d..476fb8c 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,36 @@ ## Установка (Docker) -TODO +Образы собираются из одного Dockerfile: + +```bash +docker build --build-arg GPU=nvidia -t sumka:nvidia . +docker build --build-arg GPU=amd -t sumka:amd . +``` + +Для NVIDIA нужен настроенный NVIDIA Container Toolkit. Для AMD на хосте должен +быть установлен совместимый ROCm. Рабочая директория монтируется в `/work`: + +```bash +WORK_DIR=$(mktemp -d) +cp input.mkv "$WORK_DIR/" + +# nvidia +docker run --rm --gpus all -e AI_API_KEY \ + --mount type=volume,src=sumka-whisper-cache,dst=/root/.cache/whisper \ + --mount type=bind,src="$WORK_DIR",dst=/work \ + --tmpfs /tmp:rw,size=8g sumka:nvidia + +# amd +docker run --rm --device=/dev/kfd --device=/dev/dri --group-add video \ + --security-opt seccomp=unconfined -e AI_API_KEY \ + --mount type=volume,src=sumka-whisper-cache,dst=/root/.cache/whisper \ + --mount type=bind,src="$WORK_DIR",dst=/work \ + --tmpfs /tmp:rw,size=8g sumka:amd +``` + +Все входные, промежуточные и итоговые файлы находятся в `WORK_DIR`. Для дампов +запросов заранее создайте `"$WORK_DIR/debug"`; промпты остаются внутри образа. ## Установка (Python) @@ -28,8 +57,8 @@ cd 2026-linux-sumka python3 -m venv .venv # Активировать виртуальное окружение . .venv/bin/activate -# Установить зависимости -pip install -r requirements.txt +# Установить зависимости под свою видеокарту +pip install -r requirements-nvidia.txt # либо requirements-amd.txt # Также должны быть установлены ffmpeg и Chromium/Google Chrome diff --git a/asr_eventizer.py b/asr_eventizer.py index 861c40f..8728d2b 100644 --- a/asr_eventizer.py +++ b/asr_eventizer.py @@ -1,13 +1,12 @@ from typing import Any import traceback -import shutil -import os import json from pydantic import BaseModel from agent import Agent, AgentMessage from windowizer import Windowizer, Window from asr_filter import AsrFilterResult, AsrFilterSegment +from paths import debug_path, prompt_path from utils import Timeline, Event class _PreEvent(BaseModel): @@ -41,21 +40,19 @@ class AsrEventizer: role="user" ) ] - debug_dir = None - if os.path.isdir("debug"): - debug_dir = f"debug/AsrEventizer/{AsrEventizer.DEBUG_ID}" - AsrEventizer.DEBUG_ID += 1 - os.makedirs(debug_dir, exist_ok=True) + debug_dir = debug_path("AsrEventizer", AsrEventizer.DEBUG_ID) if debug_dir: - with open(f"{debug_dir}/request.txt", "w") as f: - f.write(messages[1].content) + AsrEventizer.DEBUG_ID += 1 + (debug_dir / "request.txt").write_text(messages[1].content, encoding="utf-8") retries_left = 5 while retries_left > 0: retries_left -= 1 response = self._agent.completion(messages=messages, response_format=_EventizeResult) if debug_dir: - with open(f"{debug_dir}/{retries_left}-retries-left.txt", "w") as f: - f.write(response) + (debug_dir / f"{retries_left}-retries-left.txt").write_text( + response, + encoding="utf-8", + ) # validate data try: response = json.loads(response) @@ -85,11 +82,10 @@ class AsrEventizer: """ self._agent = agent self._windowizer = windowizer - with open("prompts/asr_eventizer.md", "r") as f: - self._system_prompt = AgentMessage( - content=f.read(), - role="system" - ) + self._system_prompt = AgentMessage( + content=prompt_path("asr_eventizer.md").read_text(encoding="utf-8"), + role="system", + ) def eventize(self, asr_filter_result: AsrFilterResult) -> Timeline: """Builds event timeline from `AsrFilterResult`. Resulting timeline @@ -125,4 +121,4 @@ class AsrEventizer: )) output_id += 1 # final timeline - return Timeline(events=events) \ No newline at end of file + return Timeline(events=events) diff --git a/paths.py b/paths.py new file mode 100644 index 0000000..18e9e50 --- /dev/null +++ b/paths.py @@ -0,0 +1,24 @@ +"""Paths for bundled resources and per-run artifacts.""" + +from pathlib import Path + + +SOURCE_DIR = Path(__file__).resolve().parent + + +def prompt_path(name: str) -> Path: + """Return a prompt bundled with the application source code.""" + path = SOURCE_DIR / "prompts" / name + if not path.is_file(): + raise FileNotFoundError(f"Prompt not found: {path}") + return path + + +def debug_path(component: str, index: int) -> Path | None: + """Create a debug directory in the current working directory when enabled.""" + root = Path.cwd() / "debug" + if not root.is_dir(): + return None + path = root / component / str(index) + path.mkdir(parents=True, exist_ok=True) + return path diff --git a/pdf_builder.py b/pdf_builder.py index 5850989..0496078 100644 --- a/pdf_builder.py +++ b/pdf_builder.py @@ -3,6 +3,7 @@ from __future__ import annotations import html +import os import re import shutil import subprocess @@ -277,6 +278,8 @@ class PdfBuilder: f"--print-to-pdf={destination}", html_path.as_uri(), ] + if os.geteuid() == 0: + command.insert(1, "--no-sandbox") try: result = subprocess.run( command, diff --git a/reference_resolver.py b/reference_resolver.py index 339879a..5b042bb 100644 --- a/reference_resolver.py +++ b/reference_resolver.py @@ -13,6 +13,7 @@ from typing import Literal from pydantic import BaseModel, ConfigDict, ValidationError, model_validator from agent import Agent, AgentMessage +from paths import debug_path, prompt_path from utils import Event, Timeline from video_references import UnresolvedReference, UnresolvedReferences @@ -94,10 +95,14 @@ class ReferenceResolver: self._image_dir = Path(image_dir) self._frame_interval = frame_interval self._max_frames = max_frames - with open("prompts/reference_resolver.md", "r", encoding="utf-8") as f: - self._system_prompt = AgentMessage(content=f.read(), role="system") - with open("prompts/reference_resolver_ocr.md", "r", encoding="utf-8") as f: - self._ocr_prompt = AgentMessage(content=f.read(), role="system") + self._system_prompt = AgentMessage( + content=prompt_path("reference_resolver.md").read_text(encoding="utf-8"), + role="system", + ) + self._ocr_prompt = AgentMessage( + content=prompt_path("reference_resolver_ocr.md").read_text(encoding="utf-8"), + role="system", + ) def _sampling_timestamps(self, start: float, end: float) -> list[float]: start = max(0.0, start) @@ -217,13 +222,10 @@ class ReferenceResolver: ), ] - debug_dir = None - if os.path.isdir("debug"): - debug_dir = f"debug/ReferenceResolver/{ReferenceResolver.DEBUG_ID}" + debug_dir = debug_path("ReferenceResolver", ReferenceResolver.DEBUG_ID) + if debug_dir: ReferenceResolver.DEBUG_ID += 1 - os.makedirs(debug_dir, exist_ok=True) - with open(f"{debug_dir}/request.txt", "w", encoding="utf-8") as f: - f.write(messages[1].content) + (debug_dir / "request.txt").write_text(messages[1].content, encoding="utf-8") retries_left = self.MAX_RETRIES while retries_left > 0: @@ -235,12 +237,10 @@ class ReferenceResolver: response_format=_FrameResult, ) if debug_dir: - with open( - f"{debug_dir}/{retries_left}-retries-left.txt", - "w", + (debug_dir / f"{retries_left}-retries-left.txt").write_text( + response, encoding="utf-8", - ) as f: - f.write(response) + ) try: result = _FrameResult.model_validate_json(response) @@ -312,14 +312,11 @@ class ReferenceResolver: ), ] - debug_dir = None - if os.path.isdir("debug"): - debug_dir = f"debug/ReferenceResolverOCR/{ReferenceResolver.OCR_DEBUG_ID}" + debug_dir = debug_path("ReferenceResolverOCR", ReferenceResolver.OCR_DEBUG_ID) + if debug_dir: ReferenceResolver.OCR_DEBUG_ID += 1 - os.makedirs(debug_dir, exist_ok=True) - with open(f"{debug_dir}/request.txt", "w", encoding="utf-8") as f: - f.write(messages[1].content) - shutil.copyfile(crop_path, f"{debug_dir}/crop.jpg") + (debug_dir / "request.txt").write_text(messages[1].content, encoding="utf-8") + shutil.copyfile(crop_path, debug_dir / "crop.jpg") retries_left = self.MAX_RETRIES while retries_left > 0: @@ -331,12 +328,10 @@ class ReferenceResolver: response_format=_OCRResult, ) if debug_dir: - with open( - f"{debug_dir}/{retries_left}-retries-left.txt", - "w", + (debug_dir / f"{retries_left}-retries-left.txt").write_text( + response, encoding="utf-8", - ) as f: - f.write(response) + ) try: result = _OCRResult.model_validate_json(response) diff --git a/requirements-amd.txt b/requirements-amd.txt new file mode 100644 index 0000000..01fd31b --- /dev/null +++ b/requirements-amd.txt @@ -0,0 +1,3 @@ +-r requirements.txt +--extra-index-url https://download.pytorch.org/whl/rocm7.2 +torch==2.14.0+rocm7.2 diff --git a/requirements-nvidia.txt b/requirements-nvidia.txt new file mode 100644 index 0000000..3a851ff --- /dev/null +++ b/requirements-nvidia.txt @@ -0,0 +1,3 @@ +-r requirements.txt +--extra-index-url https://download.pytorch.org/whl/cu130 +torch==2.14.0+cu130 diff --git a/requirements.txt b/requirements.txt index 7fa5704..ebe1490 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,5 @@ -httpx2 -latex2mathml -openai -openai-whisper -pydantic>=2 -torch +httpx2==2.13.0 +latex2mathml==3.81.1 +openai==3.14.0 +openai-whisper==20250625 +pydantic==2.13.5 diff --git a/structure_builder.py b/structure_builder.py index 997526e..d0b68ff 100644 --- a/structure_builder.py +++ b/structure_builder.py @@ -1,11 +1,11 @@ import json -import os import traceback from typing import Annotated, Literal from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator from agent import Agent, AgentMessage +from paths import debug_path, prompt_path from utils import Event, Timeline from windowizer import Window, Windowizer @@ -146,8 +146,10 @@ class StructureBuilder: """ self._agent = agent self._windowizer = windowizer - with open("prompts/structure_builder.md", "r", encoding="utf-8") as f: - self._system_prompt = AgentMessage(content=f.read(), role="system") + self._system_prompt = AgentMessage( + content=prompt_path("structure_builder.md").read_text(encoding="utf-8"), + role="system", + ) @staticmethod def _validate_result(result: _BuildResult, window: Window[Event]) -> None: @@ -192,25 +194,20 @@ class StructureBuilder: ), ] - debug_dir = None - if os.path.isdir("debug"): - debug_dir = f"debug/StructureBuilder/{StructureBuilder.DEBUG_ID}" + debug_dir = debug_path("StructureBuilder", StructureBuilder.DEBUG_ID) + if debug_dir: StructureBuilder.DEBUG_ID += 1 - os.makedirs(debug_dir, exist_ok=True) - with open(f"{debug_dir}/request.txt", "w", encoding="utf-8") as f: - f.write(messages[1].content) + (debug_dir / "request.txt").write_text(messages[1].content, encoding="utf-8") retries_left = self.MAX_RETRIES while retries_left > 0: retries_left -= 1 response = self._agent.completion(messages=messages, response_format=_BuildResult) if debug_dir: - with open( - f"{debug_dir}/{retries_left}-retries-left.txt", - "w", + (debug_dir / f"{retries_left}-retries-left.txt").write_text( + response, encoding="utf-8", - ) as f: - f.write(response) + ) try: result = _BuildResult.model_validate_json(response) diff --git a/structure_refiner.py b/structure_refiner.py index 77ee476..d57d3bf 100644 --- a/structure_refiner.py +++ b/structure_refiner.py @@ -1,10 +1,10 @@ import json -import os import traceback from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator from agent import Agent, AgentMessage +from paths import debug_path, prompt_path from structure_builder import ImageElement, Structure, StructureElement from windowizer import Window, Windowizer @@ -39,8 +39,10 @@ class StructureRefiner: def __init__(self, agent: Agent, windowizer: Windowizer[StructureElement]) -> None: self._agent = agent self._windowizer = windowizer - with open("prompts/structure_refiner.md", "r", encoding="utf-8") as f: - self._system_prompt = AgentMessage(content=f.read(), role="system") + self._system_prompt = AgentMessage( + content=prompt_path("structure_refiner.md").read_text(encoding="utf-8"), + role="system", + ) @staticmethod def _validate_result( @@ -69,13 +71,10 @@ class StructureRefiner: ), ] - debug_dir = None - if os.path.isdir("debug"): - debug_dir = f"debug/StructureRefiner/{StructureRefiner.DEBUG_ID}" + debug_dir = debug_path("StructureRefiner", StructureRefiner.DEBUG_ID) + if debug_dir: StructureRefiner.DEBUG_ID += 1 - os.makedirs(debug_dir, exist_ok=True) - with open(f"{debug_dir}/request.txt", "w", encoding="utf-8") as f: - f.write(messages[1].content) + (debug_dir / "request.txt").write_text(messages[1].content, encoding="utf-8") retries_left = self.MAX_RETRIES while retries_left > 0: @@ -85,12 +84,10 @@ class StructureRefiner: response_format=_RefineResult, ) if debug_dir: - with open( - f"{debug_dir}/{retries_left}-retries-left.txt", - "w", + (debug_dir / f"{retries_left}-retries-left.txt").write_text( + response, encoding="utf-8", - ) as f: - f.write(response) + ) try: result = _RefineResult.model_validate_json(response) diff --git a/video_references.py b/video_references.py index 8b580e7..38df2fd 100644 --- a/video_references.py +++ b/video_references.py @@ -1,11 +1,11 @@ import json -import os import traceback from typing import Literal from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator from agent import Agent, AgentMessage +from paths import debug_path, prompt_path from utils import Timeline from windowizer import Window, Windowizer @@ -71,8 +71,10 @@ class VideoReferenceBuilder: ) -> None: self._agent = agent self._windowizer = windowizer - with open("prompts/video_references.md", "r", encoding="utf-8") as f: - self._system_prompt = AgentMessage(content=f.read(), role="system") + self._system_prompt = AgentMessage( + content=prompt_path("video_references.md").read_text(encoding="utf-8"), + role="system", + ) @staticmethod def _validate_result( @@ -113,13 +115,10 @@ class VideoReferenceBuilder: ), ] - debug_dir = None - if os.path.isdir("debug"): - debug_dir = f"debug/VideoReferenceBuilder/{self.DEBUG_ID}" + debug_dir = debug_path("VideoReferenceBuilder", self.DEBUG_ID) + if debug_dir: VideoReferenceBuilder.DEBUG_ID += 1 - os.makedirs(debug_dir, exist_ok=True) - with open(f"{debug_dir}/request.txt", "w", encoding="utf-8") as f: - f.write(messages[1].content) + (debug_dir / "request.txt").write_text(messages[1].content, encoding="utf-8") retries_left = self.MAX_RETRIES while retries_left > 0: @@ -129,12 +128,10 @@ class VideoReferenceBuilder: response_format=_WindowResult, ) if debug_dir: - with open( - f"{debug_dir}/{retries_left}-retries-left.txt", - "w", + (debug_dir / f"{retries_left}-retries-left.txt").write_text( + response, encoding="utf-8", - ) as f: - f.write(response) + ) try: result = _WindowResult.model_validate_json(response)