import os, logging, torch, streamlit as st
from transformers import (
    AutoTokenizer, AutoModelForCausalLM)
st.balloons()

device = 'cpu'

@st.cache
def model_init():
    tokenizer = AutoTokenizer.from_pretrained(
        "ckip-joint/bloom-1b1-zh")
    model = AutoModelForCausalLM.from_pretrained(
        "ckip-joint/bloom-1b1-zh",
        # Ref.: Eric, Thanks!
        # torch_dtype="auto", 
        # device_map="auto",
    # Ref. for `half`: Chan-Jan, Thanks!
    ).eval().to(device)
    return tokenizer, model

tokenizer, model = model_init()

st.text(str(repr(model)))

# DONE 6.1s

# ===== Application Startup at 2023-02-23 17:51:48 =====

# 2023-02-23 18:52:26.009 INFO    matplotlib.font_manager: generated new fontManager

# Collecting usage statistics. To deactivate, set browser.gatherUsageStats to False.


#   You can now view your Streamlit app in your browser.

#   Network URL: http://10.19.49.246:8501
#   External URL: http://34.197.127.12:8501


#   A new version of Streamlit is available.

#   See what's new at https://discuss.streamlit.io/c/announcements

#   Enter the following command to upgrade:
#   $ pip install streamlit --upgrade