Merge pull request #13 from project-kxkg/MergeFix

Merge fix

Former-commit-id: 9554b4cceaef7d5a173731b003afddcd23eae48b

SRT.py

@@ -2,6 +2,7 @@ from datetime import timedelta
 from csv import reader
 from datetime import datetime
 import re
+import openai
 
 class SRT_segment(object):
     def __init__(self, *args) -> None:
@@ -105,18 +106,64 @@ class SRT_script():
             segments.append(self.merge_segs(idx_list))
 
         self.segments = segments # need memory release?
+        
+    
 
     def set_translation(self, translate:str, id_range:tuple):
         start_seg_id = id_range[0]
         end_seg_id = id_range[1]
-
+        
+        def inner_func(input_str):
+            response = openai.ChatCompletion.create(
+                    model="gpt-3.5-turbo",
+                    messages = [
+                        {"role": "system", "content": "You are a helpful assistant that help calibrates English to Chinese subtitle translations in starcraft2."},
+                        {"role": "system", "content": "You are provided with a translated Chinese transcript; you must modify or split the Chinese sentence to match the meaning and the number of the English transcript exactly one by one. You must not merge ANY Chinese lines, you can only split them but the total Chinese lines MUST equals to number of English lines."},
+                        {"role": "system", "content": "There is no need for you to add any comments or notes, and do not modify the English transcript."},
+                        {"role": "user", "content": 'You are given the English transcript and line number, your task is to merge or split the Chinese to match the exact number of lines in English transcript, no more no less. For example, if there are more Chinese lines than English lines, merge some the Chinese lines to match the number of English lines. If Chinese lines is less than English lines, split some Chinese lines to match the english lines: "{}"'.format(input_str)}
+                    ],
+                    temperature=0.7
+                )
+            return response['choices'][0]['message']['content'].strip()
+        
         lines = translate.split('\n\n')
-        if len(lines) != (end_seg_id - start_seg_id + 1):
-            print(id_range)
-            for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
-                print(seg.source_text)
-            print(translate)
-
+        if len(lines) < (end_seg_id - start_seg_id + 1):
+            count = 0
+            while count<5 and len(lines) != (end_seg_id - start_seg_id + 1):
+                
+                count += 1
+                print("Solving Unmatched Lines|iteration {}".format(count))
+                input_str = "\n"
+                #initialize GPT input
+                for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
+                    input_str += 'Sentence %d: ' %(i+1)+ seg.source_text + '\n'
+                    #Append to prompt string
+                    #Adds sentence index let GPT keep track of sentence breaks
+                input_str += translate    
+                #append translate to prompt
+                
+                flag = True
+                while flag:
+                    flag = False
+                    try:
+                        translate = inner_func(input_str)
+                    except Exception as e:
+                        print("An error has occurred during solving unmatched lines:",e)
+                        print("Retrying...")
+                        flag = True
+                        
+                lines = translate.split('\n\n')
+            if len(lines) < (end_seg_id - start_seg_id + 1):
+                print("Failed Solving unmatched lines, Manually parse needed")
+            
+            print(lines)
+            #print(id_range)
+            #for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
+            #    print(seg.source_text)
+            #print(translate)
+            
+        
+        
         for i, seg in enumerate(self.segments[start_seg_id-1:end_seg_id]):
             # naive way to due with merge translation problem
             # TODO: need a smarter solution

pipeline.py

@@ -163,10 +163,7 @@ def script_split(script_in, chunk_size = 1000):
 
 script_arr, range_arr = script_split(script_input)
 
-# Translate and save
-for s, range in tqdm(zip(script_arr, range_arr)):
-    # using chatgpt model
-    print(f"now translating sentences {range}")
+def get_response(model_name):
     if model_name == "gpt-3.5-turbo":
         # print(s + "\n")
         response = openai.ChatCompletion.create(
@@ -180,7 +177,7 @@ for s, range in tqdm(zip(script_arr, range_arr)):
             temperature=0.15
         )
 
-        translate = response['choices'][0]['message']['content'].strip()
+        return response['choices'][0]['message']['content'].strip()
 
     if model_name == "text-davinci-003":
         prompt = f"Please help me translate this into Chinese:\n\n{s}\n\n"
@@ -194,8 +191,23 @@ for s, range in tqdm(zip(script_arr, range_arr)):
             frequency_penalty=0.0,
             presence_penalty=0.0
         )
-        translate = response['choices'][0]['text'].strip()
-
+        return response['choices'][0]['text'].strip()
+    pass
+        
+        
+# Translate and save
+for s, range in tqdm(zip(script_arr, range_arr)):
+    # using chatgpt model
+    print(f"now translating sentences {range}")
+    flag = True
+    while flag:
+        flag = False
+        try:
+            translate = get_response(model_name)
+        except Exception as e:
+            print("An error has occurred during translation:",e)
+            print("Retrying...")
+            flag = True
     srt.set_translation(translate, range)
 
 srt.check_len_and_split()