GUI.py

import argparse
import tkinter as tk
from tkinter import ttk, scrolledtext

import torch
import torch.nn as nn
from transformers import (
    BertTokenizerFast,
    BertModel,
    AutoTokenizer,
    T5ForConditionalGeneration
)


# =========================
#  一、情感分类模型（BERT）
# =========================

class SentimentClassifier(nn.Module):
    """
    要和你训练/预测脚本里的结构一致：
    BERT + Dropout + Linear(num_labels)
    """
    def __init__(self, pretrained_name="bert-base-uncased", dropout_prob=0.1, num_labels=2):
        super().__init__()
        self.bert = BertModel.from_pretrained(pretrained_name)
        hidden_size = self.bert.config.hidden_size
        self.dropout = nn.Dropout(dropout_prob)
        self.classifier = nn.Linear(hidden_size, num_labels)

    def forward(self, input_ids, attention_mask):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        cls_emb = out.last_hidden_state[:, 0, :]
        logits = self.classifier(self.dropout(cls_emb))
        return logits


def load_classifier(ckpt_path, device):
    """
    从 train_sentiment.py 保存的 best_model.pt 里恢复：
    - model_state_dict
    - tokenizer_name
    - max_length
    - label2id
    """
    ckpt = torch.load(ckpt_path, map_location=device)
    tokenizer_name = ckpt["tokenizer_name"]
    max_length = ckpt["max_length"]
    label2id = ckpt["label2id"]          # dict: emotion_str -> int
    id2label = {v: k for k, v in label2id.items()}
    num_labels = len(label2id)

    model = SentimentClassifier(
        pretrained_name=tokenizer_name,
        dropout_prob=0.1,
        num_labels=num_labels
    ).to(device)
    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()

    tokenizer = BertTokenizerFast.from_pretrained(tokenizer_name)

    return model, tokenizer, max_length, id2label


@torch.no_grad()
def predict_emotion(text, model, tokenizer, max_length, id2label, device):
    """
    用情感分类模型预测一句话的情感标签（字符串）
    """
    enc = tokenizer(
        text,
        add_special_tokens=True,
        truncation=True,
        max_length=max_length,
        padding="max_length",
        return_attention_mask=True,
        return_tensors="pt"
    )
    input_ids = enc["input_ids"].to(device)
    attention_mask = enc["attention_mask"].to(device)

    logits = model(input_ids=input_ids, attention_mask=attention_mask)
    pred_id = torch.argmax(logits, dim=1).item()
    return id2label[pred_id]


# =========================
#  二、T5 改写模型（emotion+style）
# =========================

def load_t5_generator(model_path, device):
    """
    加载你用 train.py 微调并保存的 T5-small 模型：
    默认路径类似于：
      t5-emotion-style-finetuned-small/final_model
    """
    model = T5ForConditionalGeneration.from_pretrained(model_path).to(device)
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model.eval()
    return model, tokenizer


@torch.no_grad()
def generate_rewrite_t5(t5_model, t5_tokenizer, device, text, emotion, style, max_length=128):
    """
    使用微调好的 T5 模型进行改写：
      input: 原句 + 预测情感 + 选择的 style
      prompt 格式必须和训练时保持一致：
        emotion: {emotion}, rewrite as {style}: {text}
    """
    prompt = f"emotion: {emotion}, rewrite as {style}: {text}"
    inputs = t5_tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        padding=True
    ).to(device)

    output_ids = t5_model.generate(
        input_ids=inputs.input_ids,
        attention_mask=inputs.attention_mask,
        max_length=max_length,
        num_beams=4,
        early_stopping=True
    )
    rewritten = t5_tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return rewritten


# =========================
#  三、GUI 部分
# =========================

