使用 Flask 框架和 Transformers 库搭建一个简单的中英文翻译服务。我们将使用 Helsinki-NLP 提供的 MarianMTModel 模型进行中英文之间的翻译

from flask import Flask, request, jsonify
from transformers import MarianMTModel, MarianTokenizer

# 创建Flask应用实例
app = Flask(__name__)

# 指定要使用的模型名称
zh_to_en_model_name = 'Helsinki-NLP/opus-mt-zh-en'  # 中文到英文的翻译模型
en_to_zh_model_name = 'Helsinki-NLP/opus-mt-en-zh'  # 英文到中文的翻译模型

# 加载预训练的tokenizer和模型
zh_to_en_tokenizer = MarianTokenizer.from_pretrained(zh_to_en_model_name)
zh_to_en_model = MarianMTModel.from_pretrained(zh_to_en_model_name)

en_to_zh_tokenizer = MarianTokenizer.from_pretrained(en_to_zh_model_name)
en_to_zh_model = MarianMTModel.from_pretrained(en_to_zh_model_name)


# 定义一个路由,用于处理POST请求,负责中文到英文的翻译
@app.route('/translate_zh_to_en', methods=['POST'])
def translate_zh_to_en():
    # 从请求中获取JSON数据中的'text'字段
    text = request.json.get('text')
    # 如果未提供文本,则返回400错误码和错误信息
    if not text:
        return jsonify({'error': '未提供文本'}), 400

    # 使用tokenizer将输入文本编码为模型所需的张量格式
    inputs = zh_to_en_tokenizer.encode(text, return_tensors="pt", padding=True)

    # 使用翻译模型生成翻译后的文本,设置最大输出长度为100
    translated = zh_to_en_model.generate(inputs, max_length=100)

    # 使用tokenizer将生成的张量解码为可读的文本
    translated_text = zh_to_en_tokenizer.decode(translated[0], skip_special_tokens=True)

    # 返回翻译后的文本,使用JSON格式
    return jsonify({'text': translated_text})


# 定义一个路由,用于处理POST请求,负责英文到中文的翻译
@app.route('/translate_en_to_zh', methods=['POST'])
def translate_en_to_zh():
    # 从请求中获取JSON数据中的'text'字段
    text = request.json.get('text')
    # 如果未提供文本,则返回400错误码和错误信息
    if not text:
        return jsonify({'error': '未提供文本'}), 400

    # 使用tokenizer将输入文本编码为模型所需的张量格式
    inputs = en_to_zh_tokenizer.encode(text, return_tensors="pt", padding=True)

    # 使用翻译模型生成翻译后的文本,设置最大输出长度为100
    translated = en_to_zh_model.generate(inputs, max_length=100)

    # 使用tokenizer将生成的张量解码为可读的文本
    translated_text = en_to_zh_tokenizer.decode(translated[0], skip_special_tokens=True)

    # 返回翻译后的文本,使用JSON格式
    return jsonify({'text': translated_text})


# 如果脚本是直接运行的(而不是作为模块导入的),则启动Flask应用
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8082)