Update ai_sentence.py

ai_sentence.py

@@ -1,32 +1,36 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import re
 
+# 可用模型清單
 MODEL_LIST = [
     "EleutherAI/pythia-410m",
-    "EleutherAI/pythia-1b",
-    "gpt2"
+    "gpt2",
+    "mistralai/Mistral-7B-Instruct"
 ]
 
-model_cache = {}  # 緩存模型
-
+# 模型快取，避免每次重新載入
+MODEL_CACHE = {}
 
+# 加載模型
 def load_model(model_name):
-    if model_name not in model_cache:
-        print(f"⏳ 正在載入模型：{model_name}")
+    if model_name not in MODEL_CACHE:
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForCausalLM.from_pretrained(model_name)
-        model_cache[model_name] = (tokenizer, model)
-    return model_cache[model_name]
-
+        MODEL_CACHE[model_name] = (tokenizer, model)
+    return MODEL_CACHE[model_name]
 
+# 生成 AI 例句
 def generate_sentence(word, model_name):
     tokenizer, model = load_model(model_name)
 
-    prompt = f"A simple English sentence using the word '{word}' suitable for beginners. Output only the sentence."
+    prompt = f"Example sentence using '{word}':"
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=30)
+    outputs = model.generate(**inputs, max_new_tokens=20)
     sentence = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # 清理句子（可以根據需求調整）
-    sentence = sentence.split(".")[0].strip() + "."
-
+    # 清理生成句子
+    sentence = sentence.split(":")[-1].strip()
+    sentence = re.sub(r'[^a-zA-Z0-9, .!?]', '', sentence)
+    if not sentence.endswith("."):
+        sentence += "."
     return sentence