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()