app.py

import torch, os
import torch.nn as nn
import gradio as gr
import pandas as pd
from transformers import AutoTokenizer, AutoModel
from typing import Literal, Union, Dict, List, Tuple, Optional

MODEL_ID = "TAIDE-EDU/TBCL-Scorer-RoBERTa-2-12_DEMO"
HF_TOKEN = os.getenv("HF_TOKEN", None)

headers = ['id', '課文類型', 'TBCL等級', "課文標題", '課文內容']
device = "cuda" if torch.cuda.is_available() else "cpu"

model = AutoModel.from_pretrained(MODEL_ID, token=HF_TOKEN, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device).eval()

example_datas = [
    {
        "id": "L1T001\n(shi)",
        "教材名稱": "華師審查",
        "課文標題": "你好嗎？",
        "課文內容": "林春和：力冬，你今天怎麼樣？\n王力冬：很不錯，你呢？\n林春和：我也不錯。你要去哪裡？\n王力冬：我現在要去買東西。我們明天中文課見！\n林春和：明天見！",
        "list_level": "level_1",
        "類型": "對話"
    },
    {
        "id": "L1T004",
        "教材名稱": "華師審查",
        "課文標題": "今天下午做什麼？",
        "課文內容": "銘鈺：曉萱，今天下午你有什麼計畫嗎？\n曉萱：我還不知道，你呢？\n銘鈺：天氣這麼好，我們出去走走吧！\n曉萱：好啊！我們去象山怎麼樣？\n銘鈺：好，聽說那裡的風景很漂亮！可是要怎麼去？\n曉萱：我們可以坐捷運也可以坐公車。我覺得坐捷運比較快！\n銘鈺：那我們2點在中山站見面吧！\n曉萱：一會兒見！",
        "list_level": "level_1",
        "類型": "對話"
    },
    {
        "id": "L2T021",
        "教材名稱": "華師審查",
        "課文標題": "教育與學習",
        "課文內容": "小明：小美，妳最近學習的情況怎麼樣？\n小美：還不錯啦。不過有些部分還需要更努力。\n小明：妳最喜歡上哪一種課？\n小美：我比較喜歡上科學課，覺得科學課很有意思。\n小明：還有其他喜歡的課嗎？\n小美：數學課我也喜歡，因為數學課很有趣。\n小明：妳週末會看書嗎？\n小美：會啊，我有時候會複習科學和數學。\n小明：妳平常會怎麼準備呢？\n小美：我會做筆記和複習重點，讓我更了解學習的內容。",
        "list_level": "level_2",
        "類型": "對話"
    },
    {
        "id": "L2T027",
        "教材名稱": "華師審查",
        "課文標題": "寒假旅行計畫",
        "課文內容": "今天和朋友聊天，聊了寒假的計劃。我們都是第一次來台灣，想看看這裡的冬天和我們國家的有什麼不一樣。老師說台灣很少下雪，可是還是有很多漂亮的風景，我們都很想去看看。\n\n我的朋友很喜歡大海，我們打算在一月中到基隆玩。我的台灣朋友叫我們注意，雖然台灣不會下雪，可是天氣還是會非常冷，得多準備一些衣服。\n\n聽到我們的計畫，老師說我們也應該去基隆的夜市，因為有很多特別的小吃。這次旅行一定會很有意思，也希望我能拍到漂亮的照片給我的家人看。",
        "list_level": "level_2",
        "類型": "短文"
    },
    {
        "id": "L3T003",
        "教材名稱": "華師審查",
        "課文標題": "選擇一份適合的工作",
        "課文內容": "選擇一份適合自己的工作是生活中的重要決定。因為每個人的興趣和能力不同，所以選擇的工作也不同。有些人喜歡幫助別人，可能會選擇當醫生或老師。有些人對研究有興趣，這樣的人可能會選擇當研究員。選擇工作時，除了要考慮自己的興趣以外，也要考慮工作的環境，因為這些事情都會影響我們的心情。雖然每個工作都有困難的部分，可是只要努力，就能解決很多問題，工作也會越來越順利。許多大學生在畢業前就開始想以後的工作，因為找到一份適合的工作不但可以帶來收入，還能讓人過得更快樂。",
        "list_level": "level_3",
        "類型": "短文"
    },
    {
        "id": "L3T007",
        "教材名稱": "華師審查",
        "課文標題": "現代社會和傳統社會的差別",
        "課文內容": "王老師：你們覺得現代社會和傳統社會有什麼不一樣？\n天明：我覺得現代社會的科技進步很快，大家的生活也變了。以前的人寫信聯絡，現在我們用簡訊或者電子郵件。\n美心：對啊！還有，現在的年輕人不像以前那麼傳統，他們更自由了。\n王老師：傳統的價值觀，例如重視家庭，對現代社會有影響嗎？\n天明：影響很大。家人的關心和支持還是很重要。\n美心：是啊，節日時大家還是會一起慶祝，例如春節和中秋節，這些都是傳統對現代社會的影響。\n王老師：最後，請問你們覺得現代社會和傳統社會怎麼結合比較好？\n美心：我想，我們應該學習傳統的好處，同時用現代科技讓生活更方便。\n天明：我同意。",
        "list_level": "level_3",
        "類型": "對話"
    },
    {
        "id": "L4T008\n(shi)",
        "教材名稱": "華師審查",
        "課文標題": "讓小興趣變成大夢想",
        "課文內容": "每個人都有一個夢想，你希望自己長大後變成什麼樣的人呢？是站在舞台上，用歌聲感動觀眾的歌手？在醫院裡幫助病人的醫生？還是到過很多國家的的旅行家？\n選擇夢想時，興趣會帶著我們往前走。當我們對一件事有興趣時，學習就會變得更容易。例如：喜歡看書的人，可能會想當作家；喜歡做研究的人，可能想成為科學家。 \n夢想不會自己實現，它需要我們努力。就像種子要長大，需要陽光和水，為了實現夢想，我們也需要按照計劃，不斷學習和練習，才會有成果。 \n每個人的夢想都是不一樣的，就像世界上沒有兩片相同的葉子。夢想是什麼都沒關係，只要努力，就一定能實現。",
        "list_level": "level_4",
        "類型": "對話"
    },
    {
        "id": "L4T009\n(shi)",
        "教材名稱": "華師審查",
        "課文標題": "人與人的互動",
        "課文內容": "在現代社會，人與人的互動對每個人來說都非常重要。有良好的關係可以帶來許多好處，例如在工作上得到同事的幫助，在生活中得到朋友的支持。為了有好的關係，我們需要學會與別人溝通。與別人相處時，我們應該注意表達自己的想法，接受別人的意見。例如，當我們遇到困難時，可以向朋友請教，並且在他們需要幫助時，我們也應該主動提供幫助。這樣互相支持和幫助能夠加強感情。參加各種活動也是保持關係的好方法。通過這些活動，我們可以認識更多的人，增加自己的朋友。例如，參加朋友的聚餐、工作上的活動，以上這些都是認識新朋友的好機會。同時，保持良好的態度也是保持關係的重要因素。我們應該保持熱情、親切，並且在與人相處時多注意禮貌，這樣可以讓對方感覺到我們的親切。通過有效溝通、積極參加活動和保持良好的態度，我們可以和朋友有良好的關係。",
        "list_level": "level_4",
        "類型": "短文"
    },
    {
        "id": "L5T001\n(shi)",
        "教材名稱": "華師審查",
        "課文標題": "休閒活動",
        "課文內容": "在臺灣，假日是人們放鬆和享受生活的時光。每到假日，人們會選擇各種不同的休閒活動來度過這段時間。以下介紹幾種臺灣人常做的休閒活動。\n\n首先，許多人喜歡到郊外進行健行。例如，在新北市的陽明山國家公園，是許多健行愛好者的熱門去處。陽明山的七星山步道，每到假日總是吸引大批遊客。許多人會早起，踏上登山的步道，一邊呼吸新鮮空氣，一邊享受大自然的寧靜。健行不但可以鍛鍊身體，還能讓人放鬆心情，是一項非常受歡迎的活動。\n\n除了健行，騎自行車也是臺灣人假日的熱門選擇。台中的后豐鐵馬道，是一條將廢棄鐵路整修而成的自行車道。假日時，無論是家庭成員還是朋友團體，都喜歡一起騎著自行車，感受微風吹來的舒適。后豐鐵馬道全長12公里，沿路風景如畫，是台灣自行車愛好者的最愛。\n\n對喜愛美食的人來說，假日則是一個大吃大喝的好時機。士林夜市是臺北市最大的夜市，每到假日，這裡總是人山人海。無論是蚵仔煎、珍珠奶茶還是臭豆腐，都讓人口水直流。士林夜市不只提供了豐富的美食選擇，還是一個體驗當地文化的好地方。\n\n最後，看電影也是臺灣人假日的熱門活動之一。無論是獨自一人還是與親朋好友一起，走進戲院，觀看最新上映的電影，都是一種享受。特別是在炎熱的夏天，電影院內的涼爽空調更是讓人感到舒適。\n\n總之，臺灣人假日的休閒活動多種多樣，無論是親近自然、運動鍛鍊，還是享受美食、文化娛樂，每個人都能找到適合自己的假日休閒活動。這些活動不只豐富了人們的生活，還增進了家庭和朋友之間的感情，使假日變得更加充實和愉快。",
        "list_level": "level_5",
        "類型": "短文"
    },
    {
        "id": "L5T002\n(shi)",
        "教材名稱": "華師審查",
        "課文標題": "探索臺北之美",
        "課文內容": "臺北，作為臺灣的首都，不僅是政治與經濟中心，更是文化與歷史的交匯點。這座城市融合了現代與傳統，讓人每次造訪都有不同的驚喜。現在讓我們一起探索臺北之美吧！\n\n首先，臺北的夜市文化是一大亮點。士林夜市、饒河街夜市和寧夏夜市是其中最著名的幾個。這些夜市提供了豐富的街頭美食，從鹽酥雞、大腸包小腸到芒果冰、蚵仔煎，每一口都是臺灣味道的代表。夜市不僅是品嚐美食的地方，更是體驗臺灣夜生活和文化氛圍的好去處。不僅能滿足味蕾，更能讓人感受到臺北的獨特魅力。\n\n除了夜市，臺北的自然景觀也值得一提。陽明山國家公園以其壯麗的火山地形和溫泉聞名，是市民和遊客放鬆身心的好地方。而大安森林公園則被譽為「臺北的綠肺」，在這裡，你可以遠離城市的喧囂，享受自然的寧靜。這些自然景觀為城市增添了不少綠意和活力。\n\n臺北的歷史建築同樣充滿魅力。國立故宮博物院收藏了豐富的中國古代文物，讓人可以一覽千年歷史的風采。而龍山寺、保安宮等古老寺廟則展示了臺灣傳統宗教文化的深厚底蘊。這些歷史建築不僅是文化的象徵，更是了解臺北歷史的重要窗口。\n\n此外，臺北的交通系統非常便利。臺北捷運覆蓋了臺北市各個主要景點，不僅快捷而且乾淨，使得遊客可以輕鬆遊覽城市的各個角落。大多數捷運站附近有商場和餐廳，方便人們購物和用餐。這樣的交通便利性讓臺北成為一個適合觀光客探索的城市。\n\n總之，臺北是一座充滿活力與魅力的城市，無論是喜歡美食、自然還是歷史文化的人，都能在這裡找到自己的心之所向。來臺北旅行，不僅是一次視覺和味覺的享受，更是一次心靈的滿足。",
        "list_level": "level_5",
        "類型": "短文"
    }]

