Spaces:
Sleeping
Sleeping
File size: 2,393 Bytes
5bb7e34 75d06a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoConfig
from huggingface_hub import hf_hub_url
import os
# Impor kelas kustom Anda secara eksplisit
from model import IndoBERTClassifier
# --- Konfigurasi dan Pemuatan Model ---
MODEL_ID = "Hydra-RKMI/KlikBERT"
# Muat tokenizer dan config dari Hub
config = AutoConfig.from_pretrained(MODEL_ID)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
# Inisialisasi kelas kustom dan muat bobot dari Hub
model = IndoBERTClassifier(config)
model_path = hf_hub_url(repo_id=MODEL_ID, filename="pytorch_model.bin")
model.load_state_dict(torch.hub.load_state_dict_from_url(model_path, map_location="cpu"))
model.eval()
# --- Pemetaan Label ---
# Pastikan config.json Anda sudah menggunakan 'custom_id2label'
id2label_clickbait = config.custom_id2label['clickbait']
id2label_kategori = config.custom_id2label['kategori']
# --- Fungsi Prediksi ---
def predict(judul, isi):
inputs = tokenizer(
judul,
isi,
truncation=True,
padding=True,
max_length=512,
return_tensors="pt"
)
with torch.no_grad():
outputs = model(**inputs)
clickbait_logits = outputs["clickbait_logits"]
kategori_logits = outputs["kategori_logits"]
pred_clickbait_id = torch.argmax(clickbait_logits, dim=1).item()
pred_kategori_id = torch.argmax(kategori_logits, dim=1).item()
pred_clickbait_label = id2label_clickbait[str(pred_clickbait_id)]
pred_kategori_label = id2label_kategori[str(pred_kategori_id)]
# --- PERUBAHAN DI SINI ---
# Kembalikan dua nilai terpisah, bukan dictionary
return pred_clickbait_label, pred_kategori_label
# --- Antarmuka Gradio ---
inputs = [
gr.Textbox(lines=2, label="Judul Berita", placeholder="Masukkan judul berita di sini..."),
gr.Textbox(lines=10, label="Isi Berita", placeholder="Masukkan isi berita di sini...")
]
# --- PERUBAHAN DI SINI ---
# Gunakan dua komponen output terpisah
outputs = [
gr.Text(label="Prediksi Clickbait"),
gr.Text(label="Prediksi Kategori Berita")
]
title = "Model Multi-Task KlikBERT"
description = "Model ini memprediksi apakah judul clickbait dan apa kategori beritanya. Model ini dimuat dari repositori TrioF/KlikBERT."
iface = gr.Interface(fn=predict, inputs=inputs, outputs=outputs, title=title, description=description)
iface.launch() |