Update app.py

app.py

@@ -1,68 +1,34 @@
 import gradio as gr
-import torch
-from PIL import Image
-import pickle
-from torchvision import transforms
-from your_model_file import ImageCaptionAttentionModel, prepare_model, pred_caption  # 你的模型和輔助函式
+from transformers import pipeline
+
+# 1. 載入 SQuAD v2.0 預訓練模型
+qa_model = pipeline("question-answering", model="deepset/roberta-base-squad2")
+
+# 2. 定義處理邏輯
+def predict(context, question):
+    if not context or not question:
+        return "請輸入文件內容與問題。"
+    
+    # 執行問答
+    result = qa_model(question=question, context=context)
+    
+    # 如果信心分數太低，回傳無法回答（SQuAD v2.0 特色）
+    if result['score'] < 0.05:
+        return "抱歉，在文件中找不到相關答案。"
+    
+    return result['answer']
+
+# 3. 建立 Gradio 網頁介面
+demo = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Textbox(lines=10, label="Context (文件內容)", placeholder="請貼上文件內容..."),
+        gr.Textbox(lines=2, label="Question (提問)", placeholder="請問這份文件關於什麼？")
+    ],
+    outputs=gr.Textbox(label="Model Answer (模型回答)"),
+    title="Case Study: Document QA System",
+    description="這是一個基於 SQuAD v2.0 訓練的模型，能根據提供的文本回答問題。"
+)
 
-# ------------------------
-# 1️⃣ 設定裝置
-# ------------------------
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-# ------------------------
-# 2️⃣ 載入詞表
-# ------------------------
-with open('voc_dict.pkl', 'rb') as f:
-    voc_dict = pickle.load(f)
-
-encoder = lambda x: voc_dict[x]
-decoder = lambda x: [key for key, value in voc_dict.items() if value == x][0]
-
-# ------------------------
-# 3️⃣ 圖像前處理
-# ------------------------
-image_transform = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                         std=[0.229, 0.224, 0.225])
-])
-
-# ------------------------
-# 4️⃣ 載入模型
-# ------------------------
-voc_size = len(voc_dict)
-embed_dim = 256  # 根據你的訓練設定
-model = ImageCaptionAttentionModel(voc_size, embed_dim)
-model = prepare_model(model)
-model.load_state_dict(torch.load('image_caption_attention_model.pth', map_location=device))
-model.eval()
-
-# ------------------------
-# 5️⃣ Gradio 推理函式
-# ------------------------
-def generate_caption(image: Image.Image):
-    # 前處理
-    img_tensor = image_transform(image)
-    # 生成 caption
-    caption = pred_caption(model, img_tensor)
-    return caption
-
-# ------------------------
-# 6️⃣ Gradio UI
-# ------------------------
-with gr.Blocks() as demo:
-    gr.Markdown("# Image Captioning with Attention")
-    image_input = gr.Image(type="pil", label="Upload Image")
-    result = gr.Textbox(label="Generated Caption")
-    run_button = gr.Button("Generate Caption")
-
-    run_button.click(fn=generate_caption, inputs=image_input, outputs=result)
-
-# ------------------------
-# 7️⃣ 啟動介面
-# ------------------------
 if __name__ == "__main__":
     demo.launch()