探索Llama 3.1：开源模型的本地部署与创新应用实践

from transformers import AutoTokenizer, AutoModelForCausalLM

import torch

import streamlit as st



# 创建一个标题和一个副标题

st.title("💬 LLaMA3.1 Chatbot")

st.caption("🚀 A streamlit chatbot powered by Self-LLM")



# 定义模型路径

mode\_name\_or\_path = '/root/workspace/Llama-3.1-8B-Instruct'



# 定义一个函数，用于获取模型和tokenizer

@st.cache\_resource

def get\_model():

    # 从预训练的模型中获取tokenizer

    tokenizer = AutoTokenizer.from\_pretrained(mode\_name\_or\_path, trust\_remote\_code=True)

    tokenizer.pad\_token = tokenizer.eos\_token

    # 从预训练的模型中获取模型，并设置模型参数

    model = AutoModelForCausalLM.from\_pretrained(mode\_name\_or\_path, torch\_dtype=torch.bfloat16).cuda()

  

    return tokenizer, model



# 加载LLaMA3的model和tokenizer

tokenizer, model = get\_model()



# 如果session\_state中没有"messages"，则创建一个包含默认消息的列表

if "messages" not in st.session\_state:

    st.session\_state["messages"] = []



# 遍历session\_state中的所有消息，并显示在聊天界面上

for msg in st.session\_state.messages:

    st.chat\_message(msg["role"]).write(msg["content"])



# 如果用户在聊天输入框中输入了内容，则执行以下操作

if prompt := st.chat\_input():

    

    # 在聊天界面上显示用户的输入

    st.chat\_message("user").write(prompt)

    

    # 将用户输入添加到session\_state中的messages列表中

    st.session\_state.messages.append({"role": "user", "content": prompt})



    # 将对话输入模型，获得返回

    input\_ids = tokenizer.apply\_chat\_template(st.session\_state["messages"],tokenize=False,add\_generation\_prompt=True)

    model\_inputs = tokenizer([input\_ids], return\_tensors="pt").to('cuda')

    generated\_ids = model.generate(model\_inputs.input\_ids,max\_new\_tokens=512)

    generated\_ids = [

        output\_ids[len(input\_ids):] for input\_ids, output\_ids in zip(model\_inputs.input\_ids, generated\_ids)

    ]

    response = tokenizer.batch\_decode(generated\_ids, skip\_special\_tokens=True)[0]



    # 将模型的输出添加到session\_state中的messages列表中

    st.session\_state.messages.append({"role": "assistant", "content": response})

    # 在聊天界面上显示模型的输出

    st.chat\_message("assistant").write(response)

    print(st.session\_state)

streamlit run llamaBot.py --server.address 0.0.0.0 --server.port 1024