中文模型使用uer/gpt2-chinese-cluecorpussmall

app.py

@@ -7,14 +7,21 @@ set_seed(42)
 def app():
     # 创建Streamlit应用程序
     st.title("使用gpt2的文本生成")
-    input_text = st.text_input("请输入您要生成的文本，只支持英文", value="")
+
+    options = ['中文','英文']
+    choice = st.radio('不同语言使用不同模型：', options)
+
+    input_text = st.text_input("请输入您要生成的文本", value="")
     maxlen = st.text_input("请输入生成文本的最大长度，越长越慢，不要超过1000", value="30")
     button_generate = st.button("生成")
     output_text = st.empty()
 
     def generate_text(input_text):
         # 加载预训练模型
-        generator = pipeline("text-generation", model="gpt2")
+        model="gpt2"
+        if choice == '中文':
+            model = 'uer/gpt2-chinese-cluecorpussmall'
+        generator = pipeline("text-generation", model)
 
         # 生成文本
         output = generator(input_text, max_length=int(maxlen), num_return_sequences=1)