Edit model card

Traditional Chinese news classification


from transformers import BertTokenizer, AlbertForSequenceClassification
model_path = "clhuang/albert-news-classification"
model = AlbertForSequenceClassification.from_pretrained(model_path)
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
# Category index
idx2cate = { i : item for i, item in enumerate(news_categories)}

# get category probability
def get_category_proba( text ):
    max_length = 250
    # prepare token sequence
    inputs = tokenizer([text], padding=True, truncation=True, max_length=max_length, return_tensors="pt")
    # perform inference
    outputs = model(**inputs)
    # get output probabilities by doing softmax
    probs = outputs[0].softmax(1)

    # executing argmax function to get the candidate label index
    label_index = probs.argmax(dim=1)[0].tolist() # convert tensor to int
    # get the label name        
    label = idx2cate[ label_index ]

    # get the label probability
    proba = round(float(probs.tolist()[0][label_index]),2)

    response = {'label': label, 'proba': proba}

    return response
{'label': '國際', 'proba': 0.99}
Downloads last month
Hosted inference API
This model can be loaded on the Inference API on-demand.