L1-8G 显存玩转书生大模型 Demo

inscredion

已于 2024-07-30 19:12:48 修改

阅读量186

点赞数 2

分类专栏： LLM 文章标签：语言模型

于 2024-07-29 22:49:52 首次发布

本文链接：https://blog.csdn.net/inscredion/article/details/140782980

版权

LLM 专栏收录该内容

1 篇文章 0 订阅

订阅专栏

L1-8G 显存玩转书生大模型 Demo

闯关

8G cli_demo部署

1 .首先下载完权重，使用download_mini.py代码，内容如何
在这里插入图片描述

现在完成后，按照教程部署，cli_demo代码简单

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name_or_path = "/root/models/Shanghai_AI_Laboratory/internlm2-chat-1_8b"

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True, device_map='cuda:0')
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map='cuda:0')
model = model.eval()

system_prompt = """You are an AI assistant whose name is InternLM (书生·浦语).
- InternLM (书生·浦语) is a conversational language model that is developed by Shanghai AI Laboratory (上海人工智能实验室). It is designed to be helpful, honest, and harmless.
- InternLM (书生·浦语) can understand and communicate fluently in the language chosen by the user such as English and 中文.
"""

messages = [(system_prompt, '')]

print("=============Welcome to InternLM chatbot, type 'exit' to exit.=============")

while True:
    input_text = input("\nUser  >>> ")
    input_text = input_text.replace(' ', '')
    if input_text == "exit":
        break

    length = 0
    for response, _ in model.stream_chat(tokenizer, input_text, messages):
        if response is not None:
            print(response[length:], flush=True, end="")
            length = len(response)