devbox
IT小白
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

Baichuan2 源码解析之 web_demo.py_if prompt := st.chat_input("shift + enter 换行, ente

作者：IT小白 | 2024-02-21 23:50:38

赞

踩

if prompt := st.chat_input("shift + enter 换行, enter 发送"):

 github源码解析：https://github.com/ArtificialZeng/Baichuan2-Explained

 

  1. import json
  2. import torch
  3. import streamlit as st
  4. from transformers import AutoModelForCausalLM, AutoTokenizer
  5. from transformers.generation.utils import GenerationConfig
  6.  
  7.  
  8. st.set_page_config(page_title="Baichuan 2")
  9. st.title("Baichuan 2")
  10.  
  11.  
  12. @st.cache_resource
  13. def init_model():
  14.     model = AutoModelForCausalLM.from_pretrained(
  15.         "baichuan-inc/Baichuan2-13B-Chat",
  16.         torch_dtype=torch.float16,
  17.         device_map="auto",
  18.         trust_remote_code=True
  19.     )
  20.     model.generation_config = GenerationConfig.from_pretrained(
  21.         "baichuan-inc/Baichuan2-13B-Chat"
  22.     )
  23.     tokenizer = AutoTokenizer.from_pretrained(
  24.         "baichuan-inc/Baichuan2-13B-Chat",
  25.         use_fast=False,
  26.         trust_remote_code=True
  27.     )
  28.     return model, tokenizer
  29.  
  30.  
  31. def clear_chat_history():
  32.     del st.session_state.messages
  33.  
  34.  
  35. def init_chat_history():
  36.     with st.chat_message("assistant", avatar='       本文内容由网友自发贡献，转载请注明出处：https://www.wpsshop.cn/w/IT小白/article/detail/127365
    推荐阅读
    相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。