AI学习- PDF 实践

最近在极客时间学习《AI 大模型应用开发实战营》，自己一边跟着学一边开发了一个进阶版本的 OpenAI-Translator，在这里简单记录下开发过程和心得体会，供有兴趣的同学参考。

一、AI模型定义

类继承关系：

      Model  -> OpenAIModel

                  -> GLMModel

定义Model工厂类：ModelFactory， 实现从参数实例化

  

class ModelFactory:
    @classmethod
    def create_model(cls, config:Config):
        if config.common.model_type == 'OpenAIModel':
            return OpenAIModel(model=config.OpenAIModel.model, api_key=config.OpenAIModel.api_key)
        elif config.common.model_type == 'GLMModel':
            return GLMModel(model_url=config.GLMModel.model_url, timeout=config.GLMModel.timeout)
        else:
            raise Exception(f'unsupport model:{config.common.model_type}')

     

二、提示词定义

翻译PDF中文本和表格的内容， 提示词分别设计为：

class Model:
    def make_text_prompt(self, text: str, target_language: str) -> str:
        return f"翻译为{target_language}：{text}"

    def make_table_prompt(self, table: str, target_language: str) -> str:
        return f"翻译为{target_language}，保持间距（空格，分隔符），以表格形式返回：\n{table}"

三、提供API （学会利用ChatGPT生成代码）

1.支持并发：使用子线程执行翻译任务

2. 客户端通过API 轮询处理进度/结果（改进，可用websocket实现服务端推送）

from flask import Flask, request, jsonify
import os
import time
import uuid
from concurrent.futures import ThreadPoolExecutor

from ai_translator.utils import LOG

app = Flask(__name__)

UPLOAD_FOLDER = './uploads'
ALLOWED_EXTENSIONS = {'pdf'}

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024

executor = ThreadPoolExecutor(2)

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def translate_pdf(filename):
    # 进行翻译处理
    time.sleep(10) # 模拟翻译过程
    result = 'This is a translated PDF file.'
    with open(os.path.join(app.config['UPLOAD_FOLDER'], f'{filename}.txt'), 'w') as f:
        f.write(result)

@app.route('/translate', methods=['POST'])
def translate_file():
    # 检查是否有文件上传
    if 'file' not in request.files:
        LOG.error('No file uploaded')
        return 'No file uploaded', 400

    file = request.files['file']

    # 检查文件是否为PDF格式
    if not allowed_file(file.filename):
        return f'Only support type:{ALLOWED_EXTENSIONS}', 400

    # 生成唯一的任务ID
    task_id = str(uuid.uuid4())

    # 保存上传的文件
    filename = f'{task_id}.pdf'
    file.save(os.path.join(app.config['UPLOAD_FOLDER'], filename))
    LOG.info(f'go to translate file:{file.filename}, local file:{filename}')

    # 异步进行翻译处理
    executor.submit(translate_pdf, filename)

    return jsonify({'taskId': task_id, ''})

@app.route('/result/<task_id>', methods=['GET'])
def get_result(task_id):
    # 获取指定任务的翻译结果
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], f'{task_id}_translate.pdf')
    result = 'Translation is in progress...'
    if os.path.isfile(filepath):
        with open(filepath, 'r') as f:
            result = f.read()
    return jsonify({'result': result})

if __name__ == '__main__':
    app.run(port=8000)