def _to_level_num(s: str):
    try:
        return int(str(s).split("_")[-1])
    except Exception:
        return None
def tbcl_multitask_judgement(text:str) -> Tuple[Dict[str, float], int]:
    """Predicts two task distributions for the input text."""
    
    batch_x = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)

    with torch.no_grad():
        out_dict = model(**batch_x, return_dict=True)
    cls_probs = out_dict["cls_logits"].cpu().tolist()[0]
    tbcl_level = (out_dict["reg_logits"].squeeze(-1)+1).cpu().tolist()[0]
    # 任務0（程度等級）
    classfication_id2label = {0: "捨棄", 1: "保留"}
    

    class_prediction = {classfication_id2label[i]: round(cls_probs[i], 3) for i in range(len(cls_probs))}
    return class_prediction, tbcl_level

def on_row_select(evt: gr.SelectData, df):
    row_idx = evt.index[0] if isinstance(evt.index, (list, tuple)) else evt.index
    row = df.iloc[int(row_idx)]
    # 這裡決定選取後要顯示什麼
    row_id = row["id"]
    prompt_str = example_datas_id2data_str[row_id]
    #questions = [row[f"題目{i}"] for i in range(1, 7)]
    #md = f"### 範例文章\n{article}\n\n### 題目\n" + "\n".join([f"- {q}" for q in questions])
    return prompt_str#, md

