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.
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:
Language
Koos
Avg Tokens per Koo
Total Tokens
assamese
562,050
16.4414198
9,240,900
bengali
2,110,380
12.08918773
25,512,780
english
17,889,600
10.93732057
195,664,290
gujarati
1,825,770
14.33965395
26,180,910
hindi
35,948,760
16.2337502
583,583,190
kannada
2,653,860
12.04577107
31,967,790
malayalam
71,370
10.32744851
737,070
marathi
1,894,080
14.81544602
28,061,640
oriya
87,930
14.1941317
1,248,090
punjabi
940,260
18.59961075
17,488,470
tamil
1,687,710
12.12147822
20,457,540
telugu
2,471,940
10.55735576
26,097,150
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.
