A Practical Guide to Multilingual Large Language Model (RoBERTa) Classification

try:
  with open("./train.json") as f:
    train = json.load(f)

  with open("./val.json") as f:
    val = json.load(f)

  with open("./test.json") as f:
    test = json.load(f)
except:
  #Dummy Values
  train = {'Text': ["An Apple is a Fruit!", "An Apple is not a Fruit!", "An Apple has no seeds."]*64,
          'Labels': ['is_correct', 'is_incorrect', 'is_incorrect']*64}
  val = {'Text': ["An Apple is not a Fruit!", "An Apple is a Fruit!", "An Apple has no seeds."]*64,
          'Labels': ['is_incorrect', 'is_correct', 'is_incorrect']*64}
  test = {'Text': ["An Apple is not a Fruit!", "An Apple is has no seeds.", "An Apple is a Fruit."]*64,
          'Labels': ['is_incorrect', 'is_incorrect', 'is_correct']*64}


assert len(train['Text']) == len(train['Labels']), "Number of texts does not match number of labels for train data!"
assert len(val['Text']) == len(val['Labels']), "Number of texts does not match number of labels for val data!"
assert len(test['Text']) == len(test['Labels']), "Number of texts does not match number of labels for test data!"

print(f"We loaded the train data with {len(train['Text'])} texts and {len(train['Labels'])} labels,")
print(f"the validation data with {len(val['Text'])} texts and {len(val['Labels'])} labels")
print(f"and the test data with {len(test['Text'])} texts and {len(test['Labels'])} labels.")

We loaded the train data with 192 texts and 192 labels,
the validation data with 192 texts and 192 labels
and the test data with 192 texts and 192 labels.

model_config = {'pretrained_model_name_or_path': model_name,
                'num_labels': output_size,
                'problem_type': objective,
                'id2label': id2label,
                'label2id': label2id}

model = AutoModelForSequenceClassification.from_pretrained(**model_config)
tokenizer = AutoTokenizer.from_pretrained(model_name)