# 依 headers 的欄位順序組好 DataFrame
def _to_level_num(s: str):
    try:
        return int(str(s).split("_")[-1])
    except Exception:
        return None

example_datas_df = pd.DataFrame(
    [
        {
            "id": d["id"],
            "課文類型": d.get("類型", ""),
            "TBCL等級": _to_level_num(d.get("list_level", "")),
            "課文標題": d.get("課文標題", ""),
            "課文內容": d.get("課文內容", ""),
        }
        for d in example_datas
    ],
    columns=headers,  # 確保欄位順序與 headers 一致
)

# 點表格帶入文字時要用到：id -> 要帶入 textbox 的內容
example_datas_id2data_str = {d["id"]: d["課文內容"] for d in example_datas}

with gr.Blocks(title="課文內容TBCL評分器") as demo:
    gr.Markdown("## 課文內容TBCL評分器\n從下方**範例表**點選一列，系統會自動帶入並推論。")

    with gr.Row():
        inp = gr.Textbox(label="輸入文章", lines=6, placeholder="也可手動貼文後按下『送出』")
        btn = gr.Button("送出")

    with gr.Row():
        out0 = gr.Label(label="任務0：課文是否保留(限定輸入高階模型生成之內容)")
        out1 = gr.Label(label="任務1：TBCL等級預測")

    table = gr.Dataframe(
        value=example_datas_df,
        headers=headers,
        row_count=(len(example_datas_df), "fixed"),
        col_count=len(headers),
        interactive=False,
        wrap=True,
        label="課文範例表"
    )

    table.select(
        on_row_select,
        inputs=table,
        outputs=[inp]
    )

    # 也支援手動輸入後按『送出』
    btn.click(
        tbcl_multitask_judgement,
        inputs=inp,
        outputs=[out0, out1]
    )

demo.queue(max_size=32)

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