feat: new trained model + video gen

.gitignore

@@ -27,6 +27,7 @@ bert_trained/
 checkpoint-*/
 
 # Data
+data/*
 data/*.csv
 data/*.srt
 data/*.txt

data_loader/dir_loader.py

@@ -0,0 +1,21 @@
+from pathlib import Path
+import pandas as pd
+from subtitles.parser import parse_srt_to_df   # your existing parser
+
+
+def load_subtitles_from_dir(folder: str | Path, ext: str = ".srt") -> pd.DataFrame:
+    """
+    Читает каждый файл из заданной папки и возвращает обьединенный pd.DataFrame
+    """
+    folder = Path(folder)
+    all_frames: list[pd.DataFrame] = []
+
+    for file in folder.glob(f"*{ext}"):
+        df_file = parse_srt_to_df(file)
+        df_file["source_file"] = file.name
+        all_frames.append(df_file)
+
+    if not all_frames:
+        raise FileNotFoundError(f"No {ext} files found in {folder.resolve()}")
+
+    return pd.concat(all_frames, ignore_index=True)

keywords.py

@@ -0,0 +1,66 @@
+KEYWORDS: set[str] = {
+    # ── Угроза / конфликт ─────────────────────────
+    "убью", "убить", "убил", "убийство", "смерть", "грохнуть",
+    "взорвётся", "взорвать", "взрыв", "детонатор", "бомба",
+    "ловушка", "засада", "атака", "штурм", "напасть",
+    "останавливать", "надо бежать", "слишком поздно", "мы проиграли",
+    "конец света", "апокалипсис", "катастрофа",
+
+    # ── Ранение / критическое состояние ────────────
+    "ранен", "ранение", "умираю", "без сознания", "остановить кровотечение",
+    "пульс не прослушивается", "доктор срочно", "нужна помощь",
+    "реанимация", "скорая", "кровь повсюду",
+
+    # ── Секреты / скрытая информация ───────────────
+    "секрет", "засекречено", "никому не говори", "это между нами",
+    "улик нет", "улик достаточно", "доказательства", "у нас есть доказательства",
+    "двойной агент", "крота", "подслушивают", "прослушка",
+    "флешка", "пароль", "шифр", "код доступа",
+
+    # ── Планы / приказ / миссия ─────────────────────
+    "план", "план б", "новая миссия", "операция начинается",
+    "отмена операции", "действуем по плану", "приказ ясен",
+    "координаты", "время ч", "расклад таков", "приступаем",
+
+    # ── Решения / выбор / судьба ────────────────────
+    "выбирай", "нет пути назад", "жертвовать", "это мой выбор",
+    "ты со мной", "придётся рискнуть", "спасти", "пожертвовать собой",
+    "мы должны решить", "остановись пока не поздно",
+
+    # ── Сюжетные твисты / раскрытия ────────────────
+    "на самом деле", "вся правда", "это была ложь", "он жив",
+    "подмена личности", "я всё время знала", "мой отец", "его брат близнец",
+    "это был мой план", "теперь всё ясно",
+
+    # ── Отношения / эмоции ─────────────────────────
+    "я тебя люблю", "выходи за меня", "у нас будет ребёнок",
+    "мы расстаёмся", "ты мне нравишься", "я без тебя не могу",
+    "он изменяет", "прости меня", "мы начнём сначала",
+
+    # ── Юмор / неловкость (часто ключевые моменты) ─
+    "это было случайно", "я всё испортил", "не в ту дверь",
+    "костюм порвался", "конфуз", "не то письмо", "он же шутил",
+
+    # ── Sci-Fi / техно-угроза ──────────────────────
+    "портал", "временной скачок", "искусственный интеллект", "ядро реактора",
+    "система отказала", "энергощит упал", "обратный отсчёт",
+    "спутник выведен из строя", "корабль падает",
+
+    # ── Фэнтези / магия ────────────────────────────
+    "заклинание", "пророчество", "артефакт", "древний орден",
+    "пламя дракона", "тёмный лорд", "портал закрывается",
+    "круг ритуала", "великий маг",
+
+    # ── Правосудие / расследование ────────────────
+    "вы арестованы", "расследование", "ордер на обыск",
+    "признательные показания", "дело раскрыто", "дело закрыто",
+    "улик достаточно", "главный подозреваемый",
+
+    # ── Финансы / шантаж / выкуп ───────────────────
+    "выкуп", "заложники", "долг", "банковский счёт",
+    "переведи деньги", "вымогательство", "шантаж", "биткойны",
+
+    # ── Общие триггеры (англ. формы, если субтитры смешанные) ─
+    "kill", "bomb", "explode", "secret", "mission", "run!",
+    "we're out of time", "trust me", "it's a trap"
+}

main.py

@@ -2,20 +2,36 @@ import argparse
 
 from subtitles.parser import parse_srt_to_df
 from scoring.bert_ranker import BERTImportanceRanker
+from video_editor.highlight import VideoHighlighter
 import config
 
 
 def main():
     parser = argparse.ArgumentParser(
         description="Оценка важности фраз в субтитрах с помощью BERT.")
-    parser.add_argument("srt_file", help="Путь к .srt файлу субтитров")
+    parser.add_argument("--srt_path", help="Путь к .srt файлу субтитров")
+    parser.add_argument("--video_path", help="Путь к видео файлу")
     parser.add_argument(
         "--output", help="Путь к выходному CSV", default=config.CSV_OUTPUT)
     args = parser.parse_args()
 
-    df = parse_srt_to_df(args.srt_file)
+    df = parse_srt_to_df(args.srt_path)
     ranker = BERTImportanceRanker(config.MODEL_NAME)
     df_scored = ranker.apply_to_dataframe(df)
+    important_df = (
+        df[df["importance_score"] > 0.9]
+        .loc[:, ["start", "end"]]
+        .reset_index(drop=True)
+    )
+    vh = VideoHighlighter(
+        video=args.video_path,
+        segments_df=important_df,
+        pad=1.5,
+        join_gap=0.3,
+        out_dir="clips",
+        concat=True,
+    )
+    vh.cut()
     df_scored.to_csv(args.output, index=False, encoding='utf-8')
     print(f"✅ Сохранено в {args.output}")
 

requirements.txt

@@ -7,15 +7,19 @@ certifi==2025.4.26
 chardet==5.2.0
 charset-normalizer==3.4.2
 datasets==3.6.0
+decorator==5.2.1
 dill==0.3.8
 filelock==3.18.0
 frozenlist==1.6.0
 fsspec==2025.3.0
 huggingface-hub==0.31.2
 idna==3.10
+imageio==2.37.0
+imageio-ffmpeg==0.6.0
 Jinja2==3.1.6
 joblib==1.5.0
 MarkupSafe==3.0.2
+moviepy==2.1.2
 mpmath==1.3.0
 multidict==6.4.3
 multiprocess==0.70.16
@@ -37,11 +41,14 @@ nvidia-nvjitlink-cu12==12.6.85
 nvidia-nvtx-cu12==12.6.77
 packaging==25.0
 pandas==2.2.3
+pillow==10.4.0
+proglog==0.1.12
 propcache==0.3.1
 psutil==7.0.0
 pyarrow==20.0.0
 pysrt==1.1.2
 python-dateutil==2.9.0.post0
+python-dotenv==1.1.0
 pytz==2025.2
 PyYAML==6.0.2
 regex==2024.11.6

subtitles/parser.py

@@ -1,9 +1,14 @@
 import pysrt
 import pandas as pd
+import chardet
+from pathlib import Path
 
 
-def parse_srt_to_df(filepath: str, encoding: str = "cp1251") -> pd.DataFrame:
-    subs = pysrt.open(filepath, encoding=encoding)
+def parse_srt_to_df(filepath: str) -> pd.DataFrame:
+    print(f"Загружаю {filepath}")
+    raw = Path(filepath).read_bytes()[:4096]
+    enc = chardet.detect(raw)["encoding"]
+    subs = pysrt.open(filepath, encoding=enc)
     rows = []
 
     for sub in subs:

train.py

@@ -4,10 +4,11 @@ from transformers.models.auto.modeling_auto import AutoModelForSequenceClassific
 from transformers.models.bert import BertTokenizer, BertForSequenceClassification
 from transformers.trainer import Trainer
 from transformers.training_args import TrainingArguments
-from subtitles.parser import parse_srt_to_df
 import torch
 from sklearn.model_selection import train_test_split
 import torch.nn.functional as F
+from keywords import KEYWORDS
+from data_loader.dir_loader import load_subtitles_from_dir
 
 sentiment_model_name = "cointegrated/rubert-tiny-sentiment-balanced"
 sentiment_tokenizer = AutoTokenizer.from_pretrained(sentiment_model_name)
@@ -16,9 +17,11 @@ sentiment_model = AutoModelForSequenceClassification.from_pretrained(
 
 
 def is_important(text):
+    low = text.lower()
     try:
         inputs = sentiment_tokenizer(
-            text, return_tensors="pt", truncation=True, padding=True)
+            low, return_tensors="pt", truncation=True, padding=True
+        )
         with torch.no_grad():
             logits = sentiment_model(**inputs).logits
             probs = F.softmax(logits, dim=1)
@@ -26,16 +29,19 @@ def is_important(text):
             labels = ["NEGATIVE", "NEUTRAL", "POSITIVE"]
             label = labels[label_id]
     except Exception:
+        print("Ошибка расчета сентимента")
+        return 0
+    if len(low.split()) < 2:
         return 0
     if label in ("NEGATIVE", "POSITIVE"):
         return 1
-    if len(text.split()) > 6:
+    if any(kw in low for kw in KEYWORDS):
         return 1
     return 0
 
 
 def main():
-    df = parse_srt_to_df("Breaking_Bad_RUS_2008_20210402033853.srt")
+    df = load_subtitles_from_dir("./data/subtitles/train")
     df["label"] = df["text"].astype(str).apply(is_important)
 
     train_texts, test_texts = train_test_split(
@@ -45,7 +51,11 @@ def main():
     tokenizer = BertTokenizer.from_pretrained(model_name)
 
     def tokenize_function(example):
-        return tokenizer(example["text"], padding="max_length", truncation=True, max_length=64)
+        return tokenizer(
+            example["text"],
+            padding="max_length",
+            truncation=True, max_length=64
+        )
 
     train_dataset = Dataset.from_pandas(train_texts[["text", "label"]])
     test_dataset = Dataset.from_pandas(test_texts[["text", "label"]])

video_editor/highlight.py

@@ -0,0 +1,142 @@
+# highlight_cutter.py
+"""
+Класс VideoHighlighter
+======================
+
+Нарезает «важные» отрезки видео по DataFrame со столбцами
+`start` и `end` (секунды).  Умеет:
+
+• объединять соседние / перекрывающиеся интервалы (join_gap)
+• добавлять «окно» pad слева и справа
+• сохранять клипы отдельно или склеивать в единый highlights.mp4
+"""
+
+from __future__ import annotations
+from pathlib import Path
+from typing import List, Tuple, Union
+import datetime as dt
+
+import pandas as pd
+from moviepy.video.io.VideoFileClip import VideoFileClip
+from moviepy.video.compositing.CompositeVideoClip import concatenate_videoclips
+
+
+class VideoHighlighter:
+    """Высвечивает (вырезает) нужные фрагменты ролика."""
+
+    # ─────────────────────────── init ────────────────────────────
+    def __init__(
+        self,
+        video: Union[str, Path],
+        segments_df: pd.DataFrame,
+        pad: float = 1.0,
+        join_gap: float = 0.2,
+        out_dir: Union[str, Path] = "clips",
+        concat: bool = True,
+    ) -> None:
+        self.video = Path(video)
+        self.df = segments_df.copy()
+        self.pad = float(pad)
+        self.join_gap = float(join_gap)
+        self.out_dir = Path(out_dir)
+        self.concat = concat
+
+        if not {"start", "end"} <= set(self.df.columns):
+            raise ValueError(
+                "DataFrame должен содержать колонки 'start' и 'end'")
+
+        self.out_dir.mkdir(parents=True, exist_ok=True)
+
+    # ────────────────────────── public API ─────────────────────────
+    def cut(self) -> None:
+        intervals = self._prepare_intervals()
+
+        with VideoFileClip(str(self.video)) as video:
+            clips = self._make_subclips(video, intervals)
+            if self.concat:
+                self._save_concat(clips)
+
+    # ────────────────────────── helpers ───────────────────────────
+    def _to_seconds(self, x) -> float:
+        """Любое представление времени → float секунд."""
+        if isinstance(x, (int, float)):
+            return float(x)
+        if isinstance(x, dt.time):
+            return x.hour * 3600 + x.minute * 60 + x.second + x.microsecond / 1_000_000
+        if isinstance(x, str):
+            # принимает форматы HH:MM:SS, HH:MM:SS.mmm или просто число
+            parts = x.replace(",", ".").split(":")
+            if len(parts) == 1:
+                return float(parts[0])
+            h, m, s = map(float, parts) if len(
+                parts) == 3 else (0, *map(float, parts))
+            return h * 3600 + m * 60 + s
+        raise TypeError(f"Неподдерживаемый тип времени: {type(x)}")
+
+    def _prepare_intervals(self) -> List[Tuple[float, float]]:
+        df = self.df.copy()
+        df["start"] = df["start"].apply(self._to_seconds)
+        df["end"] = df["end"].apply(self._to_seconds)
+
+        raw = list(df.sort_values("start")[
+                   ["start", "end"]].itertuples(index=False, name=None))
+        merged = self._merge_intervals(raw)
+        return self._add_padding(merged, total_dur=self._get_duration())
+
+    def _merge_intervals(self, intervals: List[Tuple[float, float]]) -> List[Tuple[float, float]]:
+        if not intervals:
+            return []
+        intervals.sort()
+        out = [list(intervals[0])]
+        for s, e in intervals[1:]:
+            last_e = out[-1][1]
+            if s - last_e <= self.join_gap:          # перекрытие / стык
+                out[-1][1] = max(last_e, e)
+            else:
+                out.append([s, e])
+        return [tuple(x) for x in out]
+
+    def _add_padding(self, intervals: List[Tuple[float, float]], total_dur: float) -> List[Tuple[float, float]]:
+        return [(max(0.0, s - self.pad), min(total_dur, e + self.pad)) for s, e in intervals]
+
+    def _get_duration(self) -> float:
+        with VideoFileClip(str(self.video)) as v:
+            return v.duration
+
+    def _make_subclips(self, video: VideoFileClip, intervals: List[Tuple[float, float]]) -> List[VideoFileClip]:
+        clips: list[VideoFileClip] = []
+        for idx, (s, e) in enumerate(intervals, 1):
+            clip = video.subclipped(s, e)
+            if not self.concat:
+                fname = self.out_dir / f"clip_{idx:02d}.mp4"
+                self._save_clip(clip, fname)
+                clip.close()
+            clips.append(clip)
+        return clips
+
+    def _save_clip(self, clip: VideoFileClip, fname: Path) -> None:
+        print(f"Сохраняю {fname.name:>12}: {clip.duration:6.2f} сек")
+        clip.write_videofile(
+            fname.as_posix(),
+            codec="libx264",
+            audio_codec="aac",
+            temp_audiofile=str(fname.with_suffix('.m4a')),
+            remove_temp=True,
+            logger=None,
+        )
+
+    def _save_concat(self, clips: List[VideoFileClip]) -> None:
+        final = concatenate_videoclips(clips, method="compose")
+        outfile = self.out_dir / "highlights.mp4"
+        print(f"Сохраняю дайджест {outfile.name}: {final.duration:6.2f} сек")
+        final.write_videofile(
+            outfile.as_posix(),
+            codec="libx264",
+            audio_codec="aac",
+            temp_audiofile=str(self.out_dir / "temp_audio.m4a"),
+            remove_temp=True,
+            logger=None,
+        )
+        final.close()
+        for c in clips:
+            c.close()