Some weights of RobertaForSequenceClassification were not initialized from the model checkpoint at roberta-base and are newly initialized: ['classifier.dense.bias', 'classifier.dense.weight', 'classifier.out_proj.bias', 'classifier.out_proj.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

# Memory and system check before training
import psutil
import gc
# Trainings-loop with memory management fixes
import os

# Set tokenizers parallelism to avoid fork warnings
os.environ["TOKENIZERS_PARALLELISM"] = "false"

device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

# Check system memory
memory = psutil.virtual_memory()
print(f"System Memory: {memory.total / 1024**3:.2f} GB total, {memory.available / 1024**3:.2f} GB available")

# Check GPU memory if available
if torch.cuda.is_available():
    gpu_memory = torch.cuda.get_device_properties(0).total_memory
    gpu_memory_allocated = torch.cuda.memory_allocated(0)
    gpu_memory_reserved = torch.cuda.memory_reserved(0)
    
    print(f"GPU Memory: {gpu_memory / 1024**3:.2f} GB total")
    print(f"GPU Memory Allocated: {gpu_memory_allocated / 1024**3:.2f} GB")
    print(f"GPU Memory Reserved: {gpu_memory_reserved / 1024**3:.2f} GB")
    
    # Clear any existing GPU cache
    torch.cuda.empty_cache()
    gc.collect()
    
    print(f"Device: {device}")
else:
    print("No GPU available - using CPU")

# Check model size
model_size = sum(p.numel() for p in model.parameters())
print(f"Model has {model_size:,} parameters")

# Recommend batch size based on available memory
if torch.cuda.is_available():
    available_gpu_memory = gpu_memory - gpu_memory_reserved
    if available_gpu_memory < 4 * 1024**3:  # Less than 4GB available
        recommended_batch_size = 8
    elif available_gpu_memory < 8 * 1024**3:  # Less than 8GB available
        recommended_batch_size = 16
    else:
        recommended_batch_size = 32
    
    print(f"Recommended batch size: {recommended_batch_size}")
    
    if batch_size > recommended_batch_size:
        print(f"Warning: Current batch size ({batch_size}) may be too large. Consider reducing to {recommended_batch_size}")

System Memory: 7.76 GB total, 5.06 GB available
No GPU available - using CPU
Model has 124,647,170 parameters

best_loss = float('inf')
best_epoch = 0
already_trained = 0
best_model_path = ''
should_delete = True

# Move model to device
model.to(device)

for epoch in range(num_epochs): # Repeat num_epochs times
    model.train()
    
    for batch_idx, batch in enumerate(train_dataloader): # Train the model on the batch
        try:
            input_ids, attention_mask, y = batch[0].to(device), batch[1].to(device), batch[2].to(device)

            # First forward pass
            output = model(input_ids, attention_mask, labels=y)
            loss = output.loss
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

            # SAM optimizer first step
            optimizer.first_step(zero_grad=True)

            # Second forward pass (required by SAM)
            output2 = model(input_ids, attention_mask, labels=y)
            loss2 = output2.loss
            loss2.backward()
            
            # SAM optimizer second step
            optimizer.second_step(zero_grad=True)
            
            # Update learning rate AFTER optimizer steps
            scheduler.step()
            
            print(f"Train: Epoch {epoch}, Train step {already_trained+batch_idx}, Loss {loss.item():.4f}, learning_rate {scheduler.get_last_lr()[0]:.2e}", flush=True)
            
            # Clear cache periodically to prevent memory buildup
            if batch_idx % 5 == 0 and torch.cuda.is_available():
                torch.cuda.empty_cache()
                
        except RuntimeError as e:
            if "out of memory" in str(e).lower():
                print(f"OOM Error at batch {batch_idx}. Trying to recover...")
                torch.cuda.empty_cache()
                gc.collect()
                
                # Try with smaller effective batch size
                if batch_size > 8:
                    batch_size = batch_size // 2
                    print(f"Reducing batch size to {batch_size}")
                    train_dataloader = generate_dataloader(train_text, train_y, batch_size)
                    val_dataloader = generate_dataloader(val_text, val_y, batch_size)
                    break
                else:
                    raise e
            else:
                raise e

    already_trained += batch_idx
    
    # Validation phase
    model.eval()
    val_loss = []
    
    with torch.no_grad():
        for batch_idx, batch in enumerate(val_dataloader): # Validate the current state of the model on the validation data
            input_ids, attention_mask, y = batch[0].to(device), batch[1].to(device), batch[2].to(device)
            val_output = model(input_ids, attention_mask, labels=y)
            val_loss.append(val_output.loss)

    val_loss = torch.mean(torch.stack(val_loss))

    print(f"Validation: Epoch {epoch}, Train step {already_trained}, Loss {val_loss.item():.4f}, old best/epoch {str(best_loss)[1:6]}/{best_epoch}", flush=True)

    if val_loss < best_loss: # Save the model if the val_loss is the best loss we have seen so far
        best_loss = val_loss.item()
        best_epoch = epoch
        if should_delete and best_model_path and os.path.exists(best_model_path):
            shutil.rmtree(best_model_path)
        best_model_path = f"./my_model_epoch_{best_epoch}_val_loss_{str(val_loss.item())[1:6]}"
        model.save_pretrained(best_model_path, from_pt=True)

        print(f"**** END EPOCH {epoch} ****")
    
    # Clean up memory after each epoch
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()

print(f"**** FINISHED TRAINING FOR N={num_epochs} ****")
print(f"BEST EPOCH: {best_epoch}")
print(f"BEST LOSS: {best_loss}")

Train: Epoch 0, Train step 0, Loss 0.7092, learning_rate 8.43e-05

/srv/conda/envs/notebook/lib/python3.10/site-packages/torch/optim/lr_scheduler.py:192: UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
  warnings.warn(

Train: Epoch 0, Train step 1, Loss 0.7057, learning_rate 3.02e-05
Train: Epoch 0, Train step 2, Loss 0.6496, learning_rate 0.00e+00
Validation: Epoch 0, Train step 2, Loss 0.5889, old best/epoch nf/0
**** END EPOCH 0 ****
**** FINISHED TRAINING FOR N=1 ****
BEST EPOCH: 0
BEST LOSS: 0.5889086127281189

best_model = AutoModelForSequenceClassification.from_pretrained(best_model_path)

y_pred = []

with torch.no_grad():

  for batch_idx, batch in enumerate(test_dataloader):
    input_ids, attention_mask, y = batch[0].to(device), batch[1].to(device), batch[2].to(device)
    y_pred.append(model(input_ids, attention_mask, labels=y).logits)


y_pred = torch.cat(y_pred, dim=0)
y_pred = decision_function(y_pred)

from sklearn.metrics import classification_report

print(classification_report(test_y, y_pred, target_names=label2id.keys(), zero_division=True))

              precision    recall  f1-score   support

is_incorrect       0.67      1.00      0.80       128
  is_correct       1.00      0.00      0.00        64

    accuracy                           0.67       192
   macro avg       0.83      0.50      0.40       192
weighted avg       0.78      0.67      0.53       192