|
import streamlit as st |
|
from transformers import TFAutoModelForSequenceClassification, AutoTokenizer |
|
import tensorflow as tf |
|
import numpy as np |
|
|
|
def convert_label_to_title(label): |
|
convert_dict = { |
|
0: "SỨC KHỎE", |
|
1: "GIÁO DỤC", |
|
2: "THỂ THAO", |
|
3: "PHÁP LUẬT", |
|
4: "KHOA HỌC", |
|
5: "DU LỊCH", |
|
6: "GIẢI TRÍ", |
|
7: "KINH DOANH" |
|
} |
|
return convert_dict[label] |
|
|
|
def predict_sentence(model, tokenizer, sentence): |
|
input_data = tokenizer(sentence, return_tensors='tf', padding=True, truncation=True) |
|
logits = model(input_data['input_ids'], attention_mask=input_data['attention_mask']).logits |
|
probabilities = tf.nn.softmax(logits, axis=1) |
|
predicted_class = tf.argmax(logits, axis=1).numpy()[0] |
|
highest_probability = probabilities.numpy()[0, predicted_class] |
|
title = convert_label_to_title(predicted_class) |
|
return title, probabilities.numpy(), highest_probability |
|
|
|
@st.cache_resource |
|
def load_model(checkpoint, num_class): |
|
model = TFAutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=num_class) |
|
tokenizer = AutoTokenizer.from_pretrained(checkpoint) |
|
return model, tokenizer |
|
|
|
checkpoint = 'distilbert-base-multilingual-cased' |
|
model, tokenizer = load_model(checkpoint, 8) |
|
model.load_weights('best_model_weights.h5') |
|
|
|
text = st.text_area('Nhập tiêu đề vào đây') |
|
|
|
if text: |
|
title, probabilities, highest = predict_sentence(model, tokenizer, text) |
|
out = { |
|
'title': title, |
|
'prob': highest |
|
} |
|
st.json(out) |
|
|
|
|
|
|
