Hugging Face's logo

cssupport
/
bert-news-class

Text Classification
Transformers
PyTorch
English
bert
fill-mask
Inference Endpoints
Model card Files Community
3
Edit model card

Model Card for Model ID

Based on bert-base-uncased . This model takes in news summary/news headlines/news article and classifies into one of 40 categories (listed below) . Dataset used to traing this model is from Kaggle called News Category Dataset porvided by rishabhmisra.github.io/publications (Please cite this datatset when using this model). Contact us for more info: support@cloudsummary.com.

Below are the output labels:

  • arts = 0, arts & culture =1, black voices = 2, business = 3, college = 4, comedy = 5, crime = 6, culture & arts = 7, education = 8, entertainment = 9,environment = 10 fifty=11, food & drink = 12 ,good news = 13, green = 14, healthy living = 15, home & living = 16, impact = 17, latino voices = 18 , media = 19, money = 20 , parenting = 21 , parents = 22 politics = 23, queer voices = 24, religion = 25, science = 26, sports = 27, style = 28, style & beauty = 29 ,taste = 30 ,tech = 31, the worldpost = 32,travel = 33 u.s. news = 34, weddings = 35, weird news = 36, wellness = 37, women = 38 , world news = 39 , worldpost = 40

Model Details

Model Description

Model Sources

Please refer bert-base-uncased for Model Sources.

How to Get Started with the Model

Use the code below to get started with the model.

from transformers import BertTokenizer, BertForSequenceClassification
import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained ("cssupport/bert-news-class").to(device)

def predict(text):
    id_to_class = {0: 'arts', 1: 'arts & culture', 2: 'black voices', 3: 'business', 4: 'college', 5: 'comedy', 6: 'crime', 7: 'culture & arts', 8: 'education', 9: 'entertainment', 10: 'environment', 11: 'fifty', 12: 'food & drink', 13: 'good news', 14: 'green', 15: 'healthy living', 16: 'home & living', 17: 'impact', 18: 'latino voices', 19: 'media', 20: 'money', 21: 'parenting', 22: 'parents', 23: 'politics', 24: 'queer voices', 25: 'religion', 26: 'science', 27: 'sports', 28: 'style', 29: 'style & beauty', 30: 'taste', 31: 'tech', 32: 'the worldpost', 33: 'travel', 34: 'u.s. news', 35: 'weddings', 36: 'weird news', 37: 'wellness', 38: 'women', 39: 'world news', 40: 'worldpost'}
    # Tokenize the input text
    inputs = tokenizer(text, return_tensors='pt', truncation=True, max_length=512, padding='max_length').to(device)
    with torch.no_grad():
        logits = model(inputs['input_ids'], inputs['attention_mask'])[0]
    # Get the predicted class index
    pred_class_idx = torch.argmax(logits, dim=1).item()
    return id_to_class[pred_class_idx]


text ="The UK’s growing debt burden puts it on shaky ground ahead of upcoming assessments by the three main credit ratings agencies. A downgrade to its credit rating, which is a reflection of a country’s creditworthiness, could raise borrowing costs further still, although the impact may be limited."
predicted_class = predict(text)
print(predicted_class)
#OUTPUT : business

Uses

[More Information Needed]

Direct Use

Could used in application where natural language is to be converted into SQL queries. [More Information Needed]

Out-of-Scope Use

[More Information Needed]

Bias, Risks, and Limitations

[More Information Needed]

Recommendations

Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.

Technical Specifications

Model Architecture and Objective

bert-base-uncased

Compute Infrastructure

Hardware

one P6000 GPU

Software

Pytorch and HuggingFace

Citation

Misra, Rishabh. "News Category Dataset." arXiv preprint arXiv:2209.11429 (2022). Misra, Rishabh and Jigyasa Grover. "Sculpting Data for ML: The first act of Machine Learning." ISBN 9798585463570 (2021). Tandon, Karan. "This LLM is based on BERT (2018) a bidirectional Transformer. cssupport/bert-news-class was finetuned using AdamW with the help of NVIDIA AMP and trained in 45 minutes on one P6000 GPU. This model accepts news summary/news headlines/news article and classifies into one of 40 categories"

Downloads last month
96