Project: s41467-024-44803-4 | Function: compute_loss (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Weighted cross-entropy loss computation for imbalanced chronic vs control classification task

**KEY TERMS:** {
  "labels": "chronic/control labels",
  "logits": "model predictions",
  "weight": "class imbalance weights"
}

**CONTEXT CODE:** import comet_ml
import os
import torch
import argparse
import glob
import numpy as np

import torch.nn as nn
from datasets import load_dataset
from transformers import TrainingArguments, Trainer, BertTokenizer, BertForSequenceClassification, AdamW
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix


# this code was inspired from : https://mccormickml.com/2019/05/14/BERT-word-embeddings-tutorial/
class WeightedTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        # MASKED: Weighted cross-entropy loss computation for imbalanced chronic vs control classification implementation
        raise NotImplementedError("This function is masked for testing purposes.")

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    acc = accuracy_score(labels, preds)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='weighted')
    cm = confusion_matrix(labels, preds)
    return {
        'accuracy': acc,
        'precision': precision,
        'recall': recall,
        'f1': f1,
        'confusion_matrix': cm.tolist()
    }

class CustomTokenizer(BertTokenizer):
    def _tokenize(self, text):
        # replace : and - used by non-syn mutations and
        text = text.replace(':', 'X').replace('-','X')
        # tokenize by white space
        tokens = text.split()
        return tokens

tokenizer = CustomTokenizer.from_pretrained('/sternadi/home/volume3/chronic-corona-pred/'
                                            'sars_cov_2_mlm/models/pretrained-covBERTa/', do_lower_case=False)
def _encode(x):
    return tokenizer(x["text"], truncation=True, padding="max_length",
                     max_length=160, return_special_tokens_mask=True)


def main(args):

    experiment = comet_ml.Experiment(
        api_key='r0TcQHQkcqR8wVtCye8tTSAJj',
        project_name="cov-transformers-trainer",
    )
    experiment.add_tags(args.tags.split(','))

    # load dataset train and test splits
    files_map = {'train': glob.glob(f"{args.dataset}/fold_{args.fold}_train.csv"),
                 'test': glob.glob(f"{args.dataset}/fold_{args.fold}_test.csv")}
    if args.debug:
        files_map['train'], files_map['test'] = files_map['train'][:10], files_map['test'][:4]

    dataset = load_dataset("csv", data_files=files_map,
                           cache_dir='/sternadi/home/volume3/chronic-corona-pred/sars_cov_2_mlm/HF_DATASETS_CACHE_CLS')

    # tokenizing train & test dataset
    ncpus = args.ncpu
    train_dataset = dataset["train"].map(_encode, batched=True, num_proc=ncpus)
    test_dataset = dataset["test"].map(_encode, batched=True, num_proc=ncpus)

    # remove other columns and set input_ids and attention_mask as tensors
    train_dataset.set_format(type="torch", columns=["input_ids", "label", "attention_mask"])
    test_dataset.set_format(type="torch", columns=["input_ids", "label", "attention_mask"])

    print(f"Loaded {len(train_dataset)} train samples, {len(test_dataset)} test samples")
    parameters = {'per_device_train_batch_size': 10, 'per_device_eval_batch_size':64,
                  'num_train_epochs': args.epoches}
    experiment.log_parameters(parameters)

    # configure model output path
    model_path = os.path.join(args.model_path, f'{args.fold}')
    os.makedirs(model_path, exist_ok=True)


    # initialize the model with the config
    model = BertForSequenceClassification.from_pretrained(args.model)
    # Freeze the lower layers - do not want to overfit....
    for name, param in model.named_parameters():
        if name.startswith('bert.encoder') or name.startswith('bert.embeddings'):
            param.requires_grad = False


    training_args = TrainingArguments(
        output_dir=model_path,
        evaluation_strategy="steps",
        overwrite_output_dir=True,
        num_train_epochs=args.epoches,
        per_device_train_batch_size=6,
        per_device_eval_batch_size=64,
        logging_steps=args.logging_interval,
        save_steps=args.save_interval,
        load_best_model_at_end=True,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=test_dataset,
        compute_metrics=compute_metrics
    )

    # train the model
    trainer.train()



Complete the function implementation replacing "raise NotImplementedError".