实用模型推荐(二)中译英翻译模型:opus-mt-zh-en

1.开源地址:https://huggingface.co/Helsinki-NLP/opus-mt-zh-en

2.使用场景:中译英,多模型场景的中英转换

3.API封装

import uvicorn
from fastapi import FastAPI
from loguru import logger
from pydantic import BaseModel
from starlette.middleware.cors import CORSMiddleware
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-zh-en")

model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-zh-en")

# define the app
app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)


@app.get('/')
async def index():
    return {"message": "index, docs url: /docs"}


class InputData(BaseModel):
    text: str


@app.post('/translate')
async def emb(data: InputData):
    try:
        encoded = tokenizer([data.text], return_tensors="pt")

        translation = model.generate(**encoded)

        result = tokenizer.batch_decode(translation, skip_special_tokens=True)[0]

        logger.info("from:{} to:{}", data.text, result)

        return result
    except Exception as e:
        logger.error(e)
        return {'status': False, 'msg': str(e)}, 400


if __name__ == '__main__':
    uvicorn.run(app=app, host='0.0.0.0', port=9970)