streamlit_app.py

import os
import re
from typing import Dict, Tuple

import pandas as pd
import streamlit as st
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer

DEFAULT_PHOBERT_PATH = "outputs/phobert/checkpoint-430"
DEFAULT_GPT_MODEL_ID = "openai/gpt-oss-20b"
DEFAULT_GPT_ADAPTER_PATH = "outputs/gpt20b/checkpoint-430"
DEFAULT_TEST_PATH = "data/splits/test.csv"

PROMPT_TEMPLATE = (
    "Phân loại bài viết sau là clickbait hay không.\n\n"
    "Bài viết: {text}\n\n"
    "Nhãn:"
)

LABEL_TEXT = {0: " không clickbait", 1: " clickbait"}


@st.cache_resource(show_spinner=True)
def load_phobert_model(model_path: str):
    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    model.eval()
    return tokenizer, model


@st.cache_resource(show_spinner=True)
def load_gpt_local_model(base_model_id: str, adapter_path: str):
    tokenizer = AutoTokenizer.from_pretrained(adapter_path, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id
    tokenizer.padding_side = "left"

    if torch.cuda.is_available():
        dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        # Keep the full model on a single GPU to avoid CUDA/CPU tensor mismatch in MoE ops.
        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_id,
            torch_dtype=dtype,
            device_map={"": 0},
        )
    else:
        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_id,
            torch_dtype=torch.float32,
        )

    model = PeftModel.from_pretrained(base_model, adapter_path)
    model.eval()

    for cfg in filter(None, [model.config, getattr(model, "generation_config", None)]):
        cfg.pad_token_id = tokenizer.pad_token_id
        cfg.eos_token_id = tokenizer.eos_token_id
        if tokenizer.bos_token_id is not None:
            cfg.bos_token_id = tokenizer.bos_token_id

    return tokenizer, model


def make_input_text(title: str, lead_paragraph: str) -> str:
    title = (title or "").strip()
    lead_paragraph = (lead_paragraph or "").strip()
    return f"{title} [SEP] {lead_paragraph}".strip()


def predict_phobert(text: str, tokenizer, model) -> Dict[str, float]:
    encoded = tokenizer(
        text,
        truncation=True,
        max_length=256,
        return_tensors="pt",
    )

    with torch.no_grad():
        logits = model(**encoded).logits
        probs = torch.softmax(logits, dim=-1)[0].cpu().tolist()

    return {
        "non-clickbait": float(probs[0]),
        "clickbait": float(probs[1]),
    }


def parse_gpt_label(raw_output: str) -> str:
    normalized = raw_output.strip().lower()
    normalized = re.sub(r"\s+", " ", normalized)

    if "non-clickbait" in normalized or "khong clickbait" in normalized:
        return "non-clickbait"
    if "clickbait" in normalized:
        return "clickbait"
    return "unknown"


def score_candidate(model, input_ids, resp_start: int) -> float:
    with torch.no_grad():
        logits = model(input_ids).logits
    log_probs = torch.nn.functional.log_softmax(logits[0], dim=-1)

    total_lp = 0.0
    n_tokens = input_ids.shape[-1] - resp_start
    full_ids = input_ids[0].tolist()
    for i in range(resp_start, len(full_ids)):
        total_lp += log_probs[i - 1, full_ids[i]].item()
    return total_lp / max(n_tokens, 1)


def normalize_log_scores(scores: Dict[int, float]) -> Dict[str, float]:
    logit_tensor = torch.tensor([scores[0], scores[1]], dtype=torch.float32)
    probs = torch.softmax(logit_tensor, dim=0).tolist()
    return {
        "non-clickbait": float(probs[0]),
        "clickbait": float(probs[1]),
    }


def predict_gpt_local(text: str, tokenizer, model, max_length: int = 256) -> Tuple[str, Dict[str, float], Dict[str, float]]:
    prompt_text = PROMPT_TEMPLATE.format(text=text)
    prompt_ids = tokenizer.encode(prompt_text, add_special_tokens=False)
    device = torch.device("cuda:0") if torch.cuda.is_available() else next(model.parameters()).device
    scores = {}

    for label_id, label_text in LABEL_TEXT.items():
        full_text = prompt_text + label_text
        full_ids = tokenizer.encode(full_text, add_special_tokens=False)

        if len(full_ids) > max_length:
            resp_ids = tokenizer.encode(label_text, add_special_tokens=False)
            p_ids = prompt_ids[: max_length - len(resp_ids)]
            full_ids = p_ids + resp_ids
            resp_start = len(p_ids)
        else:
            resp_start = len(prompt_ids)

        input_ids = torch.tensor([full_ids], device=device)
        scores[label_id] = score_candidate(model, input_ids, resp_start)

    pred_id = max(scores, key=scores.get)
    pred_label = "clickbait" if pred_id == 1 else "non-clickbait"
    raw_score_view = {
        "non-clickbait": float(scores[0]),
        "clickbait": float(scores[1]),
    }
    prob_view = normalize_log_scores(scores)
    return pred_label, raw_score_view, prob_view


