-
Notifications
You must be signed in to change notification settings - Fork 0
/
split_dialogues.py
36 lines (31 loc) · 1.26 KB
/
split_dialogues.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
import os
import re
def split_dialogue(text):
# Remove anything inside parentheses and surrounding spaces
text = re.sub(r'\s*\([^)]*\)\s*', ' ', text)
# Replace multiple consecutive spaces with a single space
text = re.sub(r'\s+', ' ', text)
# Split the text into sentences
sentences = re.split(r'(?<=[.!?])\s+', text)
return sentences
def split_into_files(sentences, movie_name):
if not os.path.exists("cleaned_scripts"):
os.makedirs("cleaned_scripts")
count = 1
for i in range(0, len(sentences), 10):
filename = f"cleaned_scripts/{movie_name}_{count}.txt"
with open(filename, "w", encoding='utf-8') as file:
file.write("\n\n".join(sentences[i:i+10]))
count += 1
def clean_script(file_path):
with open(file_path, 'r', encoding='utf-8') as file:
text = file.read()
movie_name = os.path.splitext(os.path.basename(file_path))[0]
sentences = split_dialogue(text)
split_into_files(sentences, movie_name)
if __name__ == "__main__":
# Iterate through directory
for filename in os.listdir("movies_source"):
if filename.endswith(".txt"):
file_path = os.path.join("movies_source", filename)
clean_script(file_path)