DEV Community

drake
drake

Posted on

调用OpenAI (chatGPT) API 翻译PDF文档

from openai import OpenAI
from PyPDF2 import PdfReader
from fpdf import FPDF

# openai API的使用方式:https://platform.openai.com/docs/quickstart?language-preference=python
client = OpenAI()

class PDF(FPDF):
    def __init__(self):
        super().__init__()
        # 添加中文字体(以苹方为例,macOS默认字体)
        self.add_font('kaiti', '', './kaiti.ttf', uni=True)

    def footer(self):
        self.set_y(-15)
        self.set_font('kaiti', '', 8)  # 设置中文字体
        self.cell(0, 10, f'{self.page_no()}', 0, 0, 'C')  # 页脚

class Topdf:
    def __init__(self, pdf_source_path, txt_file_path):
        # PDF文件路径
        self.pdf_file = pdf_source_path
        # 输出的TXT文件路径
        self.txt_file = txt_file_path

    def to_pdf(self):
        # 创建PDF对象
        pdf = PDF()
        pdf.set_left_margin(10)
        pdf.set_right_margin(10)
        pdf.add_page()

        pdf.set_auto_page_break(auto=True, margin=15)
        pdf.set_font('kaiti', '', 12)

        # 读取TXT文件并优化文本内容
        with open(self.txt_file, 'r', encoding='utf-8') as file:
            text = file.read()

        # 将优化后的文本添加到PDF
        pdf.multi_cell(0, 10, text)

        # 保存PDF文件
        pdf.output(self.pdf_file)
        print(f"TXT内容已成功保存到 {self.pdf_file}")


class Translate:
    """
    GPT-4o-mini:输入费用为每百万 tokens $0.15,输出费用为每百万 tokens $0.60
    批处理打五折
    """
    def __init__(self):
        self.file_path = "origin.pdf"
        self.output_txt = "output_translated1.txt"
        self.output_pdf = "output_translated1.pdf"
        self.text_list = []

    def extract_text_from_pdf_translate(self, interupt=None):
        """
        抽取每一页的PDF提交给chatgpt翻译
        interupt:上一次翻译异常导致退出的页码,None表示没有任何异常导致中途退出
        """
        reader = PdfReader(self.file_path)
        text = ""
        num = 0
        for page in reader.pages:
            num += 1
            # 跳过前面已经翻译过的页面,从上一次翻译异常的页面重新开始
            if interupt:
                if num < interupt:
                    continue
            print(f'开始处理第 {num}')
            page_text = page.extract_text()
            chinese = self.translate(page_text)
            if chinese:
                text += f'{chinese}\n'
                self.text_list.append(chinese)
            else:
                print(f'{num} 页失败')
                break

        print(f'全部翻译完成\n\n{text}')
        return text

    def translate(self, text_origin):
        try:
            response = client.chat.completions.create(
                # model="gpt-4",
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": "You are a translation assistant."},
                    {"role": "user", "content": f"将该文本翻译成中文: {text_origin}"}
                ]
            )
            return response.choices[0].message.content
        except:
            return

    def save_to_pdf(self, text):
        """
        保存翻译后的内容为 txt, PDF
        """
        # 保存为txt
        with open(self.output_txt, 'w')as f:
            f.write(text)
        # 保存为pdf
        Topdf(self.output_pdf, self.output_txt).to_pdf()

    def run(self):
        text = self.extract_text_from_pdf_translate()
        self.save_to_pdf(text)

if __name__ == '__main__':
    Translate().run()
Enter fullscreen mode Exit fullscreen mode

Top comments (0)