在本地调用vicuna-13B模型(简易版)

最新推荐文章于 2025-03-07 18:08:32 发布

AIcoderAR

最新推荐文章于 2025-03-07 18:08:32 发布

阅读量711

点赞数

CC 4.0 BY-SA版权

文章标签： linux 自然语言处理

本文链接：https://blog.youkuaiyun.com/weixin_42343812/article/details/131559979

该代码段展示了如何使用transformers库从/localdata/cn-customer-engineering/xudongz/langchain/langchain-ChatGLM/llama/vicuna-13B/加载预训练的LLAMA模型。模型被转换为半精度并在GPU上运行。用户输入提示后，模型使用do_sample策略生成最多1024个新tokens的响应。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

from transformers import (AutoConfig, AutoModel, AutoModelForCausalLM,
                          AutoTokenizer, LlamaTokenizer)

model_name = "/localdata/cn-customer-engineering/xudongz/langchain/langchain-ChatGLM/llama/vicuna-13B/"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
print("load tokenizer done!")
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
model.half().cuda()
while 1:
    prompt = input()
    inputs = tokenizer(prompt, return_tensors='pt').to('cuda')
    # print(inputs.type)
    del inputs['token_type_ids']
    print("inputs", inputs)
    tokens = model.generate(
    **inputs,
    max_new_tokens=1024,
    do_sample=True,
    temperature=1.0,
    top_p=1.0,
    )
    print(tokenizer.decode(tokens[0], skip_special_tokens=True))