Model Card for KooBERT

KooBERT is a masked language model trained on data from the multilingual micro-blogging social media platform Koo India. This model was built in collaboration with Koo India and AI4Bharat.

Model Details

Model Description

On Koo platform, we have microblogs (Koos) which are limited to 400 characters and are available in multiple languages. The model was trained on a dataset that contains multilingual koos from Jan 2020 to Nov 2022 on masked language modeling task.

• Model type: BERT based pretrained model
• Language(s) (NLP): assamese, bengali, english, gujarati, hindi, kannada, malayalam, marathi, oriya, punjabi, tamil, telugu
• License: KooBERT released under the MIT License.

Uses

This model can be used to perform downstream tasks like content classification, toxicity detection, etc. for supported Indic languages It can also be used with sentence-transformers library for the creation of multilingual vector embeddings for other uses.

Bias, Risks, and Limitations

As with any machine learning model, KooBERT may have limitations and biases. It is important to keep in mind that this model was trained on Koo Social Media data and may not generalize well to other domains. It is also possible that the model may have biases in the data it was trained on, which may affect its predictions. It is recommended to evaluate the model on your specific use case and data to ensure it is appropriate for your needs.

How to Get Started with the Model

Use the code below to get started with the model for general finetuning tasks. Please note this is just a sample for finetuning.

import torch
from datasets import load_dataset, load_metric
import evaluate
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("Koodsml/KooBERT")
model = AutoModelForSequenceClassification.from_pretrained("Koodsml/KooBERT", num_labels=2)

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

def tokenize_function(examples):
    return tokenizer(examples["text"], padding='max_length', truncation=True, max_length=128)

# Load the CoLA dataset
dataset = load_dataset("glue","cola")
dataset = dataset.rename_column('sentence', 'text')

datset_tok = dataset.map(tokenize_function, batched=True)

# Set the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Define the training arguments
training_args = TrainingArguments(output_dir="test_trainer", evaluation_strategy="epoch")

# Define the trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=datset_tok['train'],
    eval_dataset=datset_tok['validation'],
    compute_metrics=compute_metrics,
)

# Fine-tune on the CoLA dataset
trainer.train()

# Evaluate on the CoLA dataset
eval_results = trainer.evaluate(eval_dataset=cola_dataset['validation'])
print(eval_results)

We can also use KooBERT with the sentence-transformers library to create multilingual vector embeddings. Here is an example:

from sentence_transformers import SentenceTransformer

# Load the KooBERT model
koo_model = SentenceTransformer('Koodsml/KooBERT', device="cuda")

# Define the text
text = "यह हमेशा से हमारी सोच है"

# Get the embedding
embedding = koo_model.encode(text)
print(embedding)

Training Details

Training Data

Following is the distribution of tokens over languages:

Total Koos = 68,143,710 Total Tokens = 966,239,820 (based on a close approximation)

Training Procedure

Preprocessing

Personal Identifiable Information (PII) was removed from data before training on microblogs. Temperature Sampling to upsample low resource languages. We used a temperature of value of 0.7 (Refer Sec 3.1 https://arxiv.org/pdf/1901.07291.pdf)

Training Hyperparameters

Training regime

• Training steps - 1M steps
• Warm - 10k steps
• Learning Rate - 5e-4
• Scheduler - Linear Decay
• Optimizer - Adam
• Batch Size of 4096 sequences
• Precision - fp32

Evaluation

The model has not been benchmarked yet. We shall be releasing the benchmark data in a future update.

Model Examination

Model Architecture and Objective

KooBERT is pretrained with BERT Architecture on Masked Language Modeling with a vocabulary size of 128k and max sequence length of 128 tokens.

Compute Infrastructure

KooBERT was trained on TPU v3 with 128 cores which took over 5 days.

Contributors

Mitesh Khapra (miteshk@cse.iitm.ac.in)- IITM AI4Bharat Sumanth Doddapaneni (dsumanth17@gmail.com)- IITM AI4Bharat Smiral Rashinkar (smiral.rashinkar@kooapp.com)- Koo India