Spaces:
Running
Running
| 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 nonesense 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-4-turbo-preview", | |
| 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() |