Update web_demo.py

baichuan-inc · Jul 11, 2023 · d12a878 · d12a878
1 parent 035f701
commit d12a878
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/web_demo.py b/web_demo.py
@@ -1,7  1,7 @@
 import json
 import torch
 import streamlit as st
-from transformers import AutoModel, AutoTokenizer
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers.generation.utils import GenerationConfig
 
 
@@ -11,7  11,7 @@
 
 @st.cache_resource
 def init_model():
-    model = AutoModel.from_pretrained(
     model = AutoModelForCausalLM.from_pretrained(
         "baichuan-inc/Baichuan-13B-Chat",
         torch_dtype=torch.float16,
         device_map="auto",