Update app.py

app.py

@@ -1,9 +1,8 @@
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
-import gradio as gr
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, pipeline
 
-# Load the pre-trained GPT2 model and tokenizer
-model = GPT2LMHeadModel.from_pretrained("gpt2")
-tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+# 使用緩存加速載入
+model = GPT2LMHeadModel.from_pretrained("gpt2", cache_dir="./model_cache")
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2", cache_dir="./model_cache")
 
 # 設置填充標記 ID
 tokenizer.pad_token = tokenizer.eos_token
@@ -28,18 +27,4 @@ def generate_command(prompt, max_length=100):
         end = len(generated_text)
     command = generated_text[start:end].strip()
     
-    return command
-
-def predict(input_text):
-    output = generate_command(input_text)
-    return output
-
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Textbox(lines=2, placeholder="請輸入你的指令生成提示..."),
-    outputs="text",
-    title="使用 GPT2 生成指令",
-    description="根據你的中文輸入提示生成 Bash 指令。"
-)
-
-iface.launch()
+    return command