马上注册，结交更多好友，享用更多功能，让你轻松玩转社区。
您需要 登录 才可以下载或查看，没有账号？立即注册
x
from transformers import pipeline
from accelerate import init_empty_weights, infer_auto_device_map
from transformers import AutoModelForCausalLM, AutoTokenizer
初始化加速器
from accelerate import Accelerator
accelerator = Accelerator()
加载模型和 tokenizer
tokenizer = AutoTokenizer.from_pretrained(“/app/model/”)
model = AutoModelForCausalLM.from_pretrained(“/app/model/”)
将模型分布到多张 GPU 上
device_map = infer_auto_device_map(model, max_memory={0: “16GB”, 1: “16GB”, 2: “16GB”, 3: “16GB”, 4: “16GB”, 5: “16GB”, 6: “16GB”, 7: “16GB”})
将模型加载到设备上
model = model.from_pretrained(“/app/model/”, device_map=device_map)
使用 pipeline
generator = pipeline(“text-generation”, model=model, tokenizer=tokenizer, device_map=device_map)
进行推理
output = generator(“Your input text”, max_length=30)
print(output)
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。
|
