-
Notifications
You must be signed in to change notification settings - Fork 4
/
main.py
336 lines (288 loc) · 10.4 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
import os
import subprocess
import configparser
from faster_whisper import WhisperModel
import datetime
import srt
import shutil
import logging
from langchain.chat_models import ChatOpenAI
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains.summarize import load_summarize_chain
from langchain.prompts import PromptTemplate
def read_config(config_file):
"""
Reads the configuration file and returns the input and output folder paths, model size,
device, compute type, and generate_transcript option as a tuple.
"""
config = configparser.ConfigParser()
config.read(config_file)
input_folder = config.get("PATHS", "input_folder")
output_folder = config.get("PATHS", "output_folder")
model_size = config.get("MODEL", "model_size")
device = config.get("MODEL", "device")
compute_type = config.get("MODEL", "compute_type")
task = config.get("MODEL", "task")
generate_transcript = config.getboolean("OPTIONS", "generate_transcript")
generate_summary = config.getboolean("OPTIONS", "generate_summary")
summary_language = config.get("OPTIONS", "summary_language")
openai_api_key = config.get("OPENAI", "API_KEY")
openai_api_base = config.get("OPENAI", "API_BASE")
openai_model = config.get("OPENAI", "model")
text_chunk_size = config.getint("OPTIONS", "text_chunk_size")
max_chunk = config.getint("OPTIONS", "max_chunk")
return (
input_folder,
output_folder,
model_size,
device,
compute_type,
task,
generate_transcript,
generate_summary,
summary_language,
openai_api_key,
openai_api_base,
openai_model,
text_chunk_size,
max_chunk
)
def convert_to_mp3(video_file, audio_file):
"""
Converts the input video file to MP3 format using ffmpeg and saves the resulting audio file to
the specified output file path. Returns True if the conversion is successful, False otherwise.
"""
if os.path.exists(audio_file):
print(f"Skipping {video_file} - audio file already exists")
return True
# if file is mp3, just use shutil to copy it
if video_file.endswith(".mp3"):
print(f"Copying {video_file} to {audio_file}")
shutil.copy2(video_file, audio_file)
return True
cmd = [
"ffmpeg",
"-i",
video_file,
"-vn",
"-acodec",
"libmp3lame",
"-ab",
"192k",
"-ac",
"2",
"-loglevel",
"quiet",
audio_file,
]
print(f"Processing: {video_file}")
print(f"Output audio path: {audio_file}")
result = subprocess.run(cmd)
if result.returncode != 0:
print(f"Conversion of {video_file} failed")
return False
print(f"Conversion of {video_file} success")
return True
def generate_subtitles(model, task, audio_file, subtitle_file, transcript_file=None):
"""
Transcribes the specified audio file using the Faster-Whisper model, generates an SRT subtitle file
at the specified output file path, and optionally generates a transcript file. Returns True if the
subtitle generation is successful, False otherwise.
"""
if os.path.exists(subtitle_file):
print(f"Skipping {audio_file} - subtitles file already exists")
return True
segments, info = model.transcribe(audio_file, vad_filter=True, task=task)
print(
"Detected language '%s' with probability %f"
% (info.language, info.language_probability)
)
subtitles = []
texts = []
for i, segment in enumerate(segments):
start_time = datetime.timedelta(milliseconds=segment.start * 1000)
end_time = datetime.timedelta(milliseconds=segment.end * 1000)
text = segment.text.strip()
texts.append(text)
if text:
# Create a subtitle object for the segment
subtitle = srt.Subtitle(
index=i + 1, start=start_time, end=end_time, content=text
)
subtitles.append(subtitle)
# Write the subtitles to the SRT file
with open(subtitle_file, "w", encoding="utf-8") as f:
f.write(srt.compose(subtitles))
print(f"Generation of {subtitle_file} successful")
# Write the transcript file if requested
print("transcript_file " + transcript_file)
if transcript_file:
with open(transcript_file, "w", encoding="utf-8") as f:
for text in texts:
f.write(text + " ")
print(f"Generation of {transcript_file} successful")
return True
def summarize(
transcript_file,
summary_file,
summary_language,
openai_api_key,
openai_api_base,
openai_model,
text_chunk_size,
max_chunk
):
# Instantiate the LLM modelI
if os.path.exists(summary_file):
print(f"Skipping {summary_file} - summary file already exists")
return True
print("Strat summaring " + transcript_file)
llm = ChatOpenAI(
model=openai_model,
temperature=0.7,
openai_api_key=openai_api_key,
openai_api_base=openai_api_base,
)
with open(transcript_file, "r") as file:
txt = file.read()
# print(txt)
# Split text
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=text_chunk_size,
)
texts = text_splitter.split_text(txt)
# Create multiple documents
docs = [Document(page_content=t) for t in texts]
# Text summarization
# Write a concise summary of the following
# CONCISE SUMMARY IN {language}:"""
if len(docs) == 1:
stuff_prompt_template = """Please use Markdown syntax to help me summarize the key information and important content. Your response should summarize the main information and important content in the original text in a clear manner, using appropriate headings, markers, and formats to facilitate readability and understanding.Please note that your response should retain the relevant details in the original text while presenting them in a concise and clear manner. You can freely choose the content to highlight and use appropriate Markdown markers to emphasize it. Now summary following content in {language}:
{text}
"""
stuff_prompt = PromptTemplate(
template=stuff_prompt_template,
input_variables=[
"text",
],
partial_variables={"language": summary_language},
)
chain = load_summarize_chain(
llm,
chain_type="stuff",
prompt=stuff_prompt
)
response = chain.run(docs)
print(response)
with open(summary_file, "w", encoding="utf-8") as f:
f.write(response)
return True
if(len(docs) > max_chunk):
print('The doc is too long, you should use gpt4-4k or calude to summarize it')
return True
map_prompt_template = """Write a concise summary of the following:
{text}
SUMMARY IN {language}:"""
map_prompt = PromptTemplate(
template=map_prompt_template,
input_variables=[
"text",
],
partial_variables={"language": summary_language},
)
combine_prompt_template = """Write a concise summary of the following:
{text}
CONCISE SUMMARY IN {language}:"""
combine_prompt = PromptTemplate(
template=combine_prompt_template,
input_variables=[
"text",
],
partial_variables={"language": summary_language},
)
chain = load_summarize_chain(
llm,
chain_type="map_reduce",
return_intermediate_steps=True,
map_prompt=map_prompt,
combine_prompt=combine_prompt,
)
response = chain({"input_documents": docs}, return_only_outputs=True)
print(response)
with open(summary_file, "w", encoding="utf-8") as f:
f.write("Summary_text:\n")
f.write(response["output_text"])
f.write("\n\nSection Contents:\n\n")
for idx, section in enumerate(response["intermediate_steps"]):
f.write(f"{idx + 1}.{section}\n")
return True
if __name__ == "__main__":
config_file = "config.ini"
(
input_folder,
output_folder,
model_size,
device,
compute_type,
task,
generate_transcript,
generate_summary,
summary_language,
openai_api_key,
openai_api_base,
openai_model,
text_chunk_size,
max_chunk
) = read_config(config_file)
logging.basicConfig()
logging.getLogger("faster_whisper").setLevel(logging.DEBUG)
model = WhisperModel(model_size, device=device, compute_type=compute_type)
if not os.path.exists(output_folder):
os.makedirs(output_folder)
for file_name in os.listdir(input_folder):
# if file is video or audio
if file_name.endswith((".mp4", ".avi", ".mov", ".wmv", ".flv", ".mkv", ".mp3")):
video_file = os.path.join(input_folder, file_name)
audio_file = os.path.join(
output_folder, os.path.splitext(file_name)[0] + ".mp3"
)
subtitle_file = os.path.join(
output_folder, os.path.splitext(file_name)[0] + ".srt"
)
transcript_file = (
os.path.join(output_folder, os.path.splitext(file_name)[0] + ".txt")
if generate_transcript
else None
)
summary_file = (
os.path.join(
output_folder, os.path.splitext(file_name)[0] + ".md"
)
if generate_summary
else None
)
if convert_to_mp3(video_file, audio_file):
generate_subtitles(
model,
task,
audio_file,
subtitle_file,
transcript_file,
)
try:
if generate_summary:
summarize(
transcript_file,
summary_file,
summary_language,
openai_api_key,
openai_api_base,
openai_model,
text_chunk_size,
max_chunk
)
except Exception as e:
print(e)
continue
print("All conversions complete")