feat: Huge improvement

This commit is contained in:
Dazed 2023-04-04 19:51:26 -05:00
parent c0401374e5
commit ef1758aba7

163
main.py
View file

@ -1,16 +1,14 @@
import os
import re
import textwrap
from googletrans import Translator
from pathlib import Path
import json
import time
from dotenv import load_dotenv
import openai
load_dotenv()
openai.organization = os.getenv('org')
openai.api_key = os.getenv('key')
pipe = '###'
def main():
# Load JSON file into a Python dictionary
@ -18,6 +16,35 @@ def main():
data = json.load(f)
parse401(data)
def count401Groups(data):
"""
Counts the number of groups of 401 in a dictionary of events.
Args:
data (list): A list of events.
Returns:
int: The number of groups of 401.
"""
# Count 401 Groups
events = data['events']
groupCount = 0
currentGroupCount = 0
for event in events:
if event is not None:
for page in event['pages']:
for command in page['list']:
if command['code'] == 401:
currentGroupCount += 1
else:
if currentGroupCount > 0:
groupCount += 1
currentGroupCount = 0
if currentGroupCount > 1:
groupCount += 1
currentGroupCount = 0
return groupCount
def parse401(data):
"""
Extracts the first parameter of all commands with code 401 from a dictionary of events.
@ -28,82 +55,70 @@ def parse401(data):
Returns:
list: A list of strings containing the first parameter of all commands with code 401.
"""
# Extract 401 Text
events = data['events']
# Extract 401 Text (Needed for Context)
translatedText = ''
untranslatedTextList = []
currentGroup = []
textHistory = []
events = data['events']
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])
try:
for i in range(len(page['list'])):
if page['list'][i]['code'] == 401:
currentGroup.append(page['list'][i]['parameters'][0])
textHistory.append(page['list'][i]['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 = wordwrapList(translatedTextList)
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:
if command['parameters'][0] == '':
del command['parameters'][0]
while page['list'][i+1]['code'] == 401:
del page['list'][i]
currentGroup.append(page['list'][i]['parameters'][0])
textHistory.append(page['list'][i]['parameters'][0])
else:
command['parameters'][0] = translatedTextList[i]
i += 1
with open('file.json', 'w') as f:
if len(currentGroup) > 0:
text = ''.join(currentGroup)
print('Translating' + text)
translatedText = translateGPT(text, ''.join(textHistory))
page['list'][i-1]['parameters'][0] = translatedText
if len(textHistory) > 50:
for i in range(len(currentGroup)):
textHistory.pop(0)
currentGroup = []
except IndexError:
print('End of List')
# Append leftover groups
if len(currentGroup) > 0:
translatedText = translateGPT(''.join(currentGroup), ' '.join(textHistory))
page['list'][i]['parameters'][0] = translatedText
currentGroup = []
with open('file.json', 'w', encoding='utf-8') 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, history):
pattern = r'[一-龠]+|[ぁ-ゔ]+|[ァ-ヴー]+'
if not re.search(pattern, t):
return t
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."
# Convert to text
pipe = '/p'
t = pipe.join(t)
system = "Context: " + history + "\n\n###\n\n You are a professional Japanese visual novel translator, editor, and localizer, \
with a writing style could only be described as being a mixture of \
John Steinbeck's, Oscar Wilde's, and Vladimir Nabokov's writing styles. \
You always manage to carry 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 read and understand the 'Context' \
You translate only what I give you. \
You include line breaks in your translation. \
You translate sound effects and Onomatopoeia in their romanji form. \
You translate anything inside '<>' as a name"
response = openai.ChatCompletion.create(
temperature=0,
max_tokens = 3000,
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": system},
@ -111,21 +126,5 @@ def translateGPT(t):
]
)
return response.choices[0].message.content
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()
return list2
def wordwrapList(list):
wrappedList = []
for item in list:
wrappedList.append(textwrap.fill(item, width=50))
return wrappedList
main()