如何让huggingface模型运行速度起飞

最新推荐文章于 2024-08-08 11:33:36 发布

狗庄欺人太甚

最新推荐文章于 2024-08-08 11:33:36 发布

阅读量2.1k

点赞数 2

分类专栏： NLP 文章标签：深度学习人工智能 python

本文链接：https://blog.csdn.net/weixin_48262500/article/details/126272169

版权

NLP 专栏收录该内容

16 篇文章 1 订阅

订阅专栏

首先调用接口的时候不要调用一次就加载一下模型，直接一开始就加载好

其次能用GPU就不要cpu

import os
import json
from tqdm import tqdm
from transformers import pipeline, AutoModelWithLMHead, AutoTokenizer
from tqdm import tqdm
import torch
torch.cuda.is_available()
def get_en_to_zh_model():
    model = AutoModelWithLMHead.from_pretrained("Helsinki-NLP/opus-mt-en-zh")
    tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-zh")
    translation = pipeline("translation_en_to_zh", model=model, tokenizer=tokenizer,device=1)
    return translation
translation = get_en_to_zh_model()




def en_to_ch(text):
    # 英文翻译成中文
    #text = "Student accommodation centres, resorts"
    translated_text = translation(text, max_length=888)[0]['translation_text']
    return translated_text

大家需要把pytorch换成GPU版本