Baichuan2 源码解析之 web_demo.py

强化学习曾小健

于 2023-09-09 16:45:13 发布

阅读量845

点赞数

分类专栏： # baichuan大模型 - 源码解析 LLM大语言模型文章标签：前端

本文链接：https://blog.csdn.net/sinat_37574187/article/details/132779405

版权

LLM大语言模型同时被 2 个专栏收录

564 篇文章 ¥199.90 ¥299.90

订阅专栏

超级会员免费看

baichuan大模型 - 源码解析

9 篇文章

订阅专栏

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

import json
import torch
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.utils import GenerationConfig


st.set_page_config(page_title="Baichuan 2")
st.title("Baichuan 2")


@st.cache_resource
def init_model():
    model = AutoModelForCausalLM.from_pretrained(
        "baichuan-inc/Baichuan2-13B-Chat",
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True
    )
    model.generation_config