Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import re | |
| from simpletransformers.classification import ClassificationModel | |
| from pythainlp import sent_tokenize | |
| from thai_tokenization import ThaiTokenizer | |
| from transformers import AutoTokenizer | |
| #tokenizer = ThaiTokenizer(vocab_file='th.wiki.bpe.op25000.vocab', spm_file='th.wiki.bpe.op25000.model') | |
| #create tokenizer | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| 'airesearch/wangchanberta-base-att-spm-uncased', | |
| revision='main', | |
| model_max_length=416,) | |
| typeId = {'การเมือง': 0, 'กีฬา': 1, 'คุณภาพชีวิต': 2, 'ทั่วไทย': 3, 'ไลฟ์สไตล์': 4, | |
| 'อื่นๆ': 5, 'อาชญากรรม': 6, 'สิ่งแวดล้อม': 7, 'บันเทิง & วัฒนธรรม': 8, 'เศรษฐกิจ': 9, | |
| 'วิทยาศาสตร์ & การศึกษา': 10, 'สังคม': 11, 'unspecified': 12, 'ต่างประเทศ': 13} | |
| loaded_model = ClassificationModel( | |
| "camembert", | |
| "SuperBigtoo/thainews-classification-wangchanberta", | |
| use_cuda=torch.cuda.is_available(), | |
| num_labels=14, | |
| ) | |
| def predict_type(title_input): | |
| title_input = title_input.lower() | |
| title_input = re.sub(r'[-:‘’“”.,%/()\[\]\'"!;\\/]', '', title_input) | |
| title_input = sent_tokenize(title_input) | |
| title_input = ' '.join(tokenizer.tokenize(' '.join(title_input))) | |
| predictions, raw_outputs = loaded_model.predict([title_input]) | |
| predicted_label_name = [type_name for type_name, type_id in typeId.items() if type_id == predictions[0]] | |
| return predicted_label_name[0] | |
| iface = gr.Interface( | |
| fn=predict_type, | |
| inputs=gr.Textbox(lines=1, max_lines=10, label="Input News's Title"), | |
| outputs=gr.Textbox(lines=1, max_lines=2, label="Predicted News's Type"), | |
| title="Thai News Classify", | |
| examples=['จบสกอร์ไม่คม หมดครึ่งแรก ยูเครน เจ๊า โปแลนด์ 0-0', | |
| 'แอรินยินดีนาฑี มีรักใหม่ ยันจบกันด้วยดี ปัดถ่ายแฟชั่นเซ็กซี่ประชดรัก อ้างถูกใจคอนเซปต์ (คลิป)', | |
| 'แนวโน้ม ราคาทอง ปี 63 ไตรมาสแรกส่อลงแรงก่อนทะยานขึ้นอาจเห็นบาทละ 23750', | |
| 'ปล้นรถ ปล้นปืน เผย 8 ศพกระบี่ เป็นไปได้ทีมฆ่าถูกสั่งตรงจากชายแดนใต้', | |
| 'อินโดฯ จัดกิจกรรมศาสนา ไม่สนโควิด-19 หวั่นซ้ำรอยมาเลเซีย', | |
| 'วงคุยว่าด้วย ชาติที่เรา (จะ) รักของ นิธิผ่านนักวิชาการ 3 รุ่น หลังเลือกตั้ง', | |
| 'ฆ่าโหด "น้องชายจุฬาราชมนตรี" หนุ่มใหญ่หลอนยา ชักมีดฟัน-แทงยับ', | |
| 'กู้เงินออมสิน 10,000-300,000 บาท ผ่านออนไลน์ ดอกเบี้ยต่ำใช่เหรอ' | |
| ] | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() |