import gradio as gr
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="根據提供的文本回答問題。"
)

if __name__ == "__main__":
    demo.launch()