Prompt Selection

SRT-gui.py

@@ -2,6 +2,7 @@
 
 import subtitlecorrector as stc
 import tkinter as tk
+from tkinter import *
 from tkinter import filedialog
 import platform
 
@@ -10,7 +11,8 @@ class SRTCorrectorGui:
     def process_subtitle(self):
         if (self.full_output_path != "" and self.subtitle_file_path != ""):
             self.sfile_status_label.config(text="Processing started...")
-            stc.correct_subtitles(self.subtitle_file_path, self.full_output_path)
+            self.selected_prompt.set(self.prompts[self.selected_prompt.get()])
+            stc.correct_subtitles(self.subtitle_file_path, self.full_output_path, self.selected_prompt.get())
             self.sfile_status_label.config(text="Processing completed.")
         else:
             self.sfile_status_label.config(text="Please select the path to the input subtitle and the path for the output.")
@@ -38,7 +40,6 @@ def choose_odir(self):
             self.ofile_path.config(state="readonly")
             open(self.full_output_path, "w+")
 
-
     # Sets up the GUI for selecting a file.
     def file_selection_GUI(self):
         self.file_frame = tk.Frame(master=self.root_window, width=100, height=100)
@@ -87,22 +88,34 @@ def processing_GUI(self):
 
     # Constructor
     def __init__(self):
+        # Variables:
         self.subtitle_file_path = ""
         self.full_output_path = ""
         self.output_dir = ""
-        # Window initialisation.
+        self.prompts = {
+            'Prompt: Subtitle Correction': 1,
+            'Prompt: Dutch to English' : 2,
+            'Prompt: English to Dutch' : 3
+        }
+
+        # Window initialisation:
         self.root_window = tk.Tk()
         self.root_window.title = "Subtitle Corrector GUI"
         self.root_window.geometry("425x150")
 
+        # Prompt Selection:
+        self.selected_prompt = StringVar(self.root_window)
+        self.selected_prompt.set(self.prompts['Prompt: Subtitle Correction'])
+        self.prompt_select = OptionMenu(self.root_window, self.selected_prompt, *self.prompts)
+        self.prompt_select.pack()
+
+        # General setup:
         self.file_selection_GUI()
         self.ofile_selection_GUI()
         self.processing_GUI()
         self.license_label = tk.Label(self.root_window, text="Subtitle Corrector is Free Software licensed under the GNU GPLv3")
         self.license_label.pack(side="bottom")
 
-
-
 def main():
     SRT_corrector = SRTCorrectorGui()
     SRT_corrector.root_window.mainloop()

SRT-tui.py

@@ -5,6 +5,12 @@
 import os
 import platform
 
+prompt_shorthand_names = {
+    'SCR': 1,
+    'DTE': 2,
+    'ETD': 3
+}
+
 # Creates the argument parser.
 # These python libraries are actually pretty cool.
 def arg_parser_init():
@@ -14,12 +20,16 @@ def arg_parser_init():
         epilog="subtitle-corrector")
     parser.add_argument('input_srt')
     parser.add_argument('--ofile', action='store', required=True)
+    parser.add_argument('--prompt', action='store')
     return (parser)
 
 # Main.
 def main():
     parser = arg_parser_init()
     args = parser.parse_args()
-    stc.correct_subtitles(args.input_srt, args.ofile)
+    if (args.prompt != "" and (args.prompt != 'SCR' or args.prompt != 'DTE' or args.prompt != 'ETD')):
+        stc.correct_subtitles(args.input_srt, args.ofile, prompt_shorthand_names[args.prompt])
+    else:
+        stc.correct_subtitles(args.input_srt, args.ofile, prompt_shorthand_names['SCR'])
 
 main()

subtitlecorrector.py

@@ -32,7 +32,8 @@
 import platform
 import tiktoken
 
-subtitle_translation_prompt = '''
+
+dutch_to_english_prompt = '''
 Je gaat optreden als een programma ontworpen om ondertitels naar het Engels te vertalen.
 Je zult automatisch gegenereerde ondertitels vertalen.
 Je krijgt een invoer in het .srt-formaat.
@@ -41,6 +42,17 @@
 Zorg ervoor dat je de ondertitel-ID behoudt
 '''
 
+english_to_dutch_prompt = '''
+You are going to act as a program designed to help translate subtitles.
+You will be translating automatically generated subtitles from Dutch to English.
+You will be given an input in the .srt format.
+You will be doing the following: Translating Dutch sentences into English. Correcting out of place words. Removing redundant and or filler words.
+Keep the content of the sentences consistent with the input.
+The number of lines in the output must be the same as the number of lines in the input.
+Make sure to preserve the subtitle id.
+Please do not use overly formal language.
+'''
+
 subtitle_correction_prompt = '''
 You are going to act as a program designed to help correct subtitles.
 You will be correcting automatically generated subtitles from a talk at a programming school.
@@ -52,6 +64,12 @@
 Please do not use overly formal language.
 '''
 
+prompt_list = {
+    1: subtitle_correction_prompt,
+    2: dutch_to_english_prompt,
+    3: english_to_dutch_prompt
+}
+
 def srt_to_text(file_name):
     print("Formatting: " + file_name)
     srt_file = open(file_name, encoding="utf-8")
@@ -65,12 +83,12 @@ def srt_to_text(file_name):
     return (subtitles_list)
 
 # Queries ChatGPT with the stripped SRT data.
-def query_chatgpt(query_str):
+def query_chatgpt(query_str, chosen_prompt):
     openai.api_key = os.environ.get('OPENAI_API_KEY')
     client = openai.ChatCompletion()
     chat_log = [{
         'role': 'system',
-        'content': subtitle_correction_prompt,
+        'content': prompt_list[int(chosen_prompt)],
     }]
     chat_log.append({'role': 'user', 'content': query_str})
     response = client.create(model='gpt-4', messages=chat_log)
@@ -83,7 +101,7 @@ def num_tokens(raw_text):
     num_tokens = len(encoding.encode(raw_text))
     return num_tokens
 
-def query_loop(subtitle_file):
+def query_loop(subtitle_file, chosen_prompt):
     full_output = ""
     query_str = ""
     query_counter = 0
@@ -129,7 +147,8 @@ def query_loop(subtitle_file):
     return (srt.compose(slist))
 
 # Reads the raw SRT data and passes it to ChatGPT to be processed.
-def correct_subtitles(subtitle_file, outputfile="output.srt"):
-    full_output = query_loop(subtitle_file)
-    ofile = open(outputfile, "w", encoding="utf-8")
-    ofile.write(full_output)
+def correct_subtitles(subtitle_file, outputfile="output.srt", chosen_prompt=1):
+    print(prompt_list[int(chosen_prompt)])
+    #full_output = query_loop(subtitle_file, chosen_prompt)
+    #ofile = open(outputfile, "w", encoding="utf-8")
+    #ofile.write(full_output)