recap_gen/main.py

25 lines
825 B
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import argparse
from subtitles.parser import parse_srt_to_df
from scoring.bert_ranker import BERTImportanceRanker
import config
def main():
parser = argparse.ArgumentParser(
description="Оценка важности фраз в субтитрах с помощью BERT.")
parser.add_argument("srt_file", help="Путь к .srt файлу субтитров")
parser.add_argument(
"--output", help="Путь к выходному CSV", default=config.CSV_OUTPUT)
args = parser.parse_args()
df = parse_srt_to_df(args.srt_file)
ranker = BERTImportanceRanker(config.MODEL_NAME)
df_scored = ranker.apply_to_dataframe(df)
df_scored.to_csv(args.output, index=False, encoding='utf-8')
print(f"✅ Сохранено в {args.output}")
if __name__ == "__main__":
main()