gen

app.py

@@ -2,7 +2,6 @@
 import os, logging, torch, streamlit as st
 from transformers import (
     AutoTokenizer, AutoModelForCausalLM)
-st.balloons()
 
 # --------------------- HELPER --------------------- #
 def C(text, color="yellow"):
@@ -18,18 +17,31 @@ def C(text, color="yellow"):
     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():
+    
+    
+    
+    from transformers import GenerationConfig
+
+    # generation_config, unused_kwargs = GenerationConfig.from_pretrained(
+    #     "ckip-joint/bloom-1b1-zh", 
+    #     max_new_tokens=200,
+    #     return_unused_kwargs=True)
+    
+    
+    
+    
+    
+    
     tokenizer = AutoTokenizer.from_pretrained(
         "ckip-joint/bloom-1b1-zh")
     model = AutoModelForCausalLM.from_pretrained(
@@ -44,14 +56,10 @@ def model_init():
     return tokenizer, model
 
 tokenizer, model = model_init()
-st.balloons()
 
 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:
@@ -59,13 +67,13 @@ try:
         with torch.no_grad():
             [texts_out] = model.generate(
                 **tokenizer(
-                    prompt, return_tensors="pt"
+                    prompt, return_tensors="pt",
+                    max_new_tokens=200,
                 ).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()