def load_demo_samples() -> pd.DataFrame:
    if os.path.exists(DEFAULT_TEST_PATH):
        df = pd.read_csv(DEFAULT_TEST_PATH)
        expected_cols = {"title", "lead_paragraph", "label", "text"}
        if expected_cols.issubset(set(df.columns)):
            return df
    return pd.DataFrame()


def main() -> None:
    st.set_page_config(page_title="ViClickbait Demo", page_icon="📰", layout="wide")
    st.title("ViClickbait-2025 Demo: PhoBERT vs GPT-OSS-20B")
    st.caption("Compare clickbait prediction between local PhoBERT and local GPT-OSS-20B (base model + LoRA adapter).")

    with st.sidebar:
        st.header("Config")
        phobert_path = st.text_input("PhoBERT path", value=DEFAULT_PHOBERT_PATH)
        gpt_model_id = st.text_input("GPT base model id", value=DEFAULT_GPT_MODEL_ID)
        gpt_adapter_path = st.text_input("GPT adapter path", value=DEFAULT_GPT_ADAPTER_PATH)
        hf_token = st.text_input("HF token (optional if base model is gated)", value=os.getenv("HF_TOKEN", ""), type="password")
        st.info("PhoBERT va GPT-OSS-20B deu duoc load local checkpoint. GPT dung base model + LoRA adapter trong outputs/gpt20b.")

    samples = load_demo_samples()

    if "title_val" not in st.session_state:
        st.session_state["title_val"] = ""
    if "lead_val" not in st.session_state:
        st.session_state["lead_val"] = ""

    col_left, col_right = st.columns(2)
    with col_left:
        st.text_area("Title", height=120, key="title_val")
    with col_right:
        st.text_area("Lead paragraph", height=120, key="lead_val")

    if not samples.empty:
        with st.expander("Load sample from test split"):
            sample_idx = st.number_input(
                "Row index",
                min_value=0,
                max_value=int(len(samples) - 1),
                value=0,
                step=1,
            )
            if st.button("Use sample", use_container_width=True):
                row = samples.iloc[int(sample_idx)]
                st.session_state["title_val"] = str(row["title"])
                st.session_state["lead_val"] = str(row["lead_paragraph"])
                st.session_state["label_val"] = str(row["label"])
                st.rerun()

    title = st.session_state.get("title_val", "")
    lead_paragraph = st.session_state.get("lead_val", "")

    article_text = make_input_text(title, lead_paragraph)

    if st.button("Predict", type="primary", use_container_width=True):
        if not article_text:
            st.error("Please enter title and/or lead paragraph.")
            return

        try:
            tokenizer, model = load_phobert_model(phobert_path)
            phobert_scores = predict_phobert(article_text, tokenizer, model)
            phobert_label = max(phobert_scores, key=phobert_scores.get)
        except Exception as exc:
            st.error(f"PhoBERT loading/prediction failed: {exc}")
            return

        gpt_label = "not-run"
        gpt_raw = ""
        gpt_raw_scores = {}
        gpt_probs = {}
        gpt_error = None

        try:
            if hf_token.strip():
                os.environ["HF_TOKEN"] = hf_token.strip()
                os.environ["HUGGINGFACEHUB_API_TOKEN"] = hf_token.strip()
            gpt_tokenizer, gpt_model = load_gpt_local_model(
                gpt_model_id.strip(),
                gpt_adapter_path.strip(),
            )
            gpt_label, gpt_raw_scores, gpt_probs = predict_gpt_local(article_text, gpt_tokenizer, gpt_model)
            gpt_raw = f"non-clickbait={gpt_raw_scores['non-clickbait']:.4f}, clickbait={gpt_raw_scores['clickbait']:.4f}"
        except Exception as exc:
            gpt_error = str(exc)

        left, right = st.columns(2)
        with left:
            st.subheader("PhoBERT")
            st.metric("Predicted label", phobert_label)
            st.progress(min(max(phobert_scores["clickbait"], 0.0), 1.0))
            st.write(
                {
                    "non-clickbait": round(phobert_scores["non-clickbait"], 4),
                    "clickbait": round(phobert_scores["clickbait"], 4),
                }
            )

        with right:
            st.subheader("GPT-OSS-20B")
            st.metric("Predicted label", gpt_label)
            if gpt_probs:
                st.progress(min(max(gpt_probs["clickbait"], 0.0), 1.0))
            if gpt_error:
                st.warning(gpt_error)
            if gpt_probs:
                st.write(
                    {
                        "non-clickbait": round(gpt_probs["non-clickbait"], 4),
                        "clickbait": round(gpt_probs["clickbait"], 4),
                    }
                )
            if gpt_raw:
                st.write("Raw log-score output:")
                st.code(gpt_raw)

        if "label_val" in st.session_state:
            st.divider()
            st.write(f"Ground truth sample label: {st.session_state['label_val']}")


if __name__ == "__main__":
    main()