Update backend/ai_output.py

backend/ai_output.py

@@ -1,7 +1,6 @@
-# backend/ai_output.py
-
 from fastapi import HTTPException
 from .model_loader import model, tokenizer
+import torch
 
 def run_ai_output(input_text: str):
 
@@ -12,11 +11,21 @@ def run_ai_output(input_text: str):
         raise HTTPException(status_code=400, detail="輸入過長，請限制在 512 字元內")
 
     try:
-        input_ids = tokenizer.encode(input_text, return_tensors='pt', truncation=True)
+        # ⭐ 正確 attention_mask
+        encoded = tokenizer(
+            input_text,
+            return_tensors="pt",
+            truncation=True,
+            padding=False
+        )
+
+        input_ids = encoded["input_ids"]
+        attention_mask = encoded["attention_mask"]
 
         output = model.generate(
-            input_ids,
-            max_length=len(input_text) + 150,
+            input_ids=input_ids,
+            attention_mask=attention_mask,   # ⭐ 加上這行
+            max_length=input_ids.shape[1] + 150,
             num_return_sequences=3,
             no_repeat_ngram_size=3,
             do_sample=True,
@@ -34,7 +43,6 @@ def run_ai_output(input_text: str):
 
         completions = sorted(list(set(completions)), key=len, reverse=True)
 
-        # ⭐ 統一回傳格式（AI skill）
         return {
             "type": "ai-multi-options",
             "options": completions or [input_text]