feat: Batch working

main.py

@@ -1,4 +1,5 @@
 import os
+import re
 from googletrans import Translator
 from pathlib import Path
 import json
@@ -10,35 +11,94 @@ load_dotenv()
 openai.organization = os.getenv('org')
 openai.api_key = os.getenv('key')
 
-# Load JSON file into a Python dictionary
-with open('Map013.json', 'r', encoding='utf-8') as f:
-    data = json.load(f)
+def main():
+    # Load JSON file into a Python dictionary
+    with open('Map013.json', 'r', encoding='utf-8') as f:
+        data = json.load(f)
+        parse401(data)
 
-# Extract 401 events text
-events = data['events']
-text = []
-for event in events:
-    if event is not None:
-        for page in event['pages']:
-            for command in page['list']:
-                if command['code'] == 401:
-                    text.append(command['parameters'][0])
+def parse401(data):
+    """
+    Extracts the first parameter of all commands with code 401 from a dictionary of events.
 
-# Split into batches
-pipe = '/p'
-script = pipe.join(text)
+    Args:
+        data (list): A list of events.
 
-# Translate text using GPT
-def translate(t):
+    Returns:
+        list: A list of strings containing the first parameter of all commands with code 401.
+    """
+    # Extract 401 Text
+    events = data['events']
+    untranslatedTextList = []
+    for event in events:
+        if event is not None:
+            for page in event['pages']:
+                for command in page['list']:
+                    if command['code'] == 401:
+                        untranslatedTextList.append(command['parameters'][0])
+
+    batchList = splitIntoBatches(untranslatedTextList, 3500)  
+
+    # Translate Batches
+    translatedBatchList = []
+    for batch in batchList:
+        translatedBatchList.append(translateGPT(batch))
+    translatedText = ''.join(translatedBatchList)
+    translatedTextList = translatedText.split('/p')
+    translatedTextList = matchListSize(untranslatedTextList, translatedTextList)
+    
+    # Write to new json file
+    i = 0
+    for event in events:
+        if event is not None:
+            for page in event['pages']:
+                for command in page['list']:
+                    if command['code'] == 401:
+                        command['parameters'][0] = translatedTextList[i]
+                        i += 1
+
+    with open('file.json', 'w') as f:
+        json.dump(data, f)
+
+def splitIntoBatches(textList, n):
+    """This function takes a piece of text and sorts it into batches based on set token length."""
+    # Initialize the batch list
+    batches = []
+    
+    # Initialize the current batch
+    current_batch = []
+    
+    # Loop through each word in the text
+    for line in textList:
+        # If adding the current word to the current batch would exceed the max length,
+        # add the current batch to the list of batches and start a new batch
+        if len(' '.join(current_batch + [line])) > n:
+            batches.append(current_batch)
+            current_batch = [line]
+        # Otherwise, add the current word to the current batch
+        else:
+            current_batch.append(line)
+    
+    # Add any remaining words to the final batch
+    if len(current_batch) > 0:
+        batches.append(current_batch)
+    
+    return batches
+
+def translateGPT(t):
+    """Translate text using GPT"""
     system = "You are a professional Japanese visual novel translator. \
     You always manages to translate all of the little nuances of the original \
     Japanese text to your output, while still making it a prose masterpiece, \
     and localizing it in a way that an average American would understand. \
-    You always include the '/p' from the original text in your translation." \
+    You always include the '/p' from the original text in your translation."
+
+    # Convert to text
+    pipe = '/p'
+    t = pipe.join(t)
 
     response = openai.ChatCompletion.create(
         temperature=0.2,
-        max_tokens=500,
         model="gpt-3.5-turbo",
         messages=[
             {"role": "system", "content": system},
@@ -47,20 +107,12 @@ def translate(t):
     )
     return response.choices[0].message.content
 
-# Make translated list the same size
-translatedText += [""] * (len(text) - len(translatedText))
-for i in range(len(translatedText)):
-    translatedText[i] = translatedText[i].strip()
-    
-# Place translated text back in its original location in the JSON file
-i = 0
-for event in events:
-    if event is not None:
-        for page in event['pages']:
-            for command in page['list']:
-                if command['code'] == 401:
-                    command['parameters'][0] = translatedText[i]
-                    i += 1
+def matchListSize(list1, list2):
+    """Make list2 the same size as list1"""
+    list2 += [""] * (len(list1) - len(list2))
+    for i in range(len(list2)):
+        list2[i] = list2[i].strip()
 
-with open('file.json', 'w') as f:
-    json.dump(data, f)
+    return list2
+
+main()