a

app.py

@@ -1,10 +1,33 @@
+# ------------------- LIBRARIES -------------------- #
 import os, logging, torch, streamlit as st
 from transformers import (
     AutoTokenizer, AutoModelForCausalLM)
 st.balloons()
 
-device = 'cpu'
+# --------------------- HELPER --------------------- #
+def C(text, color="yellow"):
+    color_dict: dict = dict(
+            red="\033[01;31m",
+          green="\033[01;32m",
+         yellow="\033[01;33m",
+           blue="\033[01;34m",
+        magenta="\033[01;35m",
+           cyan="\033[01;36m",
+    )
+    color_dict[None] = "\033[0m"
+    return (
+        f"{color_dict.get(color, None)}"
+        f"{text}{color_dict[None]}")
+st.balloons()
 
+# ------------------ ENVIORNMENT ------------------- #
+os.environ["HF_ENDPOINT"] = "https://huggingface.co"
+device = ("cuda" 
+    if torch.cuda.is_available() else "cpu")
+logging.info(C("[INFO] "f"device = {device}"))
+st.balloons()
+
+# ------------------ INITITALIZE ------------------- #
 @st.cache
 def model_init():
     tokenizer = AutoTokenizer.from_pretrained(
@@ -16,38 +39,33 @@ def model_init():
         # device_map="auto",
     # Ref. for `half`: Chan-Jan, Thanks!
     ).eval().to(device)
+    st.balloons()
+    logging.info(C("[INFO] "f"Model init success!"))
     return tokenizer, model
 
 tokenizer, model = model_init()
+st.balloons()
 
-prompt = '我是'
-with torch.no_grad():
-    [out] = model.generate(
-        **tokenizer(
-            prompt, return_tensors="pt"
-        ).to(device)
-    )
-    st.text(out)
-
-# 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
-
+try:
+    # ===================== INPUT ====================== #
+    # prompt = "\u554F\uFF1A\u53F0\u7063\u6700\u9AD8\u7684\u5EFA\u7BC9\u7269\u662F\uFF1F\u7B54\uFF1A"  #@param {type:"string"}
+    prompt = st.text_input("Prompt: ")
+    st.balloons()
+    
+
+    # =================== INFERENCE ==================== #
+    if prompt:
+        st.balloons()
+        with torch.no_grad():
+            [texts_out] = model.generate(
+                **tokenizer(
+                    prompt, return_tensors="pt"
+                ).to(device))
+        st.balloons()
+        output_text = tokenizer.decode(texts_out)
+        st.balloons()
+        st.markdown(output_text)
+        st.balloons()
+except Exception as err:
+    st.write(str(err))
+    st.snow()