Views
No views yet
1from transformers import AutoTokenizer, BertForSequenceClassification
2
3model_name = "AyoubChLin/bert_cnn_news"
4tokenizer = AutoTokenizer.from_pretrained(model_name)
5model = TFAutoModelForSequenceClassification.from_pretrained(model_name)
6
7text = "This is a news article about politics."
8inputs = tokenizer(text, padding=True, truncation=True, return_tensors="tf")
9
10outputs = model(inputs)
11predicted_class_id = tf.argmax(outputs.logits, axis=-1).numpy()[0]
12
13labels = ["business", "entertainment", "health", "news", "politics", "sport"]
14predicted_label = labels[predicted_class_id]from_pretrained methods. We then encode a news article using the tokenizer, pass the inputs through the model, and extract the predicted label using the argmax function. Finally, we map the predicted label to its corresponding category using a list of labels.