class SentimentStyleGUI:
    def __init__(
        self,
        root,
        clf_model,
        clf_tokenizer,
        clf_max_length,
        id2label,
        t5_model,
        t5_tokenizer,
        device
    ):
        self.root = root
        self.clf_model = clf_model
        self.clf_tokenizer = clf_tokenizer
        self.clf_max_length = clf_max_length
        self.id2label = id2label

        self.t5_model = t5_model
        self.t5_tokenizer = t5_tokenizer
        self.device = device

        # 最近一次输入的文本
        self.last_input_text = ""

        # Predicted Emotion 显示用
        self.predicted_emotion_var = tk.StringVar(value="(not predicted yet)")

        # style 单选按钮绑定的变量
        self.style_var = tk.StringVar(value="conversational")

        self._build_layout()

    def _build_layout(self):
        self.root.title("Emotion Prediction + Style Rewrite (T5)")

        # ---------- Input ----------
        input_frame = ttk.LabelFrame(self.root, text="Input")
        input_frame.pack(fill="both", padx=10, pady=5, expand=True)

        self.input_text = scrolledtext.ScrolledText(input_frame, height=4, wrap=tk.WORD)
        self.input_text.pack(fill="both", padx=5, pady=5, expand=True)

        btn_frame = ttk.Frame(input_frame)
        btn_frame.pack(fill="x", padx=5, pady=5)

        # Predict Emotion 按钮
        self.predict_btn = ttk.Button(btn_frame, text="Predict Emotion", command=self.on_predict)
        self.predict_btn.pack(side="right", padx=5)

        # Convert by Style 按钮
        self.convert_btn = ttk.Button(btn_frame, text="Convert by Style (T5)", command=self.on_convert)
        self.convert_btn.pack(side="right", padx=5)

        # ---------- Predicted Emotion 显示 ----------
        emo_frame = ttk.LabelFrame(self.root, text="Predicted Emotion (from BERT)")
        emo_frame.pack(fill="x", padx=10, pady=5)

        emo_label = ttk.Label(emo_frame, textvariable=self.predicted_emotion_var)
        emo_label.pack(anchor="w", padx=5, pady=5)

        # ---------- Style 选择 ----------
        style_frame = ttk.LabelFrame(self.root, text="Rewrite Style")
        style_frame.pack(fill="x", padx=10, pady=5)

        styles = ["conversational", "poetic", "formal", "narrative"]
        for st in styles:
            rb = ttk.Radiobutton(
                style_frame,
                text=st,
                value=st,
                variable=self.style_var
            )
            rb.pack(side="left", padx=5, pady=5)

        # ---------- Output ----------
        output_frame = ttk.LabelFrame(self.root, text="Output (T5 Rewritten Sentence)")
        output_frame.pack(fill="both", padx=10, pady=5, expand=True)

        self.output_text = scrolledtext.ScrolledText(output_frame, height=6, wrap=tk.WORD, state="normal")
        self.output_text.pack(fill="both", padx=5, pady=5, expand=True)

    # ===== 按钮回调 =====

    def on_predict(self):
        """
        用 BERT 情感分类器预测 Input 句子的情感类型，并更新 Predicted Emotion。
        """
        text = self.input_text.get("1.0", tk.END).strip()
        if not text:
            self._update_output("Input is empty. 请先在 Input 区输入一句话。")
            self.predicted_emotion_var.set("(not predicted yet)")
            return

        self.last_input_text = text

        emotion = predict_emotion(
            text,
            model=self.clf_model,
            tokenizer=self.clf_tokenizer,
            max_length=self.clf_max_length,
            id2label=self.id2label,
            device=self.device
        )

        self.predicted_emotion_var.set(emotion)
        self._update_output(
            f"Predicted emotion: {emotion}\n\n"
            f"Please select a style, then click “Convert by Style (T5)” to rewrite."
        )

    def on_convert(self):
        """
        根据当前 Predicted Emotion + 选中的 style，
        调用微调好的 T5 模型进行改写。
        """
        text = self.input_text.get("1.0", tk.END).strip()
        if not text:
            self._update_output("Input is empty. Please enter a sentence in the Input section first.")
            return

        self.last_input_text = text

        emotion = self.predicted_emotion_var.get()
        if emotion in ["", "(not predicted yet)"]:
            self._update_output("Emotion has not yet been predicted. Please click “Predict Emotion” first.")
            return

        style = self.style_var.get()

        rewritten = generate_rewrite_t5(
            t5_model=self.t5_model,
            t5_tokenizer=self.t5_tokenizer,
            device=self.device,
            text=self.last_input_text,
            emotion=emotion,
            style=style,
            max_length=128
        )

        self._update_output(rewritten)

    # ===== 工具函数 =====

    def _update_output(self, text):
        self.output_text.configure(state="normal")
        self.output_text.delete("1.0", tk.END)
        self.output_text.insert(tk.END, text)
        self.output_text.configure(state="normal")


# =========================
#  四、主函数
# =========================

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--clf-model-path", type=str, default="saved_model/best_model.pt",
                        help="情感分类器（BERT）checkpoint 路径")
    parser.add_argument("--t5-model-path", type=str,
                        default="final_model",
                        help="T5 微调模型路径 (train.py 保存的 final_model)")
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 1) 加载情感分类模型
    clf_model, clf_tokenizer, clf_max_length, id2label = load_classifier(args.clf_model_path, device)

    # 2) 加载 T5 改写模型
    t5_model, t5_tokenizer = load_t5_generator(args.t5_model_path, device)

    # 3) 构建 GUI
    root = tk.Tk()
    app = SentimentStyleGUI(
        root,
        clf_model=clf_model,
        clf_tokenizer=clf_tokenizer,
        clf_max_length=clf_max_length,
        id2label=id2label,
        t5_model=t5_model,
        t5_tokenizer=t5_tokenizer,
        device=device
    )
    root.mainloop()


if __name__ == "__main__":
    main()