File size: 3,793 Bytes
435fc98 |
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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import gradio as gr
from io import BytesIO
import fitz
import tempfile
import openai
class TranslationAgent:
def __init__(self, openai_key):
self.memory = []
system_msg = "You are a translator from english to Chinese.\n" \
" The only thing you do is to translate.\n" \
" You don't write anything other then the translation of the text you get.\n" \
" The user will only provide the text without asking anything, but what he wants is the translation.\n" \
" Never return the translation of a previously translated part!\n " \
"The text you will need to translate will often include none sense stuff because it is coming from a text extraction of a pdf file including images and table.\n" \
" Do your best to translate also this messy parts."
self.memory.append({"role": "system", "content": system_msg})
openai.api_key = openai_key
def fade_memory(self):
if len(self.memory) >= 5:
del self.memory[1:3]
def translate_chunk(self, chunk):
self.memory.append({"role": "user", "content": chunk})
response = openai.ChatCompletion.create(
model="gpt-3.5-turbo",
messages=self.memory
)
reply = response["choices"][0]["message"]["content"]
self.memory.append({"role": "assistant", "content": reply})
self.fade_memory()
return reply
def extract_text_from_pdf(pdf, start, stop):
text = ""
with fitz.open(stream=BytesIO(pdf), filetype='pdf') as doc: # remove .read()
for i, page in enumerate(doc):
if start <= i:
if i <= stop:
text += page.get_text()
else:
break
return text
def split_text(text, chunk_size=100):
words = text.split()
chunks = []
current_chunk_words = []
for word in words:
current_chunk_words.append(word)
if word.endswith('.') and len(current_chunk_words) >= chunk_size:
chunks.append(' '.join(current_chunk_words))
current_chunk_words = []
# add the last chunk if any words remain
if current_chunk_words:
chunks.append(' '.join(current_chunk_words))
return chunks
def translate_pdf(openai_key, pdf, start, stop):
translator = TranslationAgent(openai_key)
translated_text = ""
error_message = "Translation Successful"
try:
# extract text
if pdf is not None:
text = extract_text_from_pdf(pdf, start=start, stop=stop)
chunks = split_text(text)
translated_chunks = []
for chunk in chunks:
translated_chunk = translator.translate_chunk(chunk)
translated_chunks.append(translated_chunk + " ")
translated_text = ' '.join(translated_chunks)
except Exception as e:
error_message = f"Translation Failed: {e}"
# Create a temporary file with a specific prefix
temp = tempfile.NamedTemporaryFile(delete=False, prefix="translatedPDF_", suffix=".txt")
# Write to the temporary file
with open(temp.name, 'w', encoding='utf-8') as f:
f.write(translated_text)
return translated_text, error_message, temp.name
iface = gr.Interface(
fn=translate_pdf,
inputs=[
gr.Textbox(lines=1, label="OpenAI API key",
placeholder="Enter your OpenAI API key here"),
gr.File(type="binary", label="PDF file", ),
gr.Number(label="Starting Page", ),
gr.Number(label="Final Page")
],
outputs=["text", "text", gr.File(label="Translated Text File")],
title="Pdf Translator: English ==> Chinese",
)
iface.launch() |