ND0210/data-mining

1import gradio as gr2import torch3 4from transformers import AutoTokenizer, AutoModelForSequenceClassification5import torch.nn.functional as F6 7genres = ['Kinh tế', 'Giáo dục', 'Xe', 'Sức khoẻ', 'Công nghệ - Game']8 9tokenizer = AutoTokenizer.from_pretrained("mob2711/phoBERT_finetune_news_classification")10model = AutoModelForSequenceClassification.from_pretrained("mob2711/phoBERT_finetune_news_classification")11 12def tokenize(text):13    encoded_text = tokenizer(text, padding=True, truncation=True, return_tensors="pt")14    return encoded_text15 16def predict_proba(text_data):17    encoded_data = tokenize(text_data)18    with torch.no_grad():19        outputs = model(**encoded_data)20        logits = outputs.logits21        probabilities = F.softmax(logits, dim=-1)[0]22    label_probs = {genres[id]: prob for id, prob in enumerate(probabilities)}23    return label_probs24    25# Interface26input_text = gr.Textbox(label="Enter the title")27output_text = gr.Label(label="Predicted Probabilities")28 29demo = gr.Interface(30    fn=predict_proba, 31    inputs=input_text, 32    outputs=output_text, 33    title="Newspaper Title Classifier",34    examples=["Chủ tịch HĐQT Trường quốc tế AISVN đề xuất hỗ trợ 125 tỉ đồng", 35              "Chuyên gia tài chính Nguyễn Trí Hiếu bị 'hack' gần 500 triệu đồng, ngân hàng im lặng suốt 3 tháng?",36              "Kon Tum chi hỗ trợ 28.000 liều vắc xin tiêm phòng bệnh dại cho chó, mèo",37              "Microsoft hợp tác OpenAI phát triển siêu máy tính AI giá hơn 100 tỉ USD",38              "Triệu hồi 170.000 xe điện Hyundai, Kia bị lỗi mất điện",39              ]40)41 42demo.launch()