# embedding_metadata

key: chroma:document
string_value:
CODE:

def predict(input_sentence, model, classes, cached=False, device="cpu"):
    input_t = model.tokenizer.encode(input_sentence)
    input_t = torch.tensor([input_t], dtype=torch.long, device=device)
    if cached:
        input_t = model.avg_representation(input_t)

    log_probs = model(input_t).data.cpu().numpy().flatten().tolist()
    print("Input sentence:", input_sentence)
    print(
        "Predictions:",
        ", ".join("{}: {:.4f}".format(c, math.exp(log_prob)) for c, log_prob in zip(classes, log_probs)),


SUMMARY:
**Summary**

This function implements sentiment prediction using a pre-trained model. It takes an input sentence, a trained model, predicted classes, and optional caching parameters.

### Main Purpose and Functionality
* Tokenizes the input sentence into numerical representation
* Passes the tokenized input through the model to generate predicted probabilities of each class
* Prints the predicted probabilities and corresponding sentiment values

### Key Implementation Details
* Uses `torch.tensor` to create a batched tensor from the tokenized input
* Applies averaging pooling for cached representations, if enabled
* Generates predicted log-probabilities using the model output

### Notable Dependencies or Requirements
* `model`: Pre-trained sentiment analysis model (e.g., BERT)
* `device`: Device for computation (default: "cpu")
* `classes`: List of predicted sentiment classes


key: chroma:document
string_value:
CODE:

def get_cached_data_loader(dataset, batch_size, discriminator, shuffle=False, device="cpu"):
    data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_size, collate_fn=collate_fn)

    xs = []
    ys = []
    for batch_idx, (x, y) in enumerate(tqdm(data_loader, ascii=True)):
        with torch.no_grad():
            x = x.to(device)
            avg_rep = discriminator.avg_representation(x).cpu().detach()
            avg_rep_list = torch.unbind(avg_rep.unsqueeze(1))
            xs += avg_rep_list
            ys += y.cpu().numpy().tolist()

    data_loader = torch.utils.data.DataLoader(
        dataset=Dataset(xs, ys), batch_size=batch_size, shuffle=shuffle, collate_fn=cached_collate_fn
    )



SUMMARY:
**Summary**

This function generates a PyTorch `DataLoader` that caches and uses the average representation of each input sample as an additional feature. The `get_cached_data_loader` takes in:

* A dataset `dataset`
* Batch size `batch_size`
* Discriminator model `discriminator`
* Optional parameters: shuffle, device

It processes the dataset in two steps:

1. It iterates through the original dataset, extracting the input sample and its corresponding output label.
2. After processing each batch, it reconstructs the data loader with modified collate function that uses the cached average representation as an additional feature.

**Key Implementation Details**

* Uses `torch.utils.data.DataLoader` for iterating through the dataset
* Custom collider function `cached_collate_fn`
* Average representation extracted using `discriminator.avg_representation`

**Notable Dependencies and Requirements**

* PyTorch (`torch`)
* Tornado (`tqdm`)
[Optional]
 * Dataset class, which is not shown in this code snippet


key: chroma:document
string_value:
CODE:

def train_discriminator(
    dataset,
    dataset_fp=None,
    pretrained_model="openai-community/gpt2-medium",
    epochs=10,
    batch_size=64,
    log_interval=10,
    save_model=False,
    cached=False,
    no_cuda=False,
):
    device = "cuda" if torch.cuda.is_available() and not no_cuda else "cpu"

    print("Preprocessing {} dataset...".format(dataset))
    start = time.time()

    if dataset == "SST":
        idx2class = ["positive", "negative", "very positive", "very negative", "neutral"]
        class2idx = {c: i for i, c in enumerate(idx2class)}

        discriminator = Discriminator(
            class_size=len(idx2class), pretrained_model=pretrained_model, cached_mode=cached, device=device
        ).to(device)

        text = torchtext_data.Field()
        label = torchtext_data.Field(sequential=False)
        train_data, val_data, test_data = datasets.SST.splits(
            text,
            label,
            fine_grained=True,
            train_subtrees=True,
        )

        x = []
        y = []
        for i in trange(len(train_data), ascii=True):
            seq = TreebankWordDetokenizer().detokenize(vars(train_data[i])["text"])
            seq = discriminator.tokenizer.encode(seq)
            seq = torch.tensor([50256] + seq, device=device, dtype=torch.long)
            x.append(seq)
            y.append(class2idx[vars(train_data[i])["label"]])
        train_dataset = Dataset(x, y)

        test_x = []
        test_y = []
        for i in trange(len(test_data), ascii=True):
            seq = TreebankWordDetokenizer().detokenize(vars(test_data[i])["text"])
            seq = discriminator.tokenizer.encode(seq)
            seq = torch.tensor([50256] + seq, device=device, dtype=torch.long)
            test_x.append(seq)
            test_y.append(class2idx[vars(test_data[i])["label"]])
        test_dataset = Dataset(test_x, test_y)

        discriminator_meta = {
            "class_size": len(idx2class),
            "embed_size": discriminator.embed_size,
            "pretrained_model": pretrained_model,
            "class_vocab": class2idx,
            "default_class": 2,
        }

    elif dataset == "clickbait":
        idx2class = ["non_clickbait", "clickbait"]
        class2idx = {c: i for i, c in enumerate(idx2class)}

        discriminator = Discriminator(
            class_size=len(idx2class), pretrained_model=pretrained_model, cached_mode=cached, device=device
        ).to(device)

        with open("datasets/clickbait/clickbait_train_prefix.txt") as f:
            data = []
            for i, line in enumerate(f):
                try:
                    data.append(eval(line))
                except Exception:
                    print("Error evaluating line {}: {}".format(i, line))
                    continue
        x = []
        y = []
        with open("datasets/clickbait/clickbait_train_prefix.txt") as f:
            for i, line in enumerate(tqdm(f, ascii=True)):
                try:
                    d = eval(line)
                    seq = discriminator.tokenizer.encode(d["text"])

                    if len(seq) < max_length_seq:
                        seq = torch.tensor([50256] + seq, device=device, dtype=torch.long)
                    else:
                        print("Line {} is longer than maximum length {}".format(i, max_length_seq))
                        continue
                    x.append(seq)
                    y.append(d["label"])
                except Exception:
                    print("Error evaluating / tokenizing line {}, skipping it".format(i))
                    pass

        full_dataset = Dataset(x, y)
        train_size = int(0.9 * len(full_dataset))
        test_size = len(full_dataset) - train_size
        train_dataset, test_dataset = torch.utils.data.random_split(full_dataset, [train_size, test_size])

        discriminator_meta = {
            "class_size": len(idx2class),
            "embed_size": discriminator.embed_size,
            "pretrained_model": pretrained_model,
            "class_vocab": class2idx,
            "default_class": 1,
        }

    elif dataset == "toxic":
        idx2class = ["non_toxic", "toxic"]
        class2idx = {c: i for i, c in enumerate(idx2class)}

        discriminator = Discriminator(
            class_size=len(idx2class), pretrained_model=pretrained_model, cached_mode=cached, device=device
        ).to(device)

        x = []
        y = []
        with open("datasets/toxic/toxic_train.txt") as f:
            for i, line in enumerate(tqdm(f, ascii=True)):
                try:
                    d = eval(line)
                    seq = discriminator.tokenizer.encode(d["text"])

                    if len(seq) < max_length_seq:
                        seq = torch.tensor([50256] + seq, device=device, dtype=torch.long)
                    else:
                        print("Line {} is longer than maximum length {}".format(i, max_length_seq))
                        continue
                    x.append(seq)
                    y.append(int(np.sum(d["label"]) > 0))
                except Exception:
                    print("Error evaluating / tokenizing line {}, skipping it".format(i))
                    pass

        full_dataset = Dataset(x, y)
        train_size = int(0.9 * len(full_dataset))
        test_size = len(full_dataset) - train_size
        train_dataset, test_dataset = torch.utils.data.random_split(full_dataset, [train_size, test_size])

        discriminator_meta = {
            "class_size": len(idx2class),
            "embed_size": discriminator.embed_size,
            "pretrained_model": pretrained_model,
            "class_vocab": class2idx,
            "default_class": 0,
        }

    else:  # if dataset == "generic":
        # This assumes the input dataset is a TSV with the following structure:
        # class \t text

        if dataset_fp is None:
            raise ValueError("When generic dataset is selected, dataset_fp needs to be specified aswell.")

        classes = set()
        with open(dataset_fp) as f:
            csv_reader = csv.reader(f, delimiter="\t")
            for row in tqdm(csv_reader, ascii=True):
                if row:
                    classes.add(row[0])

        idx2class = sorted(classes)
        class2idx = {c: i for i, c in enumerate(idx2class)}

        discriminator = Discriminator(
            class_size=len(idx2class), pretrained_model=pretrained_model, cached_mode=cached, device=device
        ).to(device)

        x = []
        y = []
        with open(dataset_fp) as f:
            csv_reader = csv.reader(f, delimiter="\t")
            for i, row in enumerate(tqdm(csv_reader, ascii=True)):
                if row:
                    label = row[0]
                    text = row[1]

                    try:
                        seq = discriminator.tokenizer.encode(text)
                        if len(seq) < max_length_seq:
                            seq = torch.tensor([50256] + seq, device=device, dtype=torch.long)

                        else:
                            print("Line {} is longer than maximum length {}".format(i, max_length_seq))
                            continue

                        x.append(seq)
                        y.append(class2idx[label])

                    except Exception:
                        print("Error tokenizing line {}, skipping it".format(i))
                        pass

        full_dataset = Dataset(x, y)
        train_size = int(0.9 * len(full_dataset))
        test_size = len(full_dataset) - train_size
        train_dataset, test_dataset = torch.utils.data.random_split(full_dataset, [train_size, test_size])

        discriminator_meta = {
            "class_size": len(idx2class),
            "embed_size": discriminator.embed_size,
            "pretrained_model": pretrained_model,
            "class_vocab": class2idx,
            "default_class": 0,
        }

    end = time.time()
    print("Preprocessed {} data points".format(len(train_dataset) + len(test_dataset)))
    print("Data preprocessing took: {:.3f}s".format(end - start))

    if cached:
        print("Building representation cache...")

        start = time.time()

        train_loader = get_cached_data_loader(train_dataset, batch_size, discriminator, shuffle=True, device=device)

        test_loader = get_cached_data_loader(test_dataset, batch_size, discriminator, device=device)

        end = time.time()
        print("Building representation cache took: {:.3f}s".format(end - start))

    else:
        train_loader = torch.utils.data.DataLoader(
            dataset=train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn
        )
        test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, collate_fn=collate_fn)

    if save_model:
        with open("{}_classifier_head_meta.json".format(dataset), "w") as meta_file:
            json.dump(discriminator_meta, meta_file)

    optimizer = optim.Adam(discriminator.parameters(), lr=0.0001)

    for epoch in range(epochs):
        start = time.time()
        print("\nEpoch", epoch + 1)

        train_epoch(
            discriminator=discriminator,
            data_loader=train_loader,
            optimizer=optimizer,
            epoch=epoch,
            log_interval=log_interval,
            device=device,
        )
        evaluate_performance(data_loader=test_loader, discriminator=discriminator, device=device)

        end = time.time()
        print("Epoch took: {:.3f}s".format(end - start))

        print("\nExample prediction")
        predict(example_sentence, discriminator, idx2class, cached=cached, device=device)

        if save_model:
            # torch.save(discriminator.state_dict(),
            #           "{}_discriminator_{}.pt".format(
            #               args.dataset, epoch + 1
            #               ))
            torch.save(
                discriminator.get_classifier().state_dict(),
                "{}_classifier_head_epoch_{}.pt".format(dataset, epoch + 1),


SUMMARY:
Here is a concise summary of the provided code snippet:

**Overview**

The code preprocesses and trains a classifier using a selected dataset. The classifier is a neural network-based model that predicts a class label for a given text input.

**Preprocessing**

1. Load the selected dataset (TSV or TSV with header).
2. Preprocess the data by tokenizing the text and padding sequences to a maximum length.
3. Create a set of unique classes in the dataset.
4. Build an index mapping each class label to an integer ID.

**Model Selection**

The code supports four datasets:

1. Generic (TSV): extracts classes from the header row.
2. TSV with default labels: assumes the first column contains class names.
3. TSV with text: assumes the second column contains text inputs.
4. Custom dataset (TSV or CSV).

**Training**

1. Initialize a classifier architecture (e.g., transformer-based).
2. Load pre-trained models if caching is enabled.
3. Define an optimizer and its hyperparameters.
4. Train the model on the prepared dataset for a specified number of epochs.

**Evaluation**

1. Create test loaders to evaluate the model's performance.
2. Evaluate the model's accuracy using metrics such as macro-F1-score.

**Saving Model**

If `save_model` is enabled, save the classifier model with its state dictionary at each epoch.

Note that some parts of the code are not shown (e.g., `train_epoch`, `predict`, and `evaluate_performance`) because they are assumed to be provided elsewhere in the project.


key: chroma:document
string_value:
CODE:

def top_k_filter(logits, k, probs=False):
    """
    Masks everything but the k top entries as -infinity (1e10).
    Used to mask logits such that e^-infinity -> 0 won't contribute to the
    sum of the denominator.
    """
    if k == 0:
        return logits
    else:
        values = torch.topk(logits, k)[0]
        batch_mins = values[:, -1].view(-1, 1).expand_as(logits)
        if probs:
            return torch.where(logits < batch_mins, torch.ones_like(logits) * 0.0, logits)


SUMMARY:
**Summary**

The `top_k_filter` function masks input logits to the top `k` values by replacing them with `-infinity (1e10)` and sets all other values to 0. It is used for computational efficiency in calculations involving logits, such as calculating probabilities.

**Key Implementation Details**

* Mismatches logits with top `k` values using `torch.topk`.
* Appends a batch minimum value to each logit to create a mask.
* Replaces logit values below the mask with `-infinity (1e10)` or 0 (depending on input parameter).

**Notable Dependencies and Requirements**

* Requires PyTorch library.


key: chroma:document
string_value:
CODE:

def perturb_past(
    past,
    model,
    last,
    unpert_past=None,
    unpert_logits=None,
    accumulated_hidden=None,
    grad_norms=None,
    stepsize=0.01,
    one_hot_bows_vectors=None,
    classifier=None,
    class_label=None,
    loss_type=0,
    num_iterations=3,
    horizon_length=1,
    window_length=0,
    decay=False,
    gamma=1.5,
    kl_scale=0.01,
    device="cuda",
):
    # Generate inital perturbed past
    grad_accumulator = [(np.zeros(p.shape).astype("float32")) for p in past]

    if accumulated_hidden is None:
        accumulated_hidden = 0

    if decay:
        decay_mask = torch.arange(0.0, 1.0 + SMALL_CONST, 1.0 / (window_length))[1:]
    else:
        decay_mask = 1.0

    # TODO fix this comment (SUMANTH)
    # Generate a mask is gradient perturbated is based on a past window
    _, _, _, curr_length, _ = past[0].shape

    if curr_length > window_length and window_length > 0:
        ones_key_val_shape = tuple(past[0].shape[:-2]) + (window_length,) + tuple(past[0].shape[-1:])

        zeros_key_val_shape = tuple(past[0].shape[:-2]) + (curr_length - window_length,) + tuple(past[0].shape[-1:])

        ones_mask = torch.ones(ones_key_val_shape)
        ones_mask = decay_mask * ones_mask.permute(0, 1, 2, 4, 3)
        ones_mask = ones_mask.permute(0, 1, 2, 4, 3)

        window_mask = torch.cat((ones_mask, torch.zeros(zeros_key_val_shape)), dim=-2).to(device)
    else:
        window_mask = torch.ones_like(past[0]).to(device)

    # accumulate perturbations for num_iterations
    loss_per_iter = []
    new_accumulated_hidden = None
    for i in range(num_iterations):
        print("Iteration ", i + 1)
        curr_perturbation = [torch.from_numpy(p_).requires_grad_(True).to(device=device) for p_ in grad_accumulator]
        # make sure p_.grad is not None
        for p_ in curr_perturbation:
            p_.retain_grad()

        # Compute hidden using perturbed past
        perturbed_past = list(map(add, past, curr_perturbation))
        _, _, _, curr_length, _ = curr_perturbation[0].shape
        lm_output = model(last, past_key_values=perturbed_past)
        all_logits, all_hidden = lm_output["logits"], lm_output["hidden_states"]
        hidden = all_hidden[-1]
        new_accumulated_hidden = accumulated_hidden + torch.sum(hidden, dim=1).detach()
        # TODO: Check the layer-norm consistency of this with trained discriminator (Sumanth)
        logits = all_logits[:, -1, :]
        probs = nn.functional.softmax(logits, dim=-1)

        loss = 0.0
        loss_list = []
        if loss_type == PPLM_BOW or loss_type == PPLM_BOW_DISCRIM:
            for one_hot_bow in one_hot_bows_vectors:
                bow_logits = torch.mm(probs, torch.t(one_hot_bow))
                bow_loss = -torch.log(torch.sum(bow_logits))
                loss += bow_loss
                loss_list.append(bow_loss)
            print(" pplm_bow_loss:", loss.data.cpu().numpy())

        if loss_type == 2 or loss_type == 3:
            ce_loss = nn.CrossEntropyLoss()
            # TODO why we need to do this assignment and not just using unpert_past? (Sumanth)
            curr_unpert_past = unpert_past
            curr_probs = torch.unsqueeze(probs, dim=1)
            wte = model.resize_token_embeddings()
            for _ in range(horizon_length):
                inputs_embeds = torch.matmul(curr_probs, wte.weight.data)
                lm_output = model(past_key_values=curr_unpert_past, inputs_embeds=inputs_embeds)
                curr_all_logits, curr_unpert_past, curr_all_hidden = (
                    lm_output["logits"],
                    lm_output["past_key_values"],
                    lm_output["hidden_states"],
                )
                curr_logits = curr_all_logits[:, -1, :]
                curr_probs = nn.functional.softmax(curr_logits, dim=-1)
                curr_probs = torch.unsqueeze(curr_probs, dim=1)
                curr_hidden = curr_all_hidden[-1]
                new_accumulated_hidden = new_accumulated_hidden + torch.sum(curr_hidden, dim=1)

            prediction = classifier(new_accumulated_hidden / (curr_length + 1 + horizon_length))

            label = torch.tensor(prediction.shape[0] * [class_label], device=device, dtype=torch.long)
            discrim_loss = ce_loss(prediction, label)
            print(" pplm_discrim_loss:", discrim_loss.data.cpu().numpy())
            loss += discrim_loss
            loss_list.append(discrim_loss)

        kl_loss = 0.0
        if kl_scale > 0.0:
            unpert_probs = nn.functional.softmax(unpert_logits[:, -1, :], dim=-1)
            unpert_probs = unpert_probs + SMALL_CONST * (unpert_probs <= SMALL_CONST).float().to(device).detach()
            correction = SMALL_CONST * (probs <= SMALL_CONST).float().to(device).detach()
            corrected_probs = probs + correction.detach()
            kl_loss = kl_scale * ((corrected_probs * (corrected_probs / unpert_probs).log()).sum())
            print(" kl_loss", kl_loss.data.cpu().numpy())
            loss += kl_loss

        loss_per_iter.append(loss.data.cpu().numpy())
        print(" pplm_loss", (loss - kl_loss).data.cpu().numpy())

        # compute gradients
        loss.backward()

        # calculate gradient norms
        if grad_norms is not None and loss_type == PPLM_BOW:
            grad_norms = [
                torch.max(grad_norms[index], torch.norm(p_.grad * window_mask))
                for index, p_ in enumerate(curr_perturbation)
            ]
        else:
            grad_norms = [
                (torch.norm(p_.grad * window_mask) + SMALL_CONST) for index, p_ in enumerate(curr_perturbation)
            ]

        # normalize gradients
        grad = [
            -stepsize * (p_.grad * window_mask / grad_norms[index] ** gamma).data.cpu().numpy()
            for index, p_ in enumerate(curr_perturbation)
        ]

        # accumulate gradient
        grad_accumulator = list(map(add, grad, grad_accumulator))

        # reset gradients, just to make sure
        for p_ in curr_perturbation:
            p_.grad.data.zero_()

        # removing past from the graph
        new_past = []
        for p_ in past:
            new_past.append(p_.detach())
        past = new_past

    # apply the accumulated perturbations to the past
    grad_accumulator = [torch.from_numpy(p_).requires_grad_(True).to(device=device) for p_ in grad_accumulator]
    pert_past = list(map(add, past, grad_accumulator))



SUMMARY:
**Function Summary: `perturb_past`**

The `perturb_past` function is used to perturb the input gradient of a decoder-based language model during training. The primary purpose of this function is to accumulate gradients and update the model parameters based on these gradients, while also considering various losses such as language modeling and discrimination.

**Key Implementation Details:**

*   The function takes in a past tensor representing the input token sequence.
*   It generates a gradient perturbation system using a series of operations, including adding noise to the past tensor and computing the gradients.
*   The function uses PyTorch's `requires_grad=True` method to enable automatic differentiation for gradient computation.
*   The losses considered include language modeling, discrimination, and KL divergence.
*   The function also normalizes gradients based on their normalized magnitudes.

**Notable Parameters:**

*   `past`: the input tensor representing the past token sequence
*   `model`: the decoder-based language model
*   `last`: the last encoder output
*   `unpert_logits` and `unpert_past`: unperturbed logits and past tensors used for calculation purposes.
*   `classifier`: a classifier model used for prediction and loss calculation.
*   `class_label`: the label for the classifier model.
*   `loss_type`, `num_iterations`, `stepsize`, etc.: hyperparameters controlling the perturbation process.

**Return Values:**

*   `pert_past`, `new_accumulated_hidden`, and `grad_norms` : perturbed past tensor, updated accumulated hidden state, and gradients after perturbation.
*   `loss_per_iter`: a list of losses obtained at each iteration.


key: chroma:document
string_value:
CODE:

def get_classifier(
    name: Optional[str], class_label: Union[str, int], device: str
) -> Tuple[Optional[ClassificationHead], Optional[int]]:
    if name is None:
        return None, None

    params = DISCRIMINATOR_MODELS_PARAMS[name]
    classifier = ClassificationHead(class_size=params["class_size"], embed_size=params["embed_size"]).to(device)
    if "url" in params:
        resolved_archive_file = cached_path(params["url"])
    elif "path" in params:
        resolved_archive_file = params["path"]
    else:
        raise ValueError("Either url or path have to be specified in the discriminator model parameters")
    classifier.load_state_dict(torch.load(resolved_archive_file, map_location=device))
    classifier.eval()

    if isinstance(class_label, str):
        if class_label in params["class_vocab"]:
            label_id = params["class_vocab"][class_label]
        else:
            label_id = params["default_class"]
            print("class_label {} not in class_vocab".format(class_label))
            print("available values are: {}".format(params["class_vocab"]))
            print("using default class {}".format(label_id))

    elif isinstance(class_label, int):
        if class_label in set(params["class_vocab"].values()):
            label_id = class_label
        else:
            label_id = params["default_class"]
            print("class_label {} not in class_vocab".format(class_label))
            print("available values are: {}".format(params["class_vocab"]))
            print("using default class {}".format(label_id))

    else:
        label_id = params["default_class"]



SUMMARY:
**Function Summary**

`get_classifier(name: str, class_label: (str or int), device: str) -> (ClassificationHead, int)`
 
### Purpose and Functionality

Returns a `ClassificationHead` object and its corresponding label ID for the specified model. Sets up the classifier by loading pre-trained state from a file if provided.

### Key Implementation Details
- Loads `ClassificationHead` model based on `name` parameter
- Resolves archive file path (`url` or `path`) using cached_path function
- Loads pre-trained state using `torch.load`
- Maps labels to ID using `params["class_vocab"]` dictionary

### Notable Dependencies and Requirements
- `DISCRIMINATOR_MODELS_PARAMS`: Mapping of model names to parameters
- `cached_path`: Cache-based file path resolver (not shown in this code snippet)
- `torch`: PyTorch library for loading pre-trained state
- `ClassificationHead` class: PyTorch head for classification tasks

### Interfaces and Inherited Dependencies
- ` ClassificationHead`: Implemented from PyTorch models
- ` torch.load`: Used to load pre-trained state dictionary 

Note that the actual implementation of `cached_path` function is not shown in this code snippet.


key: chroma:document
string_value:
CODE:

def get_bag_of_words_indices(bag_of_words_ids_or_paths: List[str], tokenizer) -> List[List[List[int]]]:
    bow_indices = []
    for id_or_path in bag_of_words_ids_or_paths:
        if id_or_path in BAG_OF_WORDS_ARCHIVE_MAP:
            filepath = cached_path(BAG_OF_WORDS_ARCHIVE_MAP[id_or_path])
        else:
            filepath = id_or_path
        with open(filepath, "r") as f:
            words = f.read().strip().split("\n")
        bow_indices.append([tokenizer.encode(word.strip(), add_prefix_space=True) for word in words])


SUMMARY:
**Summary:**

`get_bag_of_words_indices` function converts a list of bag-of-words IDs/paths to numerical indexing format.

**Main Purpose and Functionality:**
The function takes a list of bag-of-words IDs (strings) or paths, tokenizes the corresponding words using the provided `tokenizer`, and returns a list of index matrices for each word ID/piath.

**Implementation Details:**

* Utilizes a dictionary (`BAG_OF_WORDS_ARCHIVE_MAP`) to map ID/paths to file paths.
* Reads word files from disk, splits lines by newlines, and strips whitespace.
* Applies tokenization using the `tokenizer` with prefix space enabled.

**Notable Dependencies and Requirements:**

* `tokenizer`: assumed tokenizer class implementing `encode` method.
* External dependency:
	+ `cached_path` function is used to cache file paths.
	+ `BAG_OF_WORDS_ARCHIVE_MAP`: a dictionary mapping ID/paths to file paths.


key: chroma:document
string_value:
CODE:

def build_bows_one_hot_vectors(bow_indices, tokenizer, device="cuda"):
    if bow_indices is None:
        return None

    one_hot_bows_vectors = []
    for single_bow in bow_indices:
        single_bow = list(filter(lambda x: len(x) <= 1, single_bow))
        single_bow = torch.tensor(single_bow).to(device)
        num_words = single_bow.shape[0]
        one_hot_bow = torch.zeros(num_words, tokenizer.vocab_size).to(device)
        one_hot_bow.scatter_(1, single_bow, 1)
        one_hot_bows_vectors.append(one_hot_bow)


SUMMARY:
**Summary:**

`build_bows_one_hot_vectors` function:

* Prepends one-hot encoding of BOW (Bag-of-Words) indices to tensor `tokenizer`.
* Filters out single-word tokens (`len(x) <= 1`) from each BOW index.

**Key Implementation Details:**

* Applies scatter operation for one-hot encoding using `torch.zeros` and `torch.scatter_`.
* Converts tensors to specified device (`device="cuda"`).

**Notable Dependencies and Requirements:**

* `tokenizer`: assumed to provide `vocab_size` attribute.
* `torch`: used for tensor operations.

**Function Signature:**
`build_bows_one_hot_vectors(bow_indices, tokenizer, device="cuda") -> list[torch.Tensor]`.


key: chroma:document
string_value:
CODE:

def full_text_generation(
    model,
    tokenizer,
    context=None,
    num_samples=1,
    device="cuda",
    bag_of_words=None,
    discrim=None,
    class_label=None,
    length=100,
    stepsize=0.02,
    temperature=1.0,
    top_k=10,
    sample=False,
    num_iterations=3,
    grad_length=10000,
    horizon_length=1,
    window_length=0,
    decay=False,
    gamma=1.5,
    gm_scale=0.9,
    kl_scale=0.01,
    repetition_penalty=1.0,
    **kwargs,
):
    classifier, class_id = get_classifier(discrim, class_label, device)

    bow_indices = []
    if bag_of_words:
        bow_indices = get_bag_of_words_indices(bag_of_words.split(";"), tokenizer)

    if bag_of_words and classifier:
        print("Both PPLM-BoW and PPLM-Discrim are on. This is not optimized.")
        loss_type = PPLM_BOW_DISCRIM

    elif bag_of_words:
        loss_type = PPLM_BOW
        print("Using PPLM-BoW")

    elif classifier is not None:
        loss_type = PPLM_DISCRIM
        print("Using PPLM-Discrim")

    else:
        raise Exception("Specify either a bag of words or a discriminator")

    unpert_gen_tok_text, _, _ = generate_text_pplm(
        model=model,
        tokenizer=tokenizer,
        context=context,
        device=device,
        length=length,
        sample=sample,
        perturb=False,
        repetition_penalty=repetition_penalty,
    )
    if device == "cuda":
        torch.cuda.empty_cache()

    pert_gen_tok_texts = []
    discrim_losses = []
    losses_in_time = []

    for i in range(num_samples):
        pert_gen_tok_text, discrim_loss, loss_in_time = generate_text_pplm(
            model=model,
            tokenizer=tokenizer,
            context=context,
            device=device,
            perturb=True,
            bow_indices=bow_indices,
            classifier=classifier,
            class_label=class_id,
            loss_type=loss_type,
            length=length,
            stepsize=stepsize,
            temperature=temperature,
            top_k=top_k,
            sample=sample,
            num_iterations=num_iterations,
            grad_length=grad_length,
            horizon_length=horizon_length,
            window_length=window_length,
            decay=decay,
            gamma=gamma,
            gm_scale=gm_scale,
            kl_scale=kl_scale,
            repetition_penalty=repetition_penalty,
        )
        pert_gen_tok_texts.append(pert_gen_tok_text)
        if classifier is not None:
            discrim_losses.append(discrim_loss.data.cpu().numpy())
        losses_in_time.append(loss_in_time)

    if device == "cuda":
        torch.cuda.empty_cache()



SUMMARY:
**Full Text Generation Function**

**Purpose and Functionality:**

This function generates full text based on input model, tokenizer, and context. It uses PPLM-BoW or PPLM-Discrim depending on the presence of bag-of-words or classification data.

**Key Implementation Details:**

* Uses `generate_text_pplm` function to generate text with perturbation
* Tries multiple samples using `num_samples` parameter
* Computes losses and accuracy metrics (discrim loss, loss in time)

**Notable Dependencies and Requirements:**

* `get_classifier` from `get_classifier.py`
* `get_bag_of_words_indices` from `tokenizer.py`
* `generate_text_pplm` from `generator.py`

**External Dependencies:**

* `torch`, `PPLM_BOW_DISCRIM`, `PPLM_BOW`, `PPLM_DISCRIM`, `RepetitionPenalty`

**Interfaces and Inheritance:**

* Implements `generative_text` interface
* Inherits from `unperturbated_text_generator` base class


key: chroma:document
string_value:
CODE:

def generate_text_pplm(
    model,
    tokenizer,
    context=None,
    past=None,
    device="cuda",
    perturb=True,
    bow_indices=None,
    classifier=None,
    class_label=None,
    loss_type=0,
    length=100,
    stepsize=0.02,
    temperature=1.0,
    top_k=10,
    sample=False,
    num_iterations=3,
    grad_length=10000,
    horizon_length=1,
    window_length=0,
    decay=False,
    gamma=1.5,
    gm_scale=0.9,
    kl_scale=0.01,
    repetition_penalty=1.0,
):
    output_so_far = None
    if context:
        context_t = torch.tensor(context, device=device, dtype=torch.long)
        while len(context_t.shape) < 2:
            context_t = context_t.unsqueeze(0)
        output_so_far = context_t

    # collect one hot vectors for bags of words
    one_hot_bows_vectors = build_bows_one_hot_vectors(bow_indices, tokenizer, device)

    grad_norms = None
    last = None
    unpert_discrim_loss = 0
    loss_in_time = []
    for i in trange(length, ascii=True):
        # Get past/probs for current output, except for last word
        # Note that GPT takes 2 inputs: past + current_token

        # run model forward to obtain unperturbed
        if past is None and output_so_far is not None:
            last = output_so_far[:, -1:]
            if output_so_far.shape[1] > 1:
                past = model(output_so_far[:, :-1])["past_key_values"]

        lm_output = model(output_so_far)
        unpert_logits, unpert_past, unpert_all_hidden = (
            lm_output["logits"],
            lm_output["past_key_values"],
            lm_output["hidden_states"],
        )
        unpert_last_hidden = unpert_all_hidden[-1]

        # check if we are abowe grad max length
        if i >= grad_length:
            current_stepsize = stepsize * 0
        else:
            current_stepsize = stepsize

        # modify the past if necessary
        if not perturb or num_iterations == 0:
            pert_past = past

        else:
            accumulated_hidden = unpert_last_hidden[:, :-1, :]
            accumulated_hidden = torch.sum(accumulated_hidden, dim=1)

            if past is not None:
                pert_past, _, grad_norms, loss_this_iter = perturb_past(
                    past,
                    model,
                    last,
                    unpert_past=unpert_past,
                    unpert_logits=unpert_logits,
                    accumulated_hidden=accumulated_hidden,
                    grad_norms=grad_norms,
                    stepsize=current_stepsize,
                    one_hot_bows_vectors=one_hot_bows_vectors,
                    classifier=classifier,
                    class_label=class_label,
                    loss_type=loss_type,
                    num_iterations=num_iterations,
                    horizon_length=horizon_length,
                    window_length=window_length,
                    decay=decay,
                    gamma=gamma,
                    kl_scale=kl_scale,
                    device=device,
                )
                loss_in_time.append(loss_this_iter)
            else:
                pert_past = past

        lm_output = model(last, past_key_values=pert_past)
        pert_logits, past = (
            lm_output["logits"],
            lm_output["past_key_values"],
        )
        pert_logits = pert_logits[:, -1, :] / temperature  # + SMALL_CONST

        for token_idx in set(output_so_far[0].tolist()):
            if pert_logits[0, token_idx] < 0:
                pert_logits[0, token_idx] *= repetition_penalty
            else:
                pert_logits[0, token_idx] /= repetition_penalty

        pert_probs = nn.functional.softmax(pert_logits, dim=-1)

        if classifier is not None:
            ce_loss = nn.CrossEntropyLoss()
            prediction = classifier(torch.mean(unpert_last_hidden, dim=1))
            label = torch.tensor([class_label], device=device, dtype=torch.long)
            unpert_discrim_loss = ce_loss(prediction, label)
            print("unperturbed discrim loss", unpert_discrim_loss.data.cpu().numpy())
        else:
            unpert_discrim_loss = 0

        # Fuse the modified model and original model
        if perturb:
            unpert_probs = nn.functional.softmax(unpert_logits[:, -1, :], dim=-1)

            pert_probs = (pert_probs**gm_scale) * (unpert_probs ** (1 - gm_scale))  # + SMALL_CONST
            pert_probs = top_k_filter(pert_probs, k=top_k, probs=True)  # + SMALL_CONST

            # rescale
            if torch.sum(pert_probs) <= 1:
                pert_probs = pert_probs / torch.sum(pert_probs)

        else:
            pert_logits = top_k_filter(pert_logits, k=top_k)  # + SMALL_CONST
            pert_probs = nn.functional.softmax(pert_logits, dim=-1)

        # sample or greedy
        if sample:
            last = torch.multinomial(pert_probs, num_samples=1)

        else:
            _, last = torch.topk(pert_probs, k=1, dim=-1)

        # update context/output_so_far appending the new token
        output_so_far = last if output_so_far is None else torch.cat((output_so_far, last), dim=1)

        print(tokenizer.decode(output_so_far.tolist()[0]))



SUMMARY:
**Function Name:** generate_text_pplm
**Primary Purpose:**

Generate text by perturbing and filtering the output of a language model.

**Main Functionality:**

The function takes input from a given model, tokenizer, and context, and proceeds as follows:

1. Checks if the context is provided; if so, it uses it as the initial output.
2. Collects one-hot vectors for bag-of-words representation.
3. Loops through a specified length (default 100) to generate text:
	* Runs the model forward to obtain unperturbed output.
	* Perturbs and filters the output based on gradients, past values, and parameters (e.g., step size, temperature).
	* Selects the top-k tokens using top_k_filter or softmax.
4. Samples or performs greedy selection based on whether the `sample` flag is set to True.

**Key Implementation Details:**

1. Perturb Past Function (`perturb_past`): modifies past values based on gradients and parameters.
2. Loss Computation:
	* Computing loss with classifier using CrossEntropyLoss.
	* Computing unperturbed discriminative loss.
3. Sampling or Greedy Selection:
	* Uses torch.multinomial for sampling, or topk for greedy selection.

**Notable Dependencies or Requirements:**

1. Torch tensor for tensor operations.
2. PyTorch Lightning for model forward passes and gradients.
3. Build Bows One Hot Vectors (`build_bows_one_hot_vectors`) function.

**External Dependencies:**

1. `torch`
2. `torch.nn.functional`


key: chroma:document
string_value:
CODE:

def set_generic_model_params(discrim_weights, discrim_meta):
    if discrim_weights is None:
        raise ValueError("When using a generic discriminator, discrim_weights need to be specified")
    if discrim_meta is None:
        raise ValueError("When using a generic discriminator, discrim_meta need to be specified")

    with open(discrim_meta, "r") as discrim_meta_file:
        meta = json.load(discrim_meta_file)
    meta["path"] = discrim_weights


SUMMARY:
**Summary:**

The `set_generic_model_params` function initializes parameters for a generic discriminator model.

* It validates that `discrim_weights` and `discrim_meta` are provided.
* Loads metadata from `discrim_meta` using JSON, updating the `path` field with `discrim_weights`.
* Updates the global dictionary `DISCRIMINATOR_MODELS_PARAMS["generic"]` with the loaded metadata.

**Key Implementation Details:**

* Validates input parameters `discrim_weights` and `discrimmeta`
* Loads metadata from file `discrim_meta` using JSON

**Dependencies/Requirements:**

* `json` library for loading metadata
* Global dictionary `DISCRIMINATOR_MODELS_PARAMS`

Note: The function assumes that the path to the JSON file matches the provided `path` in the loaded metadata.


key: chroma:document
string_value:
CODE:

def run_pplm_example(
    pretrained_model="openai-community/gpt2-medium",
    cond_text="",
    uncond=False,
    num_samples=1,
    bag_of_words=None,
    discrim=None,
    discrim_weights=None,
    discrim_meta=None,
    class_label=-1,
    length=100,
    stepsize=0.02,
    temperature=1.0,
    top_k=10,
    sample=False,
    num_iterations=3,
    grad_length=10000,
    horizon_length=1,
    window_length=0,
    decay=False,
    gamma=1.5,
    gm_scale=0.9,
    kl_scale=0.01,
    seed=0,
    no_cuda=False,
    colorama=False,
    repetition_penalty=1.0,
):
    # set Random seed
    torch.manual_seed(seed)
    np.random.seed(seed)

    # set the device
    device = "cuda" if torch.cuda.is_available() and not no_cuda else "cpu"

    if discrim == "generic":
        set_generic_model_params(discrim_weights, discrim_meta)

    if discrim is not None:
        pretrained_model = DISCRIMINATOR_MODELS_PARAMS[discrim]["pretrained_model"]
        print("discrim = {}, pretrained_model set to discriminator's = {}".format(discrim, pretrained_model))

    # load pretrained model
    model = GPT2LMHeadModel.from_pretrained(pretrained_model, output_hidden_states=True)
    model.to(device)
    model.eval()

    # load tokenizer
    tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model)

    # Freeze GPT-2 weights
    for param in model.parameters():
        param.requires_grad = False

    # figure out conditioning text
    if uncond:
        tokenized_cond_text = tokenizer.encode([tokenizer.bos_token])
    else:
        raw_text = cond_text
        while not raw_text:
            print("Did you forget to add `--cond_text`? ")
            raw_text = input("Model prompt >>> ")
        tokenized_cond_text = tokenizer.encode(tokenizer.bos_token + raw_text)

    print("= Prefix of sentence =")
    print(tokenizer.decode(tokenized_cond_text))
    print()

    # generate unperturbed and perturbed texts

    # full_text_generation returns:
    # unpert_gen_tok_text, pert_gen_tok_texts, discrim_losses, losses_in_time
    unpert_gen_tok_text, pert_gen_tok_texts, _, _ = full_text_generation(
        model=model,
        tokenizer=tokenizer,
        context=tokenized_cond_text,
        device=device,
        num_samples=num_samples,
        bag_of_words=bag_of_words,
        discrim=discrim,
        class_label=class_label,
        length=length,
        stepsize=stepsize,
        temperature=temperature,
        top_k=top_k,
        sample=sample,
        num_iterations=num_iterations,
        grad_length=grad_length,
        horizon_length=horizon_length,
        window_length=window_length,
        decay=decay,
        gamma=gamma,
        gm_scale=gm_scale,
        kl_scale=kl_scale,
        repetition_penalty=repetition_penalty,
    )

    # untokenize unperturbed text
    unpert_gen_text = tokenizer.decode(unpert_gen_tok_text.tolist()[0])

    print("=" * 80)
    print("= Unperturbed generated text =")
    print(unpert_gen_text)
    print()

    generated_texts = []

    bow_word_ids = set()
    if bag_of_words and colorama:
        bow_indices = get_bag_of_words_indices(bag_of_words.split(";"), tokenizer)
        for single_bow_list in bow_indices:
            # filtering all words in the list composed of more than 1 token
            filtered = list(filter(lambda x: len(x) <= 1, single_bow_list))
            # w[0] because we are sure w has only 1 item because previous fitler
            bow_word_ids.update(w[0] for w in filtered)

    # iterate through the perturbed texts
    for i, pert_gen_tok_text in enumerate(pert_gen_tok_texts):
        try:
            # untokenize unperturbed text
            if colorama:
                import colorama

                pert_gen_text = ""
                for word_id in pert_gen_tok_text.tolist()[0]:
                    if word_id in bow_word_ids:
                        pert_gen_text += "{}{}{}".format(
                            colorama.Fore.RED,
                            tokenizer.decode([word_id]),
                            colorama.Style.RESET_ALL,
                        )
                    else:
                        pert_gen_text += tokenizer.decode([word_id])
            else:
                pert_gen_text = tokenizer.decode(pert_gen_tok_text.tolist()[0])

            print("= Perturbed generated text {} =".format(i + 1))
            print(pert_gen_text)
            print()
        except Exception as exc:
            print("Ignoring error while generating perturbed text:", exc)

        # keep the prefix, perturbed seq, original seq for each index
        generated_texts.append((tokenized_cond_text, pert_gen_tok_text, unpert_gen_tok_text))



SUMMARY:
**Function Summary: `run_pplm_example`**

This function runs an example of the Popularly Picked Language Model (PPLM) generation using the Hugging Face Transformers library.

**Main Purpose and Functionality**

The primary purpose is to generate unperturbed and perturbed texts based on a provided conditioning text. The function takes several input parameters, including:

* `pretrained_model`: The pre-trained model to use for GPT-2.
* `cond_text`: The conditioning text to use for generation.
* `uncond`: A flag indicating whether the conditioning text is empty.
* Other optional parameters control the generation process, such as `num_samples`, `bag_of_words`, and `discrim`.

**Implementation Details**

The function uses the following key algorithms and designs:

* Loads a pre-trained GPT-2 model and tokenizer using Hugging Face's `GPT2LMHeadModel` and `GPT2Tokenizer`.
* Converts the conditioning text to token IDs using the tokenizer.
* Calls the `full_text_generation` function from the PPLM repository, passing in various parameters and models.

**Notable Dependencies and Requirements**

The function relies on the following external dependencies:

* `torch`: The PyTorch library for GPU acceleration (optional).
* `colorama`: A color-based terminal interface (optional).

It also uses several classes and functions from Hugging Face's Transformers library, including:

* `GPT2LMHeadModel`
* `GPT2Tokenizer`
* `full_text_generation`

**Key Method Signatures and Interfaces**

The function calls the following method signature:

`full_text_generation(model=model, tokenizer=tokenizer, ...)`


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/zero-shot-distillation/distill_classifier.py


key: chroma:document
string_value:
CODE:

class DistillationTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        target_p = inputs["labels"]
        outputs = model(inputs["input_ids"], attention_mask=inputs["attention_mask"])
        logits = outputs[0]

        loss = -torch.sum(target_p * logits.log_softmax(dim=-1), axis=-1).mean()

        if return_outputs:
            return loss, outputs



SUMMARY:
**Class Summary: DistillationTrainer**

* **Main Purpose:** Implementation of distillation training logic for a model, using the distillation objective formula to compute the loss between teacher and student models.

**Key Functionality:**
- `compute_loss`: Computes the distillation loss between Teacher (model) and Student (model) models. 

**Implementation Details:**
- Computes logits from the model's output with soft attention masks.
- Applies the negative log softmax on the logits, resulting in a probability distribution that aligns with the predicted labels.

**Notable Dependencies/Requirements:**

* `Trainer`: Parent class of DistillationTrainer
* `model`, `inputs`, `torch`

**Interface/Method Signatures:**
- `compute_loss(self, model, inputs, return_outputs=False):` computes loss and returns it along with student output if `return_outputs=True`.


key: chroma:document
string_value:
CODE:

def read_lines(path):
    lines = []
    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if len(line) > 0:
                lines.append(line)


SUMMARY:
**Function Summary**

Reads all non-empty lines from a specified file path and returns them as a list.

### Key Implementation Details

* The `read_lines` function reads the file line by line using a `for` loop.
* Each line is stripped of leading/trailing whitespaces using the `strip()` method.
* Non-empty lines are appended to a dynamic array `lines`.

### Notable Dependencies or Requirements

* None. This function relies solely on standard Python libraries.

### Method Signature and Interface

`read_lines(path: str) -> list[str]`

This function takes a single string argument `path`, representing the file path, and returns a list of strings (`list[str]`).


key: chroma:document
string_value:
CODE:

def get_premise_hypothesis_pairs(examples, class_names, hypothesis_template):
    premises = []
    hypotheses = []
    for example in examples:
        for name in class_names:
            premises.append(example)
            hypotheses.append(hypothesis_template.format(name))


SUMMARY:
**Summary:**

`get_premise_hypothesis_pairs`: Returns pairs of premise and hypothesis for each example in `examples`, formatted using `hypothesis_template`. The function assumes knowledge of `class_names`.

* **Purpose:** Generate data pairs for testing hypothesis template.
* **Implementation:**
	+ Iterates through `examples` and inserts each example into the premises list with no replacement.
	+ For each example, it appends a generated hypothesis to the hypotheses list by formatting the `hypothesis_template`.
	+ Returns two lists of equal length: premises and hypotheses.

**Key Details:**

* Uses `class_names` template variable for generating hypotheses.
* Formats `example` as premise with no replacement.
* Uses `get_premise_hypothesis_pairs.format()` to inject class names into the hypothesis template.


key: chroma:document
string_value:
CODE:

def get_entailment_id(config):
    for label, ind in config.label2id.items():
        if label.lower().startswith("entail"):
            return ind
    logger.warning("Could not identify entailment dimension from teacher config label2id. Setting to -1.")


SUMMARY:
**Summary:**

`get_entailment_id` function retrieves the `ind` value corresponding to an "entail" label from the `config.label2id` dictionary.

**Key Details:**

* Returns `-1` if "entail" label is not found in `teacher config label2id`.
* Uses `label.lower()` for case-insensitive comparison.
* Logs a warning message if label cannot be identified.

**Code Notes:**

* This function appears to be part of a natural language processing (NLP) or machine learning application, given the mention of "entailment" and "-1".


key: chroma:document
string_value:
CODE:

def get_teacher_predictions(
    model_path: str,
    examples: List[str],
    class_names: List[str],
    hypothesis_template: str,
    batch_size: int,
    temperature: float,
    multi_label: bool,
    use_fast_tokenizer: bool,
    no_cuda: bool,
    fp16: bool,
):
    """
    Gets predictions by the same method as the zero-shot pipeline but with DataParallel & more efficient batching
    """
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    model_config = model.config
    if not no_cuda and torch.cuda.is_available():
        model = nn.DataParallel(model.cuda())
        batch_size *= len(model.device_ids)
    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=use_fast_tokenizer)

    premises, hypotheses = get_premise_hypothesis_pairs(examples, class_names, hypothesis_template)
    logits = []

    for i in tqdm(range(0, len(premises), batch_size)):
        batch_premises = premises[i : i + batch_size]
        batch_hypotheses = hypotheses[i : i + batch_size]

        encodings = tokenizer(
            batch_premises,
            batch_hypotheses,
            padding=True,
            truncation="only_first",
            return_tensors="pt",
        )

        with torch.cuda.amp.autocast(enabled=fp16):
            with torch.no_grad():
                outputs = model(**encodings)
        logits.append(outputs.logits.detach().cpu().float())

    entail_id = get_entailment_id(model_config)
    contr_id = -1 if entail_id == 0 else 0
    logits = torch.cat(logits, dim=0)  # N*K x 3
    nli_logits = logits.reshape(len(examples), len(class_names), -1)[..., [contr_id, entail_id]]  # N x K x 2

    if multi_label:
        # softmax over (contr, entail) logits for each class independently
        nli_prob = (nli_logits / temperature).softmax(-1)
    else:
        # softmax over entail logits across classes s.t. class probabilities sum to 1.
        nli_prob = (nli_logits / temperature).softmax(1)



SUMMARY:
**get_teacher_predictions Function Summary**

### Purpose and Functionality
This function generates teacher-forced training data by feeding input premises, hypotheses, and class labels to a pre-trained model through an efficient batching mechanism using DataParallel. The outputs are stored in `logits` before being processed further.

### Key Implementation Details

*   Models an instance of `_AI_model_for_sequences_Using_AutoModelForSequenceClassification` from the pre-trained model path.
*   Utilizes `AutoTokenizer.from_pretrained()` to create a tokenizer for faster encoding.
*   Creates batches of premises and hypotheses using the batching size for efficient computation.

### Notable Dependencies or Requirements

*   Uses `torch.cuda.is_available()` with flag `no_cuda` (bool) and `fp16` (bool).
*   Inherit's from `nn.DataParallel` from the model object with `cuda()` implementation.

### Summary of Return Value
The function returns `nli_prob[..., 1]`, which represents teacher-forced training output, resulting in probabilities for entailment over contrasting classes.

### External Dependencies:

- Uses the `AutoModelForSequenceClassification`, `AutoTokenizer`, and other corresponding libraries.
  * [transformers library](https://github.com/huggingface/transformers)


key: chroma:document
string_value:
CODE:

def main():
    parser = HfArgumentParser(
        (DataTrainingArguments, TeacherModelArguments, StudentModelArguments, DistillTrainingArguments),
        description=DESCRIPTION,
    )

    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        data_args, teacher_args, student_args, training_args = parser.parse_json_file(
            json_file=os.path.abspath(sys.argv[1])
        )
    else:
        data_args, teacher_args, student_args, training_args = parser.parse_args_into_dataclasses()

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    if is_main_process(training_args.local_rank):
        utils.logging.set_verbosity_info()
        utils.logging.enable_default_handler()
        utils.logging.enable_explicit_format()

    if training_args.local_rank != -1:
        raise ValueError("Distributed training is not currently supported.")
    if training_args.tpu_num_cores is not None:
        raise ValueError("TPU acceleration is not currently supported.")

    logger.info(f"Training/evaluation parameters {training_args}")

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # 1. read in data
    examples = read_lines(data_args.data_file)
    class_names = read_lines(data_args.class_names_file)

    # 2. get teacher predictions and load into dataset
    logger.info("Generating predictions from zero-shot teacher model")
    teacher_soft_preds = get_teacher_predictions(
        teacher_args.teacher_name_or_path,
        examples,
        class_names,
        teacher_args.hypothesis_template,
        teacher_args.teacher_batch_size,
        teacher_args.temperature,
        teacher_args.multi_label,
        data_args.use_fast_tokenizer,
        training_args.no_cuda,
        training_args.fp16,
    )
    dataset = Dataset.from_dict(
        {
            "text": examples,
            "labels": teacher_soft_preds,
        }
    )

    # 3. create student
    logger.info("Initializing student model")
    model = AutoModelForSequenceClassification.from_pretrained(
        student_args.student_name_or_path, num_labels=len(class_names)
    )
    tokenizer = AutoTokenizer.from_pretrained(student_args.student_name_or_path, use_fast=data_args.use_fast_tokenizer)
    model.config.id2label = dict(enumerate(class_names))
    model.config.label2id = {label: i for i, label in enumerate(class_names)}

    # 4. train student on teacher predictions
    dataset = dataset.map(tokenizer, input_columns="text")
    dataset.set_format("torch")

    def compute_metrics(p, return_outputs=False):
        preds = p.predictions.argmax(-1)
        proxy_labels = p.label_ids.argmax(-1)  # "label_ids" are actually distributions
        return {"agreement": (preds == proxy_labels).mean().item()}

    trainer = DistillationTrainer(
        model=model,
        tokenizer=tokenizer,
        args=training_args,
        train_dataset=dataset,
        compute_metrics=compute_metrics,
    )

    if training_args.do_train:
        logger.info("Training student model on teacher predictions")
        trainer.train()

    if training_args.do_eval:
        agreement = trainer.evaluate(eval_dataset=dataset)["eval_agreement"]
        logger.info(f"Agreement of student and teacher predictions: {agreement * 100:0.2f}%")



SUMMARY:
**Main Functionality**

This is the main entry point of a script that trains a student model on top of a pre-trained teacher model. The script takes in various arguments and parses them into different namespaces, including `DataTrainingArguments`, `TeacherModelArguments`, `StudentModelArguments`, and `DistillTrainingArguments`.

The primary purpose of this function is to:

1. Parse input arguments from files or command-line arguments.
2. Set up logging and initialize the model.
3. Load data and generate teacher predictions.
4. Train a student model on top of the teacher prediction.
5. Evaluate the trained student model.

**Key Implementation Details**

The main implementation details are:

* Loading data and generating teacher predictions using `get_teacher_predictions` function.
* Creating a student model using `AutoModelForSequenceClassification.from_pretrained`.
* Training a student model on top of teacher predictions using `DistillationTrainer`.

**Important Method Signatures and Interfaces**

The important method signatures and interfaces used are:

* `parser.parse_json_file`
* `parser.parse_args_into_dataclasses`
* `get_last_checkpoint`
* `set_seed`
* `read_lines`

**Notable Dependencies or Requirements**

The notable dependencies or requirements used in this script are:

* `transformers` library
* `torch` library

Note: The code snippet assumes that the `DistillationTrainer` and `AutoModelForSequenceClassification`, `AudioTokenizer` classes are defined elsewhere in the codebase, and their exact implementation is not provided here.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/decision_transformer/run_decision_transformer.py


key: chroma:document
string_value:
CODE:

def get_action(model, states, actions, rewards, returns_to_go, timesteps):
    # we don't care about the past rewards in this model

    states = states.reshape(1, -1, model.config.state_dim)
    actions = actions.reshape(1, -1, model.config.act_dim)
    returns_to_go = returns_to_go.reshape(1, -1, 1)
    timesteps = timesteps.reshape(1, -1)

    if model.config.max_length is not None:
        states = states[:, -model.config.max_length :]
        actions = actions[:, -model.config.max_length :]
        returns_to_go = returns_to_go[:, -model.config.max_length :]
        timesteps = timesteps[:, -model.config.max_length :]

        # pad all tokens to sequence length
        attention_mask = torch.cat(
            [torch.zeros(model.config.max_length - states.shape[1]), torch.ones(states.shape[1])]
        )
        attention_mask = attention_mask.to(dtype=torch.long, device=states.device).reshape(1, -1)
        states = torch.cat(
            [
                torch.zeros(
                    (states.shape[0], model.config.max_length - states.shape[1], model.config.state_dim),
                    device=states.device,
                ),
                states,
            ],
            dim=1,
        ).to(dtype=torch.float32)
        actions = torch.cat(
            [
                torch.zeros(
                    (actions.shape[0], model.config.max_length - actions.shape[1], model.config.act_dim),
                    device=actions.device,
                ),
                actions,
            ],
            dim=1,
        ).to(dtype=torch.float32)
        returns_to_go = torch.cat(
            [
                torch.zeros(
                    (returns_to_go.shape[0], model.config.max_length - returns_to_go.shape[1], 1),
                    device=returns_to_go.device,
                ),
                returns_to_go,
            ],
            dim=1,
        ).to(dtype=torch.float32)
        timesteps = torch.cat(
            [
                torch.zeros(
                    (timesteps.shape[0], model.config.max_length - timesteps.shape[1]), device=timesteps.device
                ),
                timesteps,
            ],
            dim=1,
        ).to(dtype=torch.long)
    else:
        attention_mask = None

    _, action_preds, _ = model(
        states=states,
        actions=actions,
        rewards=rewards,
        returns_to_go=returns_to_go,
        timesteps=timesteps,
        attention_mask=attention_mask,
        return_dict=False,
    )



SUMMARY:
**Function Summary: `get_action`**

### Purpose and Functionality

The `get_action` function processes input data for an action prediction model. It takes in various inputs (states, actions, rewards, returns_to_go, timesteps) and passes them through the model to predict the next action. The function reshares and pads the input tensors according to the model's length requirements.

### Key Implementation Details

* Reshapes and reshuffles input tensors (`states`, `actions`, `returns_to_go`, `timesteps`) for efficient processing.
* Adapts lengths of input tensors based on `model.config.max_length` or pads with zeros if not specified.
* Applies attention mask to sequence length.

### Notable Dependencies and Requirements

* Model object (`model`)
	+ Has a configuration object (`model.config`) with attributes:
		- `max_length`
* External dependencies:
	+ PyTorch library for tensor operations (e.g., `torch.cat`, `torch.zeros`)

### Key Method Signatures and Interfaces

* Called on the model object (`model`) with various input arguments.

### Return Value

The predicted next action (`action_preds[0, -1]`).


key: chroma:document
string_value:
CODE:

def entropy(p):
    """Compute the entropy of a probability distribution"""
    plogp = p * torch.log(p)
    plogp[p == 0] = 0


SUMMARY:
**Summary:**

The `entropy` function computes the entropy of a probability distribution. It takes a tensor `p` as input, calculates the product `p * log(p)`, handles cases where probabilities are zero, and returns the negative sum of this product along the last dimension.

**Key Details:**

- Calculates probability-weighted log-probability (`plogp`) using `torch.log`.
- Handles zero probabilities by setting them to 0.
- Returns the entropy value: `-sum(plogp)`.

**Dependencies/Requirements:**

- Uses `torch`, a PyTorch library.


key: chroma:document
string_value:
CODE:

def print_2d_tensor(tensor):
    """Print a 2D tensor"""
    logger.info("lv, h >\t" + "\t".join(f"{x + 1}" for x in range(len(tensor))))
    for row in range(len(tensor)):
        if tensor.dtype != torch.long:
            logger.info(f"layer {row + 1}:\t" + "\t".join(f"{x:.5f}" for x in tensor[row].cpu().data))
        else:


SUMMARY:
**Function Summary: `print_2d_tensor`**

**Purpose and Functionality**

Prints a 2D tensor, logging its dimensions and values. If the tensor contains long integers, it is printed as decimal numbers; otherwise, as floating-point numbers with 5 decimal places.

**Key Implementation Details**

* Logs tensor dimensions
* Iterates over rows to print tensor values (CPU-bound)
* Uses `torch.long` type hint and performs explicit casting for numerical representation

**Notable Dependencies or Requirements**

* `torch`: PyTorch library for deep learning
* `logger`: Custom logging object not shown in the code snippet


key: chroma:document
string_value:
CODE:

def compute_heads_importance(
    args, model, eval_dataloader, compute_entropy=True, compute_importance=True, head_mask=None, actually_pruned=False
):
    """This method shows how to compute:
    - head attention entropy
    - head importance scores according to http://arxiv.org/abs/1905.10650
    """
    # Prepare our tensors
    n_layers, n_heads = model.config.num_hidden_layers, model.config.num_attention_heads
    head_importance = torch.zeros(n_layers, n_heads).to(args.device)
    attn_entropy = torch.zeros(n_layers, n_heads).to(args.device)

    if head_mask is None:
        head_mask = torch.ones(n_layers, n_heads).to(args.device)

    head_mask.requires_grad_(requires_grad=True)
    # If actually pruned attention multi-head, set head mask to None to avoid shape mismatch
    if actually_pruned:
        head_mask = None

    preds = None
    labels = None
    tot_tokens = 0.0

    for step, inputs in enumerate(tqdm(eval_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])):
        for k, v in inputs.items():
            inputs[k] = v.to(args.device)

        # Do a forward pass (not with torch.no_grad() since we need gradients for importance score - see below)
        outputs = model(**inputs, head_mask=head_mask)
        loss, logits, all_attentions = (
            outputs[0],
            outputs[1],
            outputs[-1],
        )  # Loss and logits are the first, attention the last
        loss.backward()  # Backpropagate to populate the gradients in the head mask

        if compute_entropy:
            for layer, attn in enumerate(all_attentions):
                masked_entropy = entropy(attn.detach()) * inputs["attention_mask"].float().unsqueeze(1)
                attn_entropy[layer] += masked_entropy.sum(-1).sum(0).detach()

        if compute_importance:
            head_importance += head_mask.grad.abs().detach()

        # Also store our logits/labels if we want to compute metrics afterwards
        if preds is None:
            preds = logits.detach().cpu().numpy()
            labels = inputs["labels"].detach().cpu().numpy()
        else:
            preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
            labels = np.append(labels, inputs["labels"].detach().cpu().numpy(), axis=0)

        tot_tokens += inputs["attention_mask"].float().detach().sum().data

    # Normalize
    attn_entropy /= tot_tokens
    head_importance /= tot_tokens
    # Layerwise importance normalization
    if not args.dont_normalize_importance_by_layer:
        exponent = 2
        norm_by_layer = torch.pow(torch.pow(head_importance, exponent).sum(-1), 1 / exponent)
        head_importance /= norm_by_layer.unsqueeze(-1) + 1e-20

    if not args.dont_normalize_global_importance:
        head_importance = (head_importance - head_importance.min()) / (head_importance.max() - head_importance.min())

    # Print/save matrices
    np.save(os.path.join(args.output_dir, "attn_entropy.npy"), attn_entropy.detach().cpu().numpy())
    np.save(os.path.join(args.output_dir, "head_importance.npy"), head_importance.detach().cpu().numpy())

    logger.info("Attention entropies")
    print_2d_tensor(attn_entropy)
    logger.info("Head importance scores")
    print_2d_tensor(head_importance)
    logger.info("Head ranked by importance scores")
    head_ranks = torch.zeros(head_importance.numel(), dtype=torch.long, device=args.device)
    head_ranks[head_importance.view(-1).sort(descending=True)[1]] = torch.arange(
        head_importance.numel(), device=args.device
    )
    head_ranks = head_ranks.view_as(head_importance)
    print_2d_tensor(head_ranks)



SUMMARY:
**Summary: Compute Heads Importance**

This function computes the importance scores of heads in a transformer model using the "Attention is All You Need" approach. It calculates:

1. **Head Attention Entropy**: measures the uncertainty of each attention head
2. **Head Importance Scores**: ranks the heads according to their importance

**Key Implementation Details:**

* **Forward Pass**: performs a forward pass through the model with the input data and computes the output
* **Backpropagation**: backpropagates the gradients through the loss function to update the attention mask
* **Entropy Calculation**: calculates the entropy of each attention head using the softmax formula

**Notable Parameters:**

* `head_mask`: a tensor controlling the pruning of attention heads
* `actually_pruned`: a boolean indicating whether to prune attention heads
* `compute_entropy` and `compute_importance`: booleans controlling whether to compute entropy and importance scores respectively

**Dependencies and Interfaces:**

* `torch`: PyTorch library for deep learning
* `model.config`: configuration module of the transformer model
* `eval_dataloader`: evaluation data loader for the dataset


key: chroma:document
string_value:
CODE:

def mask_heads(args, model, eval_dataloader):
    """This method shows how to mask head (set some heads to zero), to test the effect on the network,
    based on the head importance scores, as described in Michel et al. (http://arxiv.org/abs/1905.10650)
    """
    _, head_importance, preds, labels = compute_heads_importance(args, model, eval_dataloader, compute_entropy=False)
    preds = np.argmax(preds, axis=1) if args.output_mode == "classification" else np.squeeze(preds)
    original_score = glue_compute_metrics(args.task_name, preds, labels)[args.metric_name]
    logger.info("Pruning: original score: %f, threshold: %f", original_score, original_score * args.masking_threshold)

    new_head_mask = torch.ones_like(head_importance)
    num_to_mask = max(1, int(new_head_mask.numel() * args.masking_amount))

    current_score = original_score
    while current_score >= original_score * args.masking_threshold:
        head_mask = new_head_mask.clone()  # save current head mask
        # heads from least important to most - keep only not-masked heads
        head_importance[head_mask == 0.0] = float("Inf")
        current_heads_to_mask = head_importance.view(-1).sort()[1]

        if len(current_heads_to_mask) <= num_to_mask:
            break

        # mask heads
        current_heads_to_mask = current_heads_to_mask[:num_to_mask]
        logger.info("Heads to mask: %s", str(current_heads_to_mask.tolist()))
        new_head_mask = new_head_mask.view(-1)
        new_head_mask[current_heads_to_mask] = 0.0
        new_head_mask = new_head_mask.view_as(head_mask)
        new_head_mask = new_head_mask.clone().detach()
        print_2d_tensor(new_head_mask)

        # Compute metric and head importance again
        _, head_importance, preds, labels = compute_heads_importance(
            args, model, eval_dataloader, compute_entropy=False, head_mask=new_head_mask
        )
        preds = np.argmax(preds, axis=1) if args.output_mode == "classification" else np.squeeze(preds)
        current_score = glue_compute_metrics(args.task_name, preds, labels)[args.metric_name]
        logger.info(
            "Masking: current score: %f, remaining heads %d (%.1f percents)",
            current_score,
            new_head_mask.sum(),
            new_head_mask.sum() / new_head_mask.numel() * 100,
        )

    logger.info("Final head mask")
    print_2d_tensor(head_mask)
    np.save(os.path.join(args.output_dir, "head_mask.npy"), head_mask.detach().cpu().numpy())



SUMMARY:
**Summary:**

`mask_heads` method calculates the importance scores for network heads using `compute_heads_importance`, then iteratively applies masking to heads with low importance, updating the metric and head importance until the threshold is reached.

**Key Implementation Details:**

* Iterates through heads from least important to most, keeping only unmasked heads until a specified number (`num_to_mask`) is reached.
* Updates `head_importance` with masked heads set to infinity values.
* Recomputes metric and head importance after each masking step.
* Prints 2D tensor of new head mask for logging.

**Notable Dependencies/Requirements:**

* `compute_heads_importance`
* `glue_compute_metrics`
* `logger` (for logging purposes)
* `print_2d_tensor` (for logging purpose)

**Additional Information:**

The method takes in three inputs:
- `args`: contains hyperparameter settings
- `model`: neural network model
- `eval_dataloader`: evaluation data loader


key: chroma:document
string_value:
CODE:

def prune_heads(args, model, eval_dataloader, head_mask):
    """This method shows how to prune head (remove heads weights) based on
    the head importance scores as described in Michel et al. (http://arxiv.org/abs/1905.10650)
    """
    # Try pruning and test time speedup
    # Pruning is like masking but we actually remove the masked weights
    before_time = datetime.now()
    _, _, preds, labels = compute_heads_importance(
        args, model, eval_dataloader, compute_entropy=False, compute_importance=False, head_mask=head_mask
    )
    preds = np.argmax(preds, axis=1) if args.output_mode == "classification" else np.squeeze(preds)
    score_masking = glue_compute_metrics(args.task_name, preds, labels)[args.metric_name]
    original_time = datetime.now() - before_time

    original_num_params = sum(p.numel() for p in model.parameters())
    heads_to_prune = {
        layer: (1 - head_mask[layer].long()).nonzero().squeeze().tolist() for layer in range(len(head_mask))
    }

    assert sum(len(h) for h in heads_to_prune.values()) == (1 - head_mask.long()).sum().item()
    model.prune_heads(heads_to_prune)
    pruned_num_params = sum(p.numel() for p in model.parameters())

    before_time = datetime.now()
    _, _, preds, labels = compute_heads_importance(
        args,
        model,
        eval_dataloader,
        compute_entropy=False,
        compute_importance=False,
        head_mask=None,
        actually_pruned=True,
    )
    preds = np.argmax(preds, axis=1) if args.output_mode == "classification" else np.squeeze(preds)
    score_pruning = glue_compute_metrics(args.task_name, preds, labels)[args.metric_name]
    new_time = datetime.now() - before_time

    logger.info(
        "Pruning: original num of params: %.2e, after pruning %.2e (%.1f percents)",
        original_num_params,
        pruned_num_params,
        pruned_num_params / original_num_params * 100,
    )
    logger.info("Pruning: score with masking: %f score with pruning: %f", score_masking, score_pruning)


SUMMARY:
**Concise Summary**

### Functionality
The `prune_heads` function demonstrates how to prune heads in a model based on head importance scores. It takes input arguments, models, and dataset loaders as parameters.

### Implementation Details

*   Computes head importance scores using `compute_heads_importance` with pruning enabled (`head_mask=None`) and measures the time taken.
*   Prunes heads with low importance scores by calling `model.prune_heads`.
*   Repeats the computation with pruning applied to measure the speedup achieved by pruning.

### Notable Dependencies/Requirements

*   `compute_heads_importance`
*   `glue_compute_metrics`
*   `model` (specifically, its `prune_heads` method)
*   `args` (input arguments)

**Key Method Signatures:**

*   `model.prune_heads(heads_to_prune)`

### External Dependencies:

*   `glue_compute_metrics`
*   `compute_heads_importance`


key: chroma:document
string_value:
CODE:

def main():
    parser = argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--data_dir",
        default=None,
        type=str,
        required=True,
        help="The input data dir. Should contain the .tsv files (or other data files) for the task.",
    )
    parser.add_argument(
        "--model_name_or_path",
        default=None,
        type=str,
        required=True,
        help="Path to pretrained model or model identifier from huggingface.co/models",
    )
    parser.add_argument(
        "--task_name",
        default=None,
        type=str,
        required=True,
        help="The name of the task to train selected in the list: " + ", ".join(glue_processors.keys()),
    )
    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        required=True,
        help="The output directory where the model predictions and checkpoints will be written.",
    )

    # Other parameters
    parser.add_argument(
        "--config_name",
        default="",
        type=str,
        help="Pretrained config name or path if not the same as model_name_or_path",
    )
    parser.add_argument(
        "--tokenizer_name",
        default="",
        type=str,
        help="Pretrained tokenizer name or path if not the same as model_name_or_path",
    )
    parser.add_argument(
        "--cache_dir",
        default=None,
        type=str,
        help="Where do you want to store the pre-trained models downloaded from huggingface.co",
    )
    parser.add_argument(
        "--data_subset", type=int, default=-1, help="If > 0: limit the data to a subset of data_subset instances."
    )
    parser.add_argument(
        "--overwrite_output_dir", action="store_true", help="Whether to overwrite data in output directory"
    )
    parser.add_argument(
        "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
    )

    parser.add_argument(
        "--dont_normalize_importance_by_layer", action="store_true", help="Don't normalize importance score by layers"
    )
    parser.add_argument(
        "--dont_normalize_global_importance",
        action="store_true",
        help="Don't normalize all importance scores between 0 and 1",
    )

    parser.add_argument(
        "--try_masking", action="store_true", help="Whether to try to mask head until a threshold of accuracy."
    )
    parser.add_argument(
        "--masking_threshold",
        default=0.9,
        type=float,
        help="masking threshold in term of metrics (stop masking when metric < threshold * original metric value).",
    )
    parser.add_argument(
        "--masking_amount", default=0.1, type=float, help="Amount to heads to masking at each masking step."
    )
    parser.add_argument("--metric_name", default="acc", type=str, help="Metric to use for head masking.")

    parser.add_argument(
        "--max_seq_length",
        default=128,
        type=int,
        help=(
            "The maximum total input sequence length after WordPiece tokenization. \n"
            "Sequences longer than this will be truncated, sequences shorter padded."
        ),
    )
    parser.add_argument("--batch_size", default=1, type=int, help="Batch size.")

    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--local_rank", type=int, default=-1, help="local_rank for distributed training on gpus")
    parser.add_argument("--no_cuda", action="store_true", help="Whether not to use CUDA when available")
    parser.add_argument("--server_ip", type=str, default="", help="Can be used for distant debugging.")
    parser.add_argument("--server_port", type=str, default="", help="Can be used for distant debugging.")
    args = parser.parse_args()

    if args.server_ip and args.server_port:
        # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
        import ptvsd

        print("Waiting for debugger attach")
        ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
        ptvsd.wait_for_attach()

    # Setup devices and distributed training
    if args.local_rank == -1 or args.no_cuda:
        args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()
    else:
        torch.cuda.set_device(args.local_rank)
        args.device = torch.device("cuda", args.local_rank)
        args.n_gpu = 1
        torch.distributed.init_process_group(backend="nccl")  # Initializes the distributed backend

    # Setup logging
    logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
    logger.info("device: {} n_gpu: {}, distributed: {}".format(args.device, args.n_gpu, bool(args.local_rank != -1)))
    # Set the verbosity to info of the Transformers logger (on main process only):
    if is_main_process(args.local_rank):
        transformers.utils.logging.set_verbosity_info()
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()

    # Set seeds
    set_seed(args.seed)

    # Prepare GLUE task
    args.task_name = args.task_name.lower()
    if args.task_name not in glue_processors:
        raise ValueError("Task not found: %s" % (args.task_name))
    processor = glue_processors[args.task_name]()
    args.output_mode = glue_output_modes[args.task_name]
    label_list = processor.get_labels()
    num_labels = len(label_list)

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.

    config = AutoConfig.from_pretrained(
        args.config_name if args.config_name else args.model_name_or_path,
        num_labels=num_labels,
        finetuning_task=args.task_name,
        output_attentions=True,
        cache_dir=args.cache_dir,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
        cache_dir=args.cache_dir,
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_name_or_path,
        from_tf=bool(".ckpt" in args.model_name_or_path),
        config=config,
        cache_dir=args.cache_dir,
    )

    # Distributed and parallel training
    model.to(args.device)
    if args.local_rank != -1:
        model = nn.parallel.DistributedDataParallel(
            model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True
        )
    elif args.n_gpu > 1:
        model = nn.DataParallel(model)

    # Print/save training arguments
    os.makedirs(args.output_dir, exist_ok=True)
    torch.save(args, os.path.join(args.output_dir, "run_args.bin"))
    logger.info("Training/evaluation parameters %s", args)

    # Prepare dataset for the GLUE task
    eval_dataset = GlueDataset(args, tokenizer=tokenizer, mode="dev")
    if args.data_subset > 0:
        eval_dataset = Subset(eval_dataset, list(range(min(args.data_subset, len(eval_dataset)))))
    eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
    eval_dataloader = DataLoader(
        eval_dataset, sampler=eval_sampler, batch_size=args.batch_size, collate_fn=default_data_collator
    )

    # Compute head entropy and importance score
    compute_heads_importance(args, model, eval_dataloader)

    # Try head masking (set heads to zero until the score goes under a threshole)
    # and head pruning (remove masked heads and see the effect on the network)
    if args.try_masking and args.masking_threshold > 0.0 and args.masking_threshold < 1.0:
        head_mask = mask_heads(args, model, eval_dataloader)


SUMMARY:
**Summary: Code Snippet for Training a Pretrained Model**

### Main Purpose and Functionality
This code defines a function `main()` that trains a pretrained model on a specific tasks (e.g. sentiment analysis, question answering) using the HuggingFace library.

### Key Implementation Details

*   **Model Loading**: Loads pre-trained models (`AutoModelForSequenceClassification`) using the `.from_pretrained` method.
*   **Distributed Training**: Supports distributed training using `nn.parallel.DistributedDataParallel` and `torch.distributed.init_process_group`.
*   **Task Setup**: Prepares a GLUE task by loading the processor, model configuration, tokenizer, and dataset (`GlueDataset`).
*   **Loading Pretrained Model and Tokenizer**: Loads the model, tokenizer, and config from a given path or name.

### Key Functions/Modules
```markdown

* **glue_processors**: Dictionary of processors for various GLUE tasks.
* **auto_config** and **auto_tokenizer**: Auto-configuration and auto-tokenization classes provided by HuggingFace library.
* **glue_output_modes**: Dictionary mapping task names to output modes (e.g. `{'STSDetectiveTask': 'multi_label_classification'}`).
```

### Code Structure
The code is organized into several sections:

1.  Task setup and configuration
2.  Model loading and initialization
3.  Distributed training support
4.  Task-dependent functions (e.g., `compute_heads_importance`, `mask_heads`, `prune_heads`)


key: chroma:document
string_value:
CODE:

def save_model(model, dirpath):
    # save results
    if os.path.exists(dirpath):
        if os.path.exists(os.path.join(dirpath, "config.json")) and os.path.isfile(
            os.path.join(dirpath, "config.json")
        ):
            os.remove(os.path.join(dirpath, "config.json"))
        if os.path.exists(os.path.join(dirpath, "pytorch_model.bin")) and os.path.isfile(
            os.path.join(dirpath, "pytorch_model.bin")
        ):
            os.remove(os.path.join(dirpath, "pytorch_model.bin"))
    else:
        os.makedirs(dirpath)


SUMMARY:
**Summary:**

### `save_model` Function

#### Purpose
 Saves a pre-trained PyTorch model to a specified directory.

#### Implementation Details

- Checks for existing configuration and model files in the directory, removing them if they exist.
- Creates the directory if it doesn't already exist using `os.makedirs`.
- Calls `model.save_pretrained` to save the pre-trained model.

#### Dependencies/Requirements

* `os`

### Key Method Signature

* `model.save_pretrained(dirpath)`

### Notable Design Choices

* Using absolute paths for file existence checks.
* Removing existing configuration and model files before saving.


key: chroma:document
string_value:
CODE:

def entropy(p, unlogit=False):
    """Compute the entropy of a probability distribution"""
    exponent = 2
    if unlogit:
        p = torch.pow(p, exponent)
    plogp = p * torch.log(p)
    plogp[p == 0] = 0


SUMMARY:
**Summary:**

This function calculates the entropy of a probability distribution `p` using the formula `-Σ(p \* log2(p))`. It takes into account two variants:

* The default implementation (`unlogit=False`) returns a scalar value, likely to be used as an auxiliary loss term.
* An alternative implementation (`unlogit=True`) converts the input `p` to a logarithmic scale, which can help with numerical stability.

**Key details:**

* Uses PyTorch tensor operations
* Calculates -Σ(p \* log2(p))
* Ignores zero probabilities in the calculation

**Notable dependency:**
 
* torch (PyTorch library)


key: chroma:document
string_value:
CODE:

def print_2d_tensor(tensor):
    """Print a 2D tensor"""
    logger.info("lv, h >\t" + "\t".join(f"{x + 1}" for x in range(len(tensor))))
    for row in range(len(tensor)):
        if tensor.dtype != torch.long:
            logger.info(f"layer {row + 1}:\t" + "\t".join(f"{x:.5f}" for x in tensor[row].cpu().data))
        else:


SUMMARY:
**Summary**

This function, `print_2d_tensor`, prints a 2D tensor to the console. It takes two key design choices:

*   **Format switching**: It adapts the format of tensor values based on its data type:
    *   For numeric tensors (e.g., float or double), it formats the values as floating-point numbers with five decimal places.
    *   For integer tensors (represented as long int in PyTorch), it formats the values as integers without specifying a width.
*   **Printing tensor rows**: It prints each row of the tensor, along with its corresponding layer number.

**Key Implementation Details**

*   The function relies on `torch.log` for logging purposes, specifically using `logger.info()` to log timestamp (`lv`) and height (`h`) information about the tensor.
*   Tensor values are converted to CPU before printing due to PyTorch's data manipulation on GPU.

**Notable Dependencies or Requirements**

*   The function uses:
    *   `torch` library for working with tensors
    *   A logger instance (possibly part of a custom logging system) for timestamp and log information


key: chroma:document
string_value:
CODE:

def compute_heads_importance(
    args, model, eval_dataloader, compute_entropy=True, compute_importance=True, head_mask=None, actually_pruned=False
):
    """This method shows how to compute:
    - head attention entropy
    - head importance scores according to http://arxiv.org/abs/1905.10650
    """
    # Prepare our tensors
    n_layers, n_heads = model.config.num_hidden_layers, model.config.num_attention_heads
    head_importance = torch.zeros(n_layers, n_heads).to(args.device)
    attn_entropy = torch.zeros(n_layers, n_heads).to(args.device)

    if head_mask is None:
        head_mask = torch.ones(n_layers, n_heads).to(args.device)

    head_mask.requires_grad_(requires_grad=True)
    # If actually pruned attention multi-head, set head mask to None to avoid shape mismatch
    if actually_pruned:
        head_mask = None

    tot_tokens = 0.0
    total_loss = 0.0
    for step, inputs in enumerate(tqdm(eval_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])):
        inputs = tuple(t.to(args.device) for t in inputs)
        (input_ids,) = inputs

        # Do a forward pass (not with torch.no_grad() since we need gradients for importance score - see below)
        outputs = model(input_ids, labels=input_ids, head_mask=head_mask)
        #  (loss), lm_logits, presents, (all hidden_states), (attentions)
        loss, _, all_attentions = (
            outputs[0],
            outputs[1],
            outputs[-1],
        )  # Loss and logits are the first, attention the last
        loss.backward()  # Backpropagate to populate the gradients in the head mask
        total_loss += loss.detach().cpu().numpy()
        if compute_entropy:
            for layer, attn in enumerate(all_attentions):
                masked_entropy = entropy(attn.detach(), True)
                attn_entropy[layer] += masked_entropy.sum(-1).sum(0).sum(0).detach()

        if compute_importance:
            head_importance += head_mask.grad.abs().detach()
        tot_tokens += torch.ones_like(input_ids).float().detach().sum().data

    # Normalize
    attn_entropy /= tot_tokens
    head_importance /= tot_tokens
    # Layerwise importance normalization
    if not args.dont_normalize_importance_by_layer:
        exponent = 2
        norm_by_layer = torch.pow(torch.pow(head_importance, exponent).sum(-1), 1 / exponent)
        head_importance /= norm_by_layer.unsqueeze(-1) + 1e-20

    if not args.dont_normalize_global_importance:
        head_importance = (head_importance - head_importance.min()) / (head_importance.max() - head_importance.min())

    # Print matrices
    if compute_entropy:
        logger.info("Attention entropies")
        print_2d_tensor(attn_entropy)
    if compute_importance:
        logger.info("Head importance scores")
        print_2d_tensor(head_importance)
    logger.info("Head ranked by importance scores")
    head_ranks = torch.zeros(head_importance.numel(), dtype=torch.long, device=args.device)
    head_ranks[head_importance.view(-1).sort(descending=True)[1]] = torch.arange(
        head_importance.numel(), device=args.device
    )
    head_ranks = head_ranks.view_as(head_importance)
    print_2d_tensor(head_ranks)


SUMMARY:
**Concise Summary**

This function `compute_heads_importance` computes the importance scores of attention heads in a Transformer model. The main purpose is to calculate:

1. Head Attention Entropy
2. Head Importance Scores according to [1]
3. Total Loss

Key Implementation Details:

* Computes entropy and importance scores for each head
* Masks output with attention mask to avoid shape mismatch
* Does forward pass, backpropagates, computes gradients in attention mask
* Normalizes importance scores layer-wise and globally if enabled

Notable Dependencies and Requirements:

* `model` instance with hidden layers and attention heads configuration
* `args.device` for tensor placement
* External dependencies:
	+ `torch`
	+ ` entropy` function ( implementation not shown)

References:
[1] http://arxiv.org/abs/1905.10650


key: chroma:document
string_value:
CODE:

def mask_heads(args, model, eval_dataloader):
    """This method shows how to mask head (set some heads to zero), to test the effect on the network,
    based on the head importance scores, as described in Michel et al. (http://arxiv.org/abs/1905.10650)
    """
    _, head_importance, loss = compute_heads_importance(args, model, eval_dataloader, compute_entropy=False)
    original_score = 1 / loss  # instead of downsteam score use the LM loss
    logger.info("Pruning: original score: %f, threshold: %f", original_score, original_score * args.masking_threshold)

    new_head_mask = torch.ones_like(head_importance)
    num_to_mask = max(1, int(new_head_mask.numel() * args.masking_amount))

    current_score = original_score
    while current_score >= original_score * args.masking_threshold:
        head_mask = new_head_mask.clone().detach()  # save current head mask
        # heads from least important to most - keep only not-masked heads
        head_importance[head_mask == 0.0] = float("Inf")
        current_heads_to_mask = head_importance.view(-1).sort()[1]

        if len(current_heads_to_mask) <= num_to_mask:
            print("BREAK BY num_to_mask")
            break

        # mask heads
        current_heads_to_mask = current_heads_to_mask[:num_to_mask]
        logger.info("Heads to mask: %s", str(current_heads_to_mask.tolist()))
        new_head_mask = new_head_mask.view(-1)
        new_head_mask[current_heads_to_mask] = 0.0
        new_head_mask = new_head_mask.view_as(head_mask)
        new_head_mask = new_head_mask.clone().detach()
        print_2d_tensor(new_head_mask)

        # Compute metric and head importance again
        _, head_importance, loss = compute_heads_importance(
            args, model, eval_dataloader, compute_entropy=False, head_mask=new_head_mask
        )
        current_score = 1 / loss
        logger.info(
            "Masking: current score: %f, remaining heads %d (%.1f percents)",
            current_score,
            new_head_mask.sum(),
            new_head_mask.sum() / new_head_mask.numel() * 100,
        )

    logger.info("Final head mask")
    print_2d_tensor(head_mask)
    np.save(os.path.join(args.output_dir, "head_mask.npy"), head_mask.detach().cpu().numpy())



SUMMARY:
Here's a concise summary of the code chunk:

### Function Name: `mask_heads`

#### Purpose
The function masks heads in a neural network model based on their importance scores, as described in Michel et al. (2019).

#### Functionality

1. Compute head importance scores using `compute_heads_importance` with `compute_entropy=False`.
2. Initialize a masking threshold and apply it to the original score.
3. Iterate until the current score meets or drops below the threshold:
	* Save the current head mask.
	* Mask heads with lower importance scores.
	* Recompute head importance scores using the new head mask.
	* Update the current score.

#### Implementation Details

* The function uses a threshold-based approach to determine which heads to mask.
* It employs a greedy search strategy by iterating from least important to most important heads.

#### Dependencies and Requirements
```python
- `compute_heads_importance` (a separate function)
- `eval_dataloader`
- `args.masking_threshold` (threshold value)
- `args.masking_amount` (ratio of heads to mask)
```

#### Output

* The final masked head mask is saved as an RGB tensor in `npy` file format.
* A 2D tensor representation is printed using the `print_2d_tensor` function.


key: chroma:document
string_value:
CODE:

def prune_heads(args, model, eval_dataloader, head_mask):
    """This method shows how to prune head (remove heads weights) based on
    the head importance scores as described in Michel et al. (http://arxiv.org/abs/1905.10650)
    """
    # Try pruning and test time speedup
    # Pruning is like masking but we actually remove the masked weights
    before_time = datetime.now()
    _, _, loss = compute_heads_importance(
        args, model, eval_dataloader, compute_entropy=False, compute_importance=False, head_mask=head_mask
    )
    score_masking = 1 / loss
    original_time = datetime.now() - before_time

    original_num_params = sum(p.numel() for p in model.parameters())
    heads_to_prune = {
        layer: (1 - head_mask[layer].long()).nonzero().squeeze().tolist() for layer in range(len(head_mask))
    }

    for k, v in heads_to_prune.items():
        if isinstance(v, int):
            heads_to_prune[k] = [
                v,
            ]

    assert sum(len(h) for h in heads_to_prune.values()) == (1 - head_mask.long()).sum().item()
    model.prune_heads(heads_to_prune)
    pruned_num_params = sum(p.numel() for p in model.parameters())

    before_time = datetime.now()
    _, _, loss = compute_heads_importance(
        args,
        model,
        eval_dataloader,
        compute_entropy=False,
        compute_importance=False,
        head_mask=None,
        actually_pruned=True,
    )

    score_pruning = 1 / loss
    new_time = datetime.now() - before_time

    logger.info(
        "Pruning: original num of params: %.2e, after pruning %.2e (%.1f percents)",
        original_num_params,
        pruned_num_params,
        pruned_num_params / original_num_params * 100,
    )
    logger.info("Pruning: score with masking: %f score with pruning: %f", score_masking, score_pruning)
    logger.info("Pruning: speed ratio (original timing / new timing): %f percents", original_time / new_time * 100)


SUMMARY:
**Function Summary**

`prune_heads` method:

* Prunes head weights based on importance scores using the Michel et al. (http://arxiv.org/abs/1905.10650) technique.

### Main Purpose and Functionality

This function prunes head weights in a model by removing unimportant heads and updates the model's parameters count. It uses two approaches:
1. Score-based pruning, where important heads are retained.
2. Time-based pruning, which measures the speedup after pruning.

The function takes five inputs:

* `args`: hyperparameters
* `model`: neural network model to be pruned
* `eval_dataloader`: evaluation data loader
* `head_mask`: head importance scores

### Key Implementation Details

1. **Before and After Pruning**:
	* Computes head importance scores using `compute_heads_importance`.
	* Prints original number of parameters, score with masking, and score with pruning.
2. **Pruning**:
	* Calls `model.prune_heads` to prune unimportant heads.

### Notable Dependencies or Requirements

1. `torch`: PyTorch library for neural networks.
2. `datetime`: Python's built-in datetime module for timing measurements.
3. `logger`: Log management system (e.g., tensorboard, logging).
4. `compute_heads_importance` function: internally computes head importance scores.

### Model and Hyperparameters

* Model is a PyTorch neural network model (`model`).
* Hyperparameters are stored in `args`.

### Time Complexity Analysis

The time complexity of this function is dominated by the two calls to `compute_heads_importance`, which have a computational complexity of O(N), where N is the number of heads.

### Security Considerations

This code does not appear to introduce any security vulnerabilities. However, proper input validation and sanitization should always be applied when working with untrusted data.


key: chroma:document
string_value:
CODE:

def main():
    parser = argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--data_dir",
        default=None,
        type=str,
        required=True,
        help="The input data dir. Should contain the .tsv files (or other data files) for the task.",
    )
    parser.add_argument(
        "--model_name_or_path",
        default=None,
        type=str,
        required=True,
        help="Path to pretrained model or model identifier from huggingface.co/models",
    )
    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        required=True,
        help="The output directory where the model predictions and checkpoints will be written.",
    )

    # Other parameters
    parser.add_argument(
        "--config_name",
        default="",
        type=str,
        help="Pretrained config name or path if not the same as model_name_or_path",
    )
    parser.add_argument(
        "--tokenizer_name",
        default="",
        type=str,
        help="Pretrained tokenizer name or path if not the same as model_name_or_path",
    )
    parser.add_argument(
        "--cache_dir",
        default=None,
        type=str,
        help="Where do you want to store the pre-trained models downloaded from s3",
    )
    parser.add_argument(
        "--data_subset", type=int, default=-1, help="If > 0: limit the data to a subset of data_subset instances."
    )
    parser.add_argument(
        "--overwrite_output_dir", action="store_true", help="Whether to overwrite data in output directory"
    )
    parser.add_argument(
        "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
    )

    parser.add_argument(
        "--dont_normalize_importance_by_layer", action="store_true", help="Don't normalize importance score by layers"
    )
    parser.add_argument(
        "--dont_normalize_global_importance",
        action="store_true",
        help="Don't normalize all importance scores between 0 and 1",
    )

    parser.add_argument(
        "--try_masking", action="store_true", help="Whether to try to mask head until a threshold of accuracy."
    )
    parser.add_argument(
        "--masking_threshold",
        default=0.9,
        type=float,
        help="masking threshold in term of metrics (stop masking when metric < threshold * original metric value).",
    )
    parser.add_argument(
        "--masking_amount", default=0.1, type=float, help="Amount to heads to masking at each masking step."
    )
    parser.add_argument("--metric_name", default="acc", type=str, help="Metric to use for head masking.")

    parser.add_argument(
        "--max_seq_length",
        default=128,
        type=int,
        help=(
            "The maximum total input sequence length after WordPiece tokenization. \n"
            "Sequences longer than this will be truncated, sequences shorter padded."
        ),
    )
    parser.add_argument("--batch_size", default=1, type=int, help="Batch size.")

    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--local_rank", type=int, default=-1, help="local_rank for distributed training on gpus")
    parser.add_argument("--no_cuda", action="store_true", help="Whether not to use CUDA when available")
    parser.add_argument("--server_ip", type=str, default="", help="Can be used for distant debugging.")
    parser.add_argument("--server_port", type=str, default="", help="Can be used for distant debugging.")
    args = parser.parse_args()

    if args.server_ip and args.server_port:
        # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
        import ptvsd

        print("Waiting for debugger attach")
        ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
        ptvsd.wait_for_attach()

    # Setup devices and distributed training
    if args.local_rank == -1 or args.no_cuda:
        args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()
    else:
        torch.cuda.set_device(args.local_rank)
        args.device = torch.device("cuda", args.local_rank)
        args.n_gpu = 1
        torch.distributed.init_process_group(backend="nccl")  # Initializes the distributed backend

    # Setup logging
    logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
    logger.info("device: {} n_gpu: {}, distributed: {}".format(args.device, args.n_gpu, bool(args.local_rank != -1)))

    model = GPT2LMHeadModel.from_pretrained(args.model_name_or_path)

    # Distributed and parallel training
    model.to(args.device)
    if args.local_rank != -1:
        model = nn.parallel.DistributedDataParallel(
            model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True
        )
    elif args.n_gpu > 1:
        model = nn.DataParallel(model)

    # Print/save training arguments
    os.makedirs(args.output_dir, exist_ok=True)
    torch.save(args, os.path.join(args.output_dir, "run_args.bin"))
    logger.info("Training/evaluation parameters %s", args)

    # Prepare dataset
    numpy_data = np.concatenate(
        [
            np.loadtxt(args.data_dir, dtype=np.int64),
        ]
    )
    train_tensor_dataset = (torch.from_numpy(numpy_data),)
    train_data = TensorDataset(*train_tensor_dataset)
    train_sampler = RandomSampler(train_data)
    eval_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.batch_size)

    # Compute head entropy and importance score
    compute_heads_importance(args, model, eval_dataloader)

    # Try head masking (set heads to zero until the score goes under a threshole)
    # and head pruning (remove masked heads and see the effect on the network)
    if args.try_masking and args.masking_threshold > 0.0 and args.masking_threshold < 1.0:
        head_mask = mask_heads(args, model, eval_dataloader)


SUMMARY:
**Function Code Summary**

This function is the main entry point for training and evaluation of a GPT2LMHeadModel. It sets up command-line arguments using argparse, initializes logging, and prepares the dataset.

**Key Functions and Details:**

1.  **`parse_args()`**: Parse command-line arguments to configure model, data, and optimization settings.
2.  **`setup devices and training`**: Determine device (GPU or CPU), set CUDA if available, and setup distributed training using NCCL backend.
3.  **`model initialization and loading`**: Load pre-trained GPT2LMHeadModel from `huggingface.co/models`.
4.  **`dataset preparation`**: Create datasets for training and evaluation, including data concatenation, tokenization, and batching.
5.  **`head masking and pruning`**: Implement try masking (masking heads until a threshold of accuracy) and head pruning (removing masked heads) using custom functions `compute_heads_importance`, `.mask_heads`, and `prune_heads`.

**Notable Parameters and Dependencies:**

*   `--model_name_or_path`: Pre-trained model name or path to load.
*   `--config_name` and `--tokenizer_name`: Config and tokenizer names/path for fine-tuning.
*   `--data_dir` and `--output_dir`: Input data directory and output directory for save predictions/checkpoints.
*   Importance mask threshold (`--masking_threshold`) and masking amount (`--masking_amount`) configure the head masking strategy.
*   Datasets created using PyTorch DataLoader, torch from numpy, and torch DDP (DataParallel).

**References:**

*   GPT2LMHeadModel (from `huggingface.co/models`)
*   argparse
*   logging (Python standard library)
*   PyTorch libraries (`torchcuda`, `torch.nn.parallel.DistributedDataParallel`, etc.)
*   numpy


key: chroma:document
string_value:
CODE:

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:


SUMMARY:
**Summary**

This function, `set_seed`, initializes the seeds for various random number generators:
- PyTorch: `torch.manual_seed` and `torch.cuda.manual_seed_all`
- NumPy: `np.random.seed`
- Python's built-in `random`: `random.seed`

The primary purpose is to reproduce the same sequence of random numbers when running an experiment multiple times.

**Key Implementation Details**

* Uses consistent seed value from input arguments (`args.seed`)
* Initializes PyTorch's GPU seeds if `num_gpus > 0` (`torch.cuda.manual_seed_all`)
* Uses NumPy and Python's built-in `random` modules for initialization (`np.random.seed`, `random.seed`)

**Notable Dependencies or Requirements**

- Args (`args`: an object containing the seed value and other arguments)
- PyTorch (`torch`)
- NumPy (`np`)
- Python's built-in `random` module


key: chroma:document
string_value:
CODE:

def get_wanted_result(result):
    if "spearmanr" in result:
        print_result = result["spearmanr"]
    elif "f1" in result:
        print_result = result["f1"]
    elif "mcc" in result:
        print_result = result["mcc"]
    elif "acc" in result:
        print_result = result["acc"]
    else:
        raise ValueError("Primary metric unclear in the results")


SUMMARY:
**Summary:**

This function, `get_wanted_result`, retrieves and returns a specific evaluation metric from a given result dictionary. It supports extracting "spearmanr", "f1", "mcc", and "acc" metrics from the input dictionary.

**Key Implementation Details:**

* Uses conditional statements to determine which metric to extract based on its presence in the result.
* Prints an error message if none of the supported metrics are found in the results.

**Dependencies/Requirements:**

* None mentioned, but assumes access to a dictionary containing evaluation metric data (e.g., result = {"spearmanr": ...}).

**Method Signature and Interface:** 
```python
def get_wanted_result(result) -> any:
    ...
```
This function takes an input `result` without specifying its type, and returns the extracted evaluation metric as any compatible type (`any`).


key: chroma:document
string_value:
CODE:

def train(args, train_dataset, model, tokenizer, train_highway=False):
    """Train the model"""
    if args.local_rank in [-1, 0]:
        tb_writer = SummaryWriter()

    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)

    if args.max_steps > 0:
        t_total = args.max_steps
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
    else:
        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs

    # Prepare optimizer and schedule (linear warmup and decay)
    no_decay = ["bias", "LayerNorm.weight"]
    if train_highway:
        optimizer_grouped_parameters = [
            {
                "params": [
                    p
                    for n, p in model.named_parameters()
                    if ("highway" in n) and (not any(nd in n for nd in no_decay))
                ],
                "weight_decay": args.weight_decay,
            },
            {
                "params": [
                    p for n, p in model.named_parameters() if ("highway" in n) and (any(nd in n for nd in no_decay))
                ],
                "weight_decay": 0.0,
            },
        ]
    else:
        optimizer_grouped_parameters = [
            {
                "params": [
                    p
                    for n, p in model.named_parameters()
                    if ("highway" not in n) and (not any(nd in n for nd in no_decay))
                ],
                "weight_decay": args.weight_decay,
            },
            {
                "params": [
                    p
                    for n, p in model.named_parameters()
                    if ("highway" not in n) and (any(nd in n for nd in no_decay))
                ],
                "weight_decay": 0.0,
            },
        ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
    )
    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = nn.DataParallel(model)

    # Distributed training (should be after apex fp16 initialization)
    if args.local_rank != -1:
        model = nn.parallel.DistributedDataParallel(
            model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True
        )

    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
    logger.info(
        "  Total train batch size (w. parallel, distributed & accumulation) = %d",
        args.train_batch_size
        * args.gradient_accumulation_steps
        * (torch.distributed.get_world_size() if args.local_rank != -1 else 1),
    )
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    global_step = 0
    tr_loss, logging_loss = 0.0, 0.0
    model.zero_grad()
    train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
    set_seed(args)  # Added here for reproducibility (even between python 2 and 3)
    for _ in train_iterator:
        epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
        for step, batch in enumerate(epoch_iterator):
            model.train()
            batch = tuple(t.to(args.device) for t in batch)
            inputs = {"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
            if args.model_type != "distilbert":
                inputs["token_type_ids"] = (
                    batch[2] if args.model_type in ["bert", "xlnet"] else None
                )  # XLM, DistilBERT and RoBERTa don't use segment_ids
            inputs["train_highway"] = train_highway
            outputs = model(**inputs)
            loss = outputs[0]  # model outputs are always tuple in transformers (see doc)

            if args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

            tr_loss += loss.item()
            if (step + 1) % args.gradient_accumulation_steps == 0:
                if args.fp16:
                    nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                else:
                    nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                optimizer.step()
                scheduler.step()  # Update learning rate schedule
                model.zero_grad()
                global_step += 1

                if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    # Log metrics
                    if (
                        args.local_rank == -1 and args.evaluate_during_training
                    ):  # Only evaluate when single GPU otherwise metrics may not average well
                        results = evaluate(args, model, tokenizer)
                        for key, value in results.items():
                            tb_writer.add_scalar("eval_{}".format(key), value, global_step)
                    tb_writer.add_scalar("lr", scheduler.get_lr()[0], global_step)
                    tb_writer.add_scalar("loss", (tr_loss - logging_loss) / args.logging_steps, global_step)
                    logging_loss = tr_loss

                if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
                    # Save model checkpoint
                    output_dir = os.path.join(args.output_dir, "checkpoint-{}".format(global_step))
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    model_to_save = (
                        model.module if hasattr(model, "module") else model
                    )  # Take care of distributed/parallel training
                    model_to_save.save_pretrained(output_dir)
                    torch.save(args, os.path.join(output_dir, "training_args.bin"))
                    logger.info("Saving model checkpoint to %s", output_dir)

            if args.max_steps > 0 and global_step > args.max_steps:
                epoch_iterator.close()
                break
        if args.max_steps > 0 and global_step > args.max_steps:
            train_iterator.close()
            break

    if args.local_rank in [-1, 0]:
        tb_writer.close()



SUMMARY:
**Summary:**

This is the training function for a deep learning model. It takes input from various parameters and data structures, including:

* `args`: command line arguments
* `train_dataset`: dataset to train on
* `model`: pre-trained model architecture
* `tokenizer`: tokenizer for tokenizing text inputs
* `train_highway`: flag indicating whether to use the Highway connection layer in the model

**Primary Purpose:**

The primary purpose of this function is to train a deep learning model using a given dataset and hyperparameters.

**Key Implementation Details:**

1. **Model Initialization:** The model is initialized with certain parameters, including weight decay levels for different layers.
2. **Optimizer and Schedule Initialization:** The optimizer (AdamW) and schedule (linear warmup and decay) are set up according to the `args` object.
3. **Gradient Accumulation:** Gradient accumulation steps are handled based on the number of GPUs and gradient accumulation steps.
4. **Multi-GPU Training:** Trained models can be parallelized across multiple GPUs using the `DataParallel` module.

**Notable Dependencies:**

1. `torch`: PyTorch library
2. `transformers`: Transformer library for natural language processing tasks
3. Various other libraries and modules for data loading, evaluation, logging, and saving/checking models

**Function flow:**

The function starts by initializing the model and optimizer, then enters a loop that trains the model on the given dataset with gradient accumulation steps. The weights are updated after each batch using the optimizer and scheduler. Logging and saving checkpoints occur at regular intervals.

**Return Values:**

The function returns two values:

* `global_step`: iteration number during training
* `tr_loss / global_step`: average loss per step


key: chroma:document
string_value:
CODE:

def evaluate(args, model, tokenizer, prefix="", output_layer=-1, eval_highway=False):
    # Loop to handle MNLI double evaluation (matched, mis-matched)
    eval_task_names = ("mnli", "mnli-mm") if args.task_name == "mnli" else (args.task_name,)
    eval_outputs_dirs = (args.output_dir, args.output_dir + "-MM") if args.task_name == "mnli" else (args.output_dir,)

    results = {}
    for eval_task, eval_output_dir in zip(eval_task_names, eval_outputs_dirs):
        eval_dataset = load_and_cache_examples(args, eval_task, tokenizer, evaluate=True)

        if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
            os.makedirs(eval_output_dir)

        args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
        # Note that DistributedSampler samples randomly
        eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
        eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)

        # multi-gpu eval
        if args.n_gpu > 1:
            model = nn.DataParallel(model)

        # Eval!
        logger.info("***** Running evaluation {} *****".format(prefix))
        logger.info("  Num examples = %d", len(eval_dataset))
        logger.info("  Batch size = %d", args.eval_batch_size)
        eval_loss = 0.0
        nb_eval_steps = 0
        preds = None
        out_label_ids = None
        exit_layer_counter = {(i + 1): 0 for i in range(model.num_layers)}
        st = time.time()
        for batch in tqdm(eval_dataloader, desc="Evaluating"):
            model.eval()
            batch = tuple(t.to(args.device) for t in batch)

            with torch.no_grad():
                inputs = {"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
                if args.model_type != "distilbert":
                    inputs["token_type_ids"] = (
                        batch[2] if args.model_type in ["bert", "xlnet"] else None
                    )  # XLM, DistilBERT and RoBERTa don't use segment_ids
                if output_layer >= 0:
                    inputs["output_layer"] = output_layer
                outputs = model(**inputs)
                if eval_highway:
                    exit_layer_counter[outputs[-1]] += 1
                tmp_eval_loss, logits = outputs[:2]

                eval_loss += tmp_eval_loss.mean().item()
            nb_eval_steps += 1
            if preds is None:
                preds = logits.detach().cpu().numpy()
                out_label_ids = inputs["labels"].detach().cpu().numpy()
            else:
                preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
                out_label_ids = np.append(out_label_ids, inputs["labels"].detach().cpu().numpy(), axis=0)
        eval_time = time.time() - st
        logger.info("Eval time: {}".format(eval_time))

        eval_loss = eval_loss / nb_eval_steps
        if args.output_mode == "classification":
            preds = np.argmax(preds, axis=1)
        elif args.output_mode == "regression":
            preds = np.squeeze(preds)
        result = compute_metrics(eval_task, preds, out_label_ids)
        results.update(result)

        if eval_highway:
            logger.info("Exit layer counter: {}".format(exit_layer_counter))
            actual_cost = sum([l * c for l, c in exit_layer_counter.items()])
            full_cost = len(eval_dataloader) * model.num_layers
            logger.info("Expected saving: {}".format(actual_cost / full_cost))
            if args.early_exit_entropy >= 0:
                save_fname = (
                    args.plot_data_dir
                    + "/"
                    + args.model_name_or_path[2:]
                    + "/entropy_{}.npy".format(args.early_exit_entropy)
                )
                if not os.path.exists(os.path.dirname(save_fname)):
                    os.makedirs(os.path.dirname(save_fname))
                print_result = get_wanted_result(result)
                np.save(save_fname, np.array([exit_layer_counter, eval_time, actual_cost / full_cost, print_result]))
                logger.info("Entropy={}\tResult={:.2f}".format(args.early_exit_entropy, 100 * print_result))

        output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt")
        with open(output_eval_file, "w") as writer:
            logger.info("***** Eval results {} *****".format(prefix))
            for key in sorted(result.keys()):
                logger.info("  %s = %s", key, str(result[key]))
                writer.write("%s = %s\n" % (key, str(result[key])))



SUMMARY:
**Function Summary: Evaluate**

This function evaluates the performance of a model on multiple datasets using different evaluation tasks. It takes in various arguments to customize the evaluation process.

**Main Purpose and Functionality**

The primary purpose of this function is to perform supervised learning tasks such as MNLI (Multi-National Language Impact) evaluations, which involve predicting labels or classifications for input data samples.

**Key Implementation Details**

1. **Evaluation Loop**: The function loops over multiple evaluation tasks (MLNI) with different output directories.
2. **Model Evaluation**: It evaluates the model's performance on each task by iterating over batches of input data and computing loss and predictions.
3. **Data Preparation**: The function prepares the dataset, generates batched inputs, and samples batches using a sequential or distributed sampler, depending on the number of GPUs.

**Notable Dependencies or Requirements**

- `model`: An instance of a neural network model.
- `tokenizer`: An object that tokenizes input data.
- `output_layer` (optional): A specific output layer index for evaluation.
- `eval_highway` (optional): Enables "exit layer counter" tracking.

``` Plain Text Syntax


def evaluate(
    args,
    model,
    tokenizer,
    prefix="",
    output_layer=-1,
    eval_highway=False
)

```


key: chroma:document
string_value:
CODE:

def load_and_cache_examples(args, task, tokenizer, evaluate=False):
    if args.local_rank not in [-1, 0] and not evaluate:
        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache

    processor = processors[task]()
    output_mode = output_modes[task]
    # Load data features from cache or dataset file
    cached_features_file = os.path.join(
        args.data_dir,
        "cached_{}_{}_{}_{}".format(
            "dev" if evaluate else "train",
            list(filter(None, args.model_name_or_path.split("/"))).pop(),
            str(args.max_seq_length),
            str(task),
        ),
    )
    if os.path.exists(cached_features_file) and not args.overwrite_cache:
        logger.info("Loading features from cached file %s", cached_features_file)
        features = torch.load(cached_features_file)
    else:
        logger.info("Creating features from dataset file at %s", args.data_dir)
        label_list = processor.get_labels()
        if task in ["mnli", "mnli-mm"] and args.model_type in ["roberta"]:
            # HACK(label indices are swapped in RoBERTa pretrained model)
            label_list[1], label_list[2] = label_list[2], label_list[1]
        examples = (
            processor.get_dev_examples(args.data_dir) if evaluate else processor.get_train_examples(args.data_dir)
        )
        features = convert_examples_to_features(
            examples,
            tokenizer,
            label_list=label_list,
            max_length=args.max_seq_length,
            output_mode=output_mode,
        )
        if args.local_rank in [-1, 0]:
            logger.info("Saving features into cached file %s", cached_features_file)
            torch.save(features, cached_features_file)

    if args.local_rank == 0 and not evaluate:
        torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache

    # Convert to Tensors and build dataset
    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long)

    if features[0].token_type_ids is None:
        # For RoBERTa (a potential bug!)
        all_token_type_ids = torch.tensor([[0] * args.max_seq_length for f in features], dtype=torch.long)
    else:
        all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long)
    if output_mode == "classification":
        all_labels = torch.tensor([f.label for f in features], dtype=torch.long)
    elif output_mode == "regression":
        all_labels = torch.tensor([f.label for f in features], dtype=torch.float)

    dataset = TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_labels)


SUMMARY:
**Summary:**

The `load_and_cache_examples` function loads and caches dataset examples for a given task using a RoBERTa-like tokenizer. It:

*   Ensures only one process in distributed training processes the dataset by checking the local_rank and evaluate flag
*   Loads or creates features from cache or dataset file, depending on whether it already exists and whether overwrite_cache is set to False
*   Conversions examples to features using the `convert_examples_to_features` function, which takes into account specific tokenization and classification/typing strategies
*   Saves features to a cached file if necessary (on local_rank 0 and not in evaluate mode)
*   Finalizes dataset creation by building tensors for input IDs, attention masks, token type IDs, and labels using TensorDataset

**Key Implementation Details:**

*   `processor.get_labels()`: Returns label list based on the task
*   `convert_examples_to_features(examples, tokenizer, max_length=...)`:
    *   Processes examples to create tokens
    *   Uses output_mode classification or regression for correct labelling
*   `torch.distributed.barrier()`:
    +   Ensures first process in training handles dataset loading
    - `output_modes`: Set classification or regression


**Notable Dependencies and Requirements:**

*   Torch
*   Hugging Face processors (e.g. "RoBERTaTokenizer")
*   TorchVision

**Interfaces Implemented/ Inherited:**

*   `TensorDataset`: dataset class returned by function

Please note that for full clarity, consider implementing some custom data type as defined inside the convert_examples_to_features().


key: chroma:document
string_value:
CODE:

def main():
    parser = argparse.ArgumentParser()

    # Required parameters
    parser.add_argument(
        "--data_dir",
        default=None,
        type=str,
        required=True,
        help="The input data dir. Should contain the .tsv files (or other data files) for the task.",
    )
    parser.add_argument(
        "--model_type",
        default=None,
        type=str,
        required=True,
        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
    )
    parser.add_argument(
        "--model_name_or_path",
        default=None,
        type=str,
        required=True,
        help="Path to pre-trained model or shortcut name.",
    )
    parser.add_argument(
        "--task_name",
        default=None,
        type=str,
        required=True,
        help="The name of the task to train selected in the list: " + ", ".join(processors.keys()),
    )
    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        required=True,
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument(
        "--plot_data_dir",
        default="./plotting/",
        type=str,
        required=False,
        help="The directory to store data for plotting figures.",
    )

    # Other parameters
    parser.add_argument(
        "--config_name", default="", type=str, help="Pretrained config name or path if not the same as model_name"
    )
    parser.add_argument(
        "--tokenizer_name",
        default="",
        type=str,
        help="Pretrained tokenizer name or path if not the same as model_name",
    )
    parser.add_argument(
        "--cache_dir",
        default="",
        type=str,
        help="Where do you want to store the pre-trained models downloaded from huggingface.co",
    )
    parser.add_argument(
        "--max_seq_length",
        default=128,
        type=int,
        help=(
            "The maximum total input sequence length after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded."
        ),
    )
    parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
    parser.add_argument("--do_eval", action="store_true", help="Whether to run eval on the dev set.")
    parser.add_argument(
        "--evaluate_during_training", action="store_true", help="Rul evaluation during training at each logging step."
    )
    parser.add_argument(
        "--do_lower_case", action="store_true", help="Set this flag if you are using an uncased model."
    )
    parser.add_argument("--eval_each_highway", action="store_true", help="Set this flag to evaluate each highway.")
    parser.add_argument(
        "--eval_after_first_stage",
        action="store_true",
        help="Set this flag to evaluate after training only bert (not highway).",
    )
    parser.add_argument("--eval_highway", action="store_true", help="Set this flag if it's evaluating highway models")

    parser.add_argument("--per_gpu_train_batch_size", default=8, type=int, help="Batch size per GPU/CPU for training.")
    parser.add_argument(
        "--per_gpu_eval_batch_size", default=8, type=int, help="Batch size per GPU/CPU for evaluation."
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
    parser.add_argument("--weight_decay", default=0.0, type=float, help="Weight decay if we apply some.")
    parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
    parser.add_argument(
        "--num_train_epochs", default=3.0, type=float, help="Total number of training epochs to perform."
    )
    parser.add_argument(
        "--max_steps",
        default=-1,
        type=int,
        help="If > 0: set total number of training steps to perform. Override num_train_epochs.",
    )
    parser.add_argument("--warmup_steps", default=0, type=int, help="Linear warmup over warmup_steps.")
    parser.add_argument("--early_exit_entropy", default=-1, type=float, help="Entropy threshold for early exit.")

    parser.add_argument("--logging_steps", type=int, default=50, help="Log every X updates steps.")
    parser.add_argument("--save_steps", type=int, default=50, help="Save checkpoint every X updates steps.")
    parser.add_argument(
        "--eval_all_checkpoints",
        action="store_true",
        help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number",
    )
    parser.add_argument("--no_cuda", action="store_true", help="Avoid using CUDA when available")
    parser.add_argument(
        "--overwrite_output_dir", action="store_true", help="Overwrite the content of the output directory"
    )
    parser.add_argument(
        "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
    )
    parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")

    parser.add_argument(
        "--fp16",
        action="store_true",
        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",
    )
    parser.add_argument(
        "--fp16_opt_level",
        type=str,
        default="O1",
        help=(
            "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']. "
            "See details at https://nvidia.github.io/apex/amp.html"
        ),
    )
    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
    parser.add_argument("--server_ip", type=str, default="", help="For distant debugging.")
    parser.add_argument("--server_port", type=str, default="", help="For distant debugging.")
    args = parser.parse_args()

    if (
        os.path.exists(args.output_dir)
        and os.listdir(args.output_dir)
        and args.do_train
        and not args.overwrite_output_dir
    ):
        raise ValueError(
            "Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(
                args.output_dir
            )
        )

    # Setup distant debugging if needed
    if args.server_ip and args.server_port:
        # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
        import ptvsd

        print("Waiting for debugger attach")
        ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
        ptvsd.wait_for_attach()

    # Setup CUDA, GPU & distributed training
    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        args.n_gpu = torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of synchronizing nodes/GPUs
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend="nccl")
        args.n_gpu = 1
    args.device = device

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
        args.local_rank,
        device,
        args.n_gpu,
        bool(args.local_rank != -1),
        args.fp16,
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    if is_main_process(args.local_rank):
        transformers.utils.logging.set_verbosity_info()
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()
    # Set seed
    set_seed(args)

    # Prepare GLUE task
    args.task_name = args.task_name.lower()
    if args.task_name not in processors:
        raise ValueError("Task not found: %s" % (args.task_name))
    processor = processors[args.task_name]()
    args.output_mode = output_modes[args.task_name]
    label_list = processor.get_labels()
    num_labels = len(label_list)

    # Load pretrained model and tokenizer
    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab

    args.model_type = args.model_type.lower()
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config = config_class.from_pretrained(
        args.config_name if args.config_name else args.model_name_or_path,
        num_labels=num_labels,
        finetuning_task=args.task_name,
        cache_dir=args.cache_dir if args.cache_dir else None,
    )
    tokenizer = tokenizer_class.from_pretrained(
        args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
        do_lower_case=args.do_lower_case,
        cache_dir=args.cache_dir if args.cache_dir else None,
    )
    model = model_class.from_pretrained(
        args.model_name_or_path,
        from_tf=bool(".ckpt" in args.model_name_or_path),
        config=config,
        cache_dir=args.cache_dir if args.cache_dir else None,
    )

    if args.model_type == "bert":
        model.bert.encoder.set_early_exit_entropy(args.early_exit_entropy)
        model.bert.init_highway_pooler()
    elif args.model_type == "roberta":
        model.roberta.encoder.set_early_exit_entropy(args.early_exit_entropy)
        model.roberta.init_highway_pooler()
    else:
        raise NotImplementedError()

    if args.local_rank == 0:
        torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab

    model.to(args.device)

    logger.info("Training/evaluation parameters %s", args)

    # Training
    if args.do_train:
        train_dataset = load_and_cache_examples(args, args.task_name, tokenizer, evaluate=False)
        global_step, tr_loss = train(args, train_dataset, model, tokenizer)
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)

        if args.eval_after_first_stage:
            result = evaluate(args, model, tokenizer, prefix="")
            print_result = get_wanted_result(result)

        train(args, train_dataset, model, tokenizer, train_highway=True)

    # Saving best-practices: if you use defaults names for the model, you can reload it using from_pretrained()
    if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
        # Create output directory if needed
        if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
            os.makedirs(args.output_dir)

        logger.info("Saving model checkpoint to %s", args.output_dir)
        # Save a trained model, configuration and tokenizer using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        model_to_save = (
            model.module if hasattr(model, "module") else model
        )  # Take care of distributed/parallel training
        model_to_save.save_pretrained(args.output_dir)
        tokenizer.save_pretrained(args.output_dir)

        # Good practice: save your training arguments together with the trained model
        torch.save(args, os.path.join(args.output_dir, "training_args.bin"))

        # Load a trained model and vocabulary that you have fine-tuned
        model = model_class.from_pretrained(args.output_dir)
        tokenizer = tokenizer_class.from_pretrained(args.output_dir)
        model.to(args.device)

    # Evaluation
    results = {}
    if args.do_eval and args.local_rank in [-1, 0]:
        tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
        checkpoints = [args.output_dir]
        if args.eval_all_checkpoints:
            checkpoints = [
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
            ]

        logger.info("Evaluate the following checkpoints: %s", checkpoints)
        for checkpoint in checkpoints:
            global_step = checkpoint.split("-")[-1] if len(checkpoints) > 1 else ""
            prefix = checkpoint.split("/")[-1] if checkpoint.find("checkpoint") != -1 else ""

            model = model_class.from_pretrained(checkpoint)
            if args.model_type == "bert":
                model.bert.encoder.set_early_exit_entropy(args.early_exit_entropy)
            elif args.model_type == "roberta":
                model.roberta.encoder.set_early_exit_entropy(args.early_exit_entropy)
            else:
                raise NotImplementedError()

            model.to(args.device)
            result = evaluate(args, model, tokenizer, prefix=prefix, eval_highway=args.eval_highway)
            print_result = get_wanted_result(result)
            logger.info("Result: {}".format(print_result))
            if args.eval_each_highway:
                last_layer_results = print_result
                each_layer_results = []
                for i in range(model.num_layers):
                    logger.info("\n")
                    _result = evaluate(
                        args, model, tokenizer, prefix=prefix, output_layer=i, eval_highway=args.eval_highway
                    )
                    if i + 1 < model.num_layers:
                        each_layer_results.append(get_wanted_result(_result))
                each_layer_results.append(last_layer_results)
                save_fname = args.plot_data_dir + "/" + args.model_name_or_path[2:] + "/each_layer.npy"
                if not os.path.exists(os.path.dirname(save_fname)):
                    os.makedirs(os.path.dirname(save_fname))
                np.save(save_fname, np.array(each_layer_results))
                info_str = "Score of each layer:"
                for i in range(model.num_layers):
                    info_str += " {:.2f}".format(100 * each_layer_results[i])
                logger.info(info_str)
            result = {k + "_{}".format(global_step): v for k, v in result.items()}
            results.update(result)



SUMMARY:
Here's a succinct summary of the provided code:

**Overview**

This is a PyTorch implementation of a natural language processing (NLP) model, specifically BERT and RoBERTa architectures. The script trains, evaluates, and fine-tunes these models using various objectives and evaluation metrics.

**Key Components**

1. **Training**: Loads and caches training examples, sets up the training loop, and performs epochs of training.
2. **Evaluation**: Evaluates the models on a validation set or specific checkpoints and calculates the per-layer results for some architectures.
3. **Model Loading and Saving**: Saves trained model weights, configuration, and tokenizers to disk using `save_pretrained()` method.
4. **Hyperparameter Tuning**: Supports hyperparameter tuning using `args` object.

**Notable Features**

1. Custom evaluation functions (`evaluate()`, `get_wanted_result()`) for calculating per-layer results and aggregate scores.
2. Support for multiple model architectures (BERT, RoBERTa) with tailored initialization and early exit entropy settings.
3. Fine-tuning pre-trained models using `from_pretrained()` method.

**Unknowns**

1. Lack of clarity on specific hyperparameter configurations or optimization algorithms used in the training loop.
2. Assumption that evaluation metrics are provided in an external dictionary (no explicit code).
3. No documentation or comments explaining the custom functions (`evaluate()`, `get_wanted_result()`) in detail.

To improve this summary, more context about the actual NLP tasks being performed and specific details on hyperparameters or optimization strategies used would be helpful.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/deebert/test_glue_deebert.py


key: chroma:document
string_value:
CODE:

def get_setup_file():
    parser = argparse.ArgumentParser()
    parser.add_argument("-f")
    args = parser.parse_args()


SUMMARY:
**Summary:**

`get_setup_file()` is a utility function that takes an optional command-line argument `-f` and returns its value. The function utilizes `argparse` to parse the input, which does not support positional arguments. 

**Key Details:**

* Utilizes `argparse` for attribute parsing
* Expects single input with optional flag `-f`

**Notable Dependencies:**

* `argparse`
* ` sys` (not explicitly used in this snippet)

**Interface/Class Inheritance:**
This function does not inherit any specific interfaces; it leverages the internal functionality of the `argparse` module.


key: chroma:document
string_value:
CODE:

class DeeBertTests(TestCasePlus):
    def setup(self) -> None:
        stream_handler = logging.StreamHandler(sys.stdout)
        logger.addHandler(stream_handler)

    def run_and_check(self, args):
        n_gpu = get_gpu_count()

        if n_gpu > 1:
            pass
            # XXX: doesn't quite work with n_gpu > 1 https://github.com/huggingface/transformers/issues/10560
            # script = f"{self.examples_dir_str}/research_projects/deebert/run_glue_deebert.py"
            # distributed_args = f"-m torch.distributed.launch --nproc_per_node={n_gpu} {script}".split()
            # cmd = [sys.executable] + distributed_args + args
            # execute_subprocess_async(cmd, env=self.get_env())
            # XXX: test the results - need to save them first into .json file
        else:
            args.insert(0, "run_glue_deebert.py")
            with patch.object(sys, "argv", args):
                result = run_glue_deebert.main()
                for value in result.values():
                    self.assertGreaterEqual(value, 0.666)

    @slow
    @require_torch_non_multi_gpu
    def test_glue_deebert_train(self):
        train_args = """
            --model_type roberta
            --model_name_or_path FacebookAI/roberta-base
            --task_name MRPC
            --do_train
            --do_eval
            --do_lower_case
            --data_dir ./tests/fixtures/tests_samples/MRPC/
            --max_seq_length 128
            --per_gpu_eval_batch_size=1
            --per_gpu_train_batch_size=8
            --learning_rate 2e-4
            --num_train_epochs 3
            --overwrite_output_dir
            --seed 42
            --output_dir ./examples/deebert/saved_models/FacebookAI/roberta-base/MRPC/two_stage
            --plot_data_dir ./examples/deebert/results/
            --save_steps 0
            --overwrite_cache
            --eval_after_first_stage
            """.split()
        self.run_and_check(train_args)

        eval_args = """
            --model_type roberta
            --model_name_or_path ./examples/deebert/saved_models/FacebookAI/roberta-base/MRPC/two_stage
            --task_name MRPC
            --do_eval
            --do_lower_case
            --data_dir ./tests/fixtures/tests_samples/MRPC/
            --output_dir ./examples/deebert/saved_models/FacebookAI/roberta-base/MRPC/two_stage
            --plot_data_dir ./examples/deebert/results/
            --max_seq_length 128
            --eval_each_highway
            --eval_highway
            --overwrite_cache
            --per_gpu_eval_batch_size=1
            """.split()
        self.run_and_check(eval_args)

        entropy_eval_args = """
            --model_type roberta
            --model_name_or_path ./examples/deebert/saved_models/FacebookAI/roberta-base/MRPC/two_stage
            --task_name MRPC
            --do_eval
            --do_lower_case
            --data_dir ./tests/fixtures/tests_samples/MRPC/
            --output_dir ./examples/deebert/saved_models/FacebookAI/roberta-base/MRPC/two_stage
            --plot_data_dir ./examples/deebert/results/
            --max_seq_length 128
            --early_exit_entropy 0.1
            --eval_highway
            --overwrite_cache
            --per_gpu_eval_batch_size=1
            """.split()


SUMMARY:
**Class Summary:**

```
class DeeBertTests(TestCasePlus):
    ...
```

**Main Purpose and Functionality:**

This class is designed to test the DeeBert suite of models for natural language processing tasks. It inherits from `TestCasePlus` and provides a comprehensive set of tests for training, evaluation, and inference.

**Key Implementation Details:**

1. The `run_and_check` method runs the specified command with distributed arguments (if applicable) or standalone.
2. The `test_glue_deebert_train`, `test_glue_deebert_eval`, and `test_glue_deebert_entropy_eval` methods test different aspects of the DeeBert model, including training, evaluation, and inference.

**Notable Dependencies and Requirements:**

* `torch.distributed.launch`: Used for distributed testing.
* `get_gpu_count`: Returns the number of available GPUs.
* `sys.argv`, `sys.executable`: Used to pass command-line arguments and run the script.
* `logging` and `logging.StreamHandler`: Used for logging output to standard output.

**Method Signatures:**

1. `run_and_check(self, args)`;
2. `test_glue_deebert_train(self)`;
3. `test_glue_deebert_eval(self)`;
4. `test_glue_deebert_entropy_eval(self)`.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/deebert/src/modeling_highway_bert.py


key: chroma:document
string_value:
CODE:

def entropy(x):
    """Calculate entropy of a pre-softmax logit Tensor"""
    exp_x = torch.exp(x)
    A = torch.sum(exp_x, dim=1)  # sum of exp(x_i)
    B = torch.sum(x * exp_x, dim=1)  # sum of x_i * exp(x_i)


SUMMARY:
**Concise Summary**

This function calculates the entropy of a pre-softmax logit tensor.

### Primary Purpose and Functionality
Calculates the entropy of input tensor `x`, which represents soft-logits scores in a neural network. Entropy measures the uncertainty or unpredictability of these scores.

### Key Implementation Details

* The formula implements the definition of Shannon entropy (E[log(p)]), where p is the probability distribution represented by exp(x_i) / A.
* The sum of squared log-explicit probabilities `B` is computed, along with the sum of each individual logarithmic probability `A`.

### Notable Dependencies and Requirements

* `torch.exp(x)` calculates exponential function
* `torch.sum` computes sum across `dim=1`
* Interface: returns a tensor of shape (batch_size,)


key: chroma:document
string_value:
CODE:

class DeeBertEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
        self.highway = nn.ModuleList([BertHighway(config) for _ in range(config.num_hidden_layers)])

        self.early_exit_entropy = [-1 for _ in range(config.num_hidden_layers)]

    def set_early_exit_entropy(self, x):
        if isinstance(x, (float, int)):
            for i in range(len(self.early_exit_entropy)):
                self.early_exit_entropy[i] = x
        else:
            self.early_exit_entropy = x

    def init_highway_pooler(self, pooler):
        loaded_model = pooler.state_dict()
        for highway in self.highway:
            for name, param in highway.pooler.state_dict().items():
                param.copy_(loaded_model[name])

    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
    ):
        all_hidden_states = ()
        all_attentions = ()
        all_highway_exits = ()
        for i, layer_module in enumerate(self.layer):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_outputs = layer_module(
                hidden_states, attention_mask, head_mask[i], encoder_hidden_states, encoder_attention_mask
            )
            hidden_states = layer_outputs[0]

            if self.output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

            current_outputs = (hidden_states,)
            if self.output_hidden_states:
                current_outputs = current_outputs + (all_hidden_states,)
            if self.output_attentions:
                current_outputs = current_outputs + (all_attentions,)

            highway_exit = self.highway[i](current_outputs)
            # logits, pooled_output

            if not self.training:
                highway_logits = highway_exit[0]
                highway_entropy = entropy(highway_logits)
                highway_exit = highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
                all_highway_exits = all_highway_exits + (highway_exit,)

                if highway_entropy < self.early_exit_entropy[i]:
                    new_output = (highway_logits,) + current_outputs[1:] + (all_highway_exits,)
                    raise HighwayException(new_output, i + 1)
            else:
                all_highway_exits = all_highway_exits + (highway_exit,)

        # Add last layer
        if self.output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
        if self.output_attentions:
            outputs = outputs + (all_attentions,)

        outputs = outputs + (all_highway_exits,)


SUMMARY:
**DeeBertEncoder Class Summary**

**Main Purpose and Functionality**
The `DeeBertEncoder` class is an implementation of the BERT model's encoder, utilizing the PyTorch library. It encodes input sequences and generates contextualized representations.

**Key Implementation Details**

* Initializes multiple BERT layers and highway modules with the specified number of hidden layers
* Implements early exit strategy for premature termination during training
* Updates highway exit outputs with entropy values during training

**Notable Dependencies or Requirements**

* PyTorch library (`nn.Module`, `ModuleList`)
* BertLayer module (with `config` as input)
* Highway module (with `config` as input)
* Entropy calculation function (`entropy`)

**Method Signatures and Interfaces**

* `set_early_exit_entropy(x)`: updates early exit entropy values
* `init_highway_pooler(pooler)`: initializes highway pooler outputs
* `forward(hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask)`: encodes input sequence and generates contextualized representations.


key: chroma:document
string_value:
CODE:

class HighwayException(Exception):
    def __init__(self, message, exit_layer):
        self.message = message


SUMMARY:
**Summary**

HighwayException Class (RAG Code Summary)

*   **Purpose**: Custom exception class for handling errors related to highway layers (exit layer numbers)
*   **Key Implementation Details**:
    *   Initializes an exception with a message and exit layer number
    *   Exits at entry point 1 (`self.exit_layer = exit_layer # start from 1!`)
*   **Notable Dependencies or Requirements**: None explicitly mentioned, relies on standard Exception class for inheritance.


key: chroma:document
string_value:
CODE:

class BertHighway(nn.Module):
    """A module to provide a shortcut
    from (the output of one non-final BertLayer in BertEncoder) to (cross-entropy computation in BertForSequenceClassification)
    """

    def __init__(self, config):
        super().__init__()
        self.pooler = BertPooler(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

    def forward(self, encoder_outputs):
        # Pooler
        pooler_input = encoder_outputs[0]
        pooler_output = self.pooler(pooler_input)
        # "return" pooler_output

        # BertModel
        bmodel_output = (pooler_input, pooler_output) + encoder_outputs[1:]
        # "return" bmodel_output

        # Dropout and classification
        pooled_output = bmodel_output[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)



SUMMARY:
**Summary**

The `BertHighway` class is a neural network module designed to provide a shortcut from the output of one non-final BertLayer in BertEncoder to cross-entropy computation in BertForSequenceClassification. It combines features of:

1.  **Pooler Module**: Applies the BERT pooler to the input sequence, producing `pooler_input` and `pooler_output`.
2.  **BertModel Module**: Passes the original inputs (`encoder_outputs`) through the first layer, preserving additional layers (except the first one).
3.  **Dropout and Classification**: Applies dropout to the output of the second BertLayer and uses a linear classification layer to produce logits.

**Key Implementation Details**

*   Method Signature: `forward(encoder_outputs)`, returning `(logits, pooled_output)`
*   Key Algorithms: BERT pooler application, concatenation pattern with original layers
*   Design Choices: Combination of multiple pooling and classification operations

**Notable Dependencies or Requirements**

*   `nn.Module`: Parent class
*   `BertPooler`: Module inherited from parent class
*   `nn.Dropout`, `nn.Linear`: Modules imported from PyTorch library
*   `BertTokenizer` (assuming BertModel usage): No exact dependency listed in this code snippet.


key: chroma:document
string_value:
CODE:

def _split_thousands(delimiter, value):
    split = value.split(delimiter)


SUMMARY:
**Summary**

This function `_split_thousands` takes two parameters: `delimiter` and `value`. It splits the input `value` using the provided `delimiter`, checks if the resulting list has more than one element, and then verifies if any of these elements have a length of exactly 3. If both conditions are true, it returns `True`.

**Key Details**

* Splits `value` into substrings using `delimiter`
* Checks for more than one substring (`len(split) > 1`)
* Verifies if any substring has exactly 3 characters (`any((len(x) == 3 for x in split))`)
* Method signature: `_split_thousands(delimiter, value)`

**Notable Dependencies**

* None explicitly referenced
* Assumes a `delimiter` parameter passing to the function


key: chroma:document
string_value:
CODE:

def convert_to_float(value):
    """Converts value to a float using a series of increasingly complex heuristics.
    Args:
      value: object that needs to be converted. Allowed types include
        float/int/strings.
    Returns:
      A float interpretation of value.
    Raises:
      ValueError if the float conversion of value fails.
    """
    if isinstance(value, float):
        return value
    if isinstance(value, int):
        return float(value)
    if not isinstance(value, str):
        raise TypeError("Argument value is not a string. Can't parse it as float")
    sanitized = value

    try:
        # Example: 1,000.7
        if "." in sanitized and "," in sanitized:
            return float(sanitized.replace(",", ""))
        # 1,000
        if "," in sanitized and _split_thousands(",", sanitized):
            return float(sanitized.replace(",", ""))
        # 5,5556
        if "," in sanitized and sanitized.count(",") == 1 and not _split_thousands(",", sanitized):
            return float(sanitized.replace(",", "."))
        # 0.0.0.1
        if sanitized.count(".") > 1:
            return float(sanitized.replace(".", ""))
        # 0,0,0,1
        if sanitized.count(",") > 1:
            return float(sanitized.replace(",", ""))
        return float(sanitized)
    except ValueError:
        # Avoid adding the sanitized value in the error message.


SUMMARY:
**Function Summary: `convert_to_float`**

Converting object to float using various heuristics.

```python
def convert_to_float(value):
    """
    Args:
      value (object): Value to be converted.
    Returns:
      float
    Raises:
      ValueError: If conversion fails, providing limited error context.
    """
    
    # Base cases
    if isinstance(value, float) or isinstance(value, int):
        return float(value)
    
    # Handling non-string values with a TypeError
    if not isinstance(value, str):
        raise TypeError("Argument value is not a string. Can't parse it as float")
    
    sanitized = value
    
    # Complex heuristics:
    try:
        # Replaces comma separators for US decimal formatting
        if "," in sanitized and _split_thousands(",", sanitized):
            return float(sanitized.replace(",", ""))
        
        # Handles comma delimited large numbers
        if "," in sanitized and sanitized.count(",") == 1 and not _split_thousands(",", sanitized):
            return float(sanitized.replace(",", "."))
        
        # Replaces dot separators for thousandal separation
        if.","in sanitized:
            return float(sanitized)
        
        # Converts from any other format
        return float(sanitized)
    except ValueError as e:
        # Avoid exposing original value in error context
        raise ValueError("Unable to convert value to float")
```

**Key Implementation Details:**

1. **Complex Heuristics**: Multiple rules are used to handle comma-separated decimal and thousandal formats.
2. **`_split_thousands` function (not shown)**: Used to separate thousands from a number with comma separators.

**Notable Dependencies and Requirements:**

* `_split_thousands` function

Note: The functions `isinstance`, `try-except` block, and the variable names (`_split_thousands`). are preserved exactly.


key: chroma:document
string_value:
CODE:

def _normalize_float(answer):
    if answer is None:
        return None
    try:
        value = convert_to_float(answer)
        if isinstance(value, float) and math.isnan(value):
            return None
        return value
    except ValueError:


SUMMARY:
**`_normalize_float` Function Summary**

**Main Purpose and Functionality**
The `_normalize_float` function normalizes a given `answer` by converting it to a float, handling potential errors, and applying additional cleanup operations as needed.

**Implementation Details**

* Converts the input `answer` to a float using `convert_to_float`.
* Checks if the resulting value is NaN (Not a Number) and returns None in such cases.
* If a `ValueError` occurs during conversion, converts the answer to lowercase instead.

**Note:** External dependency: The function relies on an external `convert_to_float` function, which is not defined in this snippet. Additionally, it inherits from `math` namespace, using its `isnan` and `isNaN` functions to detect NaN values.


key: chroma:document
string_value:
CODE:

class _Aggregation(enum.Enum):
    """Aggregations as defined by WikiSQL. Indexes match the data."""

    NONE = 0
    MAX = 1
    MIN = 2
    COUNT = 3
    SUM = 4


SUMMARY:
**Aggregate Functions Enum**

This class defines an enumeration of aggregate functions used in WikiSQL indexes:

*   `Aggregation` Enum:
    *   Defines 6 distinct aggregations: `NONE`, `MAX`, `MIN`, `COUNT`, `SUM`, and `AVERAGE`.

    **Purpose:** Provides a standardized representation of common aggregation operations.

**Key Implementation Detail:**
This enum is used to standardize the classification of WikiSQL indexes.


key: chroma:document
string_value:
CODE:

class _Operator(enum.Enum):
    """The boolean operators used by WikiSQL. Indexes match the data."""

    EQUALS = 0
    GREATER = 1


SUMMARY:
**Summary:**

This code defines an enumeration class `_Operator` with three values: `EQUALS`, `GREATER`, and `LESSER`. These represent boolean operators used by WikiSQL, mapping index names to their corresponding logical operations. The primary purpose is to provide a consistent and standardized way of representing Boolean logic in the application.

**Key Details:**

* Enum type: `_Operator`
* Enum values: `EQUALS=0`, `GREATER=1`, `LESSER=2`

**No notable dependencies or requirements**


key: chroma:document
string_value:
CODE:

def _normalize_for_match(x):


SUMMARY:
**Summary:**

The `_normalize_for_match` function takes an input string `x`, converts it to lowercase, and then uses the `_TOKENIZER.findall` method to match its characters against predefined tokens. It returns these matches as a list of strings.

**Key Implementation Details:**

* The input string is converted to lowercase using the `lower()` method.
* The function relies on the `_TOKENIZER.foundall` method, which is not defined in this snippet (presumably part of an external library).

**Notable Dependencies or Requirements:**

* `_TOKENIZER` appears to be an object with a `findall` attribute, but its definition is not shown here.


key: chroma:document
string_value:
CODE:

def _compare(operator, src, tgt):
    if operator == _Operator.EQUALS:
        return src == tgt
    elif operator == _Operator.GREATER:
        return src > tgt
    elif operator == _Operator.LESSER:
        return src < tgt


SUMMARY:
**Summary**

`_compare` function compares two values based on an `operator`. It returns a boolean result for equality (`_OperaTor.EQUALS`) greater than, less than, or value comparison for other operators.

* Primary Purpose: Compare values based on operator
* Key Implementation Details:
	+ Uses pre-defined `_Operator` enum to map operator values to comparison functions
	+ Handles three main operations: equality (`src == tgt`), greater by one (`src > tgt`), and lesser by one (`src < tgt`)
	+ Rases `ValueError` if an unknown operator is provided

**No notable dependencies, requirements or inherited/introduced interfaces.**

* `_Operator` enum not defined in this snippet but assumed to be defined elsewhere.


key: chroma:document
string_value:
CODE:

def _parse_value(table, column, cell_value):
    """Convert numeric values to floats and keeps everything else as string."""
    types = table["types"]


SUMMARY:
**Summary**

*Purpose: Converts a cell value to its corresponding type, either as a float or string.*

**Key Implementation Details:**

* Uses a `_TYPE_CONVERTER` dictionary to determine the conversion function based on the table's `types` column for the specified `column`.
* Applies the determined conversion function to the `cell_value`.

**Notable Dependencies/Requirements:**

* Assumes existence of `_TYPE_CONVERTER` dictionary with type-conversion functions.


key: chroma:document
string_value:
CODE:

def _is_string(x):


SUMMARY:
**Summary:**

* **Function Name:** `_is_string`
* **Primary Purpose:** Checks if the input `x` is an instance of the `str` type (`isinstance(x, str)`).
* **Notable Implementation Detail:** Utilizes Python's built-in `isinstance` function to make a simple and efficient string verification.
 
**Note:** This implementation uses a private underscore prefix `_`, suggesting it may be intended for internal use within a class or module.


key: chroma:document
string_value:
CODE:

def _respect_conditions(table, row, conditions):
    """True if 'row' satisfies all 'conditions'."""
    for cond in conditions:
        table_value = row[cond.column]

        cmp_value = _parse_value(table, cond.column, cond.cmp_value)

        if _is_string(table_value) and _is_string(cmp_value):
            table_value = _normalize_for_match(table_value)
            cmp_value = _normalize_for_match(cmp_value)

        if not isinstance(table_value, type(cmp_value)):
            raise TypeError("Type difference {} != {}".format(type(table_value), type(cmp_value)))

        if not _compare(cond.operator, table_value, cmp_value):
            return False


SUMMARY:
**Summary:**

`_respect_conditions` function evaluates if a row (`row`) meets all conditions (`conditions`) in a given table (`table`). It compares values based on specified operators and types. The primary purpose is to check for truthfulness against certain criteria.

**Key Implementation Details:**

* Iterates through conditions, normalizing strings and comparing values using `_compare` function with type checks.
* Raises `TypeError` if data types do not match.

**Notable Dependencies or Requirements:**

* `_parse_value`
* `_is_string`
* `_normalize_for_match`
* `_compare`
* `type`


key: chroma:document
string_value:
CODE:

def _get_float_answer(table, answer_coordinates, aggregation_op):
    """Applies operation to produce reference float answer."""
    if not answer_coordinates:
        if aggregation_op == _Aggregation.COUNT:
            return 0.0
        else:
            return EMPTY_ANSWER_AGG

    # Count can support non numeric answers.
    if aggregation_op == _Aggregation.COUNT:
        return float(len(answer_coordinates))

    # If we have just one answer, if float returns it or try a conversion.
    values = [table["rows"][i][j] for (i, j) in answer_coordinates]
    if len(answer_coordinates) == 1:
        try:
            return convert_to_float(values[0])
        except ValueError as e:
            if aggregation_op != _Aggregation.NONE:
                raise e

    if aggregation_op == _Aggregation.NONE:
        return None

    # Other aggregation only support numeric values. Bail out if we have strings.
    if not all((isinstance(v, (int, float)) for v in values)):
        return None

    if aggregation_op == _Aggregation.SUM:
        return float(sum(values))
    elif aggregation_op == _Aggregation.AVERAGE:
        return sum(values) / len(answer_coordinates)
    else:


SUMMARY:
**Summary**

`_get_float_answer` function:

Computes a reference float answer for given table data based on `answer_coordinates`, aggregating values using `_Aggregation.COUNT`, `_Aggregation.SUM`, or `_Aggregation.AVERAGE`.

**Key Functionality**

* Applies operation to produce reference float answer (e.g., counting, summing, averaging).
* Supports non-numeric answers with `_Aggregation.COUNT`.
* Handles single value with `try`-`except` block for conversion.
* Ensures only numeric values are processed for aggregation operations.

**Method Signatures and Interfaces**

* Takes three parameters: `table`, `answer_coordinates`, and `aggregation_op`.

**Key Implementation Details**

* Uses list comprehension to extract table rows at specified coordinates (`values = ...`).
* Utilizes Python's built-in `sum` function for aggregating values.
* Throws exceptions for unknown aggregation operations or non-numeric value types.

**Notable Dependencies and Requirements**

* `_Aggregation` (enumeration) imported from another module.


key: chroma:document
string_value:
CODE:

def _get_answer_coordinates(table, sql_query):
    """Retrieves references coordinates by executing SQL."""
    # MAX and MIN are automatically supported by the model.
    aggregation_op_index = sql_query["agg"]
    if aggregation_op_index >= 3:
        aggregation_op = _Aggregation(aggregation_op_index)
    else:
        aggregation_op = _Aggregation.NONE

    target_column = sql_query["sel"]
    conditions = [
        _Condition(column, _Operator(operator), cmp_value)
        for column, operator, cmp_value in zip(
            sql_query["conds"]["column_index"], sql_query["conds"]["operator_index"], sql_query["conds"]["condition"]
        )
    ]

    indices = []
    for row in range(len(table["rows"])):
        if _respect_conditions(table, table["rows"][row], conditions):
            indices.append((row, target_column))

    if not indices:
        return [], aggregation_op

    if len(indices) == 1:
        return indices, aggregation_op

    # Parsing of MIN/MAX.
    if aggregation_op_index in (1, 2):
        operators = {2: min, 1: max}
        values = [(table["rows"][i][j], index) for index, (i, j) in enumerate(indices)]
        reduced = functools.reduce(operators[sql_query["agg"]], values)

        ret = [indices[reduced[1]]]
        return ret, _Aggregation.NONE



SUMMARY:
**Summary: `_get_answer_coordinates` Function**

This function retrieves reference coordinates from a table by executing an SQL query.

**Purpose:** Given an SQL query with aggregation and condition specifications, it executes the query on a table to find relevant rows based on conditions and returns their coordinates.

**Key Implementation Details:**

* Uses `agg` and `sel` parameters to specify aggregation operation (max/min) and target column.
* Applies conditions using `_Condition` class to filter rows.
* Reduces values for min/max operations using `functools.reduce`.
* Returns a list of indices with corresponding coordinates if multiple matching rows are found.

**Notable Dependencies/Requirements:**

* `_Aggregation` enum ( Aggregation<NONE>, Aggregation<max>, Aggregation<min> )
* `_Operator` class
* `_Condition` class
* `functools.reduce`
* External dependencies:
	+ `functools` for reduce function
	+ `_Aggregation` class
	+ `_Operator` and `_Condition` classes


key: chroma:document
string_value:
CODE:

def _get_answer_text(table, answer_coordinates, float_answer):
    if float_answer is not None:
        return [str(float_answer)]


SUMMARY:
**Summary**

Function: `_get_answer_text`

Purpose: Retrieves the text representation of answers from a table based on given coordinates.

*   Primary purpose: Returns a list of strings representing the answers at specified coordinates or a single floating-point value if provided.
*   Key implementation details:
    *   Uses list comprehension to extract values from "real_rows" in the table using answer coordinates (r, c).
    *   If `float_answer` is not None, returns a list containing its string representation.

Notable dependencies:
*   Table data accessed through object attributes ("table", "real_rows")
*   External dependency: None

Interfaces/Methods involved:
*   List comprehension
*   String conversion (str())


key: chroma:document
string_value:
CODE:

def retrieve_wikisql_query_answer_tapas(table, example) -> List:
    answer_coordinates, aggregation_op = _get_answer_coordinates(table, example)
    float_answer = _get_float_answer(table, answer_coordinates, aggregation_op)
    answer_text = _get_answer_text(table, answer_coordinates, float_answer)
    # keep the original data the same with TaPas
    if len(answer_text) == 0:
        answer_text = [EMPTY_ANSWER]


SUMMARY:
Here is a concise and focused summary of the function code chunk:

**Function Summary: `retrieve_wikisql_query_answer_tapas`**

Purpose: Retrieves a Wikipedia SQL query answer with Tapas-specific formatting.

Returns a list of strings representing the formatted answer coordinates and text.

Key Implementation Details:
- Calls `_get_answer_coordinates`, `_get_float_answer`, and `_get_answer_text` functions to generate intermediate results.
- Appends an empty answer string (`EMPTY_ANSWER`) to `answer_text` if its length is 0.

Notable Dependencies:
- `_get_answer_coordinates`
- `_get_float_answer`
- `_get_answer_text`

External Dependencies: List of related classes/functions (_get_answer_coordinates, _get_float_answer, and _get_answer_text are assumed to be defined elsewhere in the codebase).

 Interface Inheritance/Implementation: Returns a list (`List` is an assumed external class/interface).


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/tapex/run_wikisql_with_tapex.py


key: chroma:document
string_value:
CODE:

def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    if is_main_process(training_args.local_rank):
        transformers.utils.logging.set_verbosity_info()
    logger.info(f"Training/evaluation parameters {training_args}")

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)
    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
    # (the dataset will be downloaded automatically from the datasets Hub).
    #
    # For JSON files, this script will use the `question` column for the input question and `table` column for the corresponding table.
    #
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
    if data_args.dataset_name is not None:
        # Downloading and loading a dataset from the hub.
        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir)
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
            extension = data_args.train_file.split(".")[-1]
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
            extension = data_args.validation_file.split(".")[-1]
        if data_args.test_file is not None:
            data_files["test"] = data_args.test_file
            extension = data_args.test_file.split(".")[-1]
        datasets = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)

    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
    # https://huggingface.co/docs/datasets/loading_datasets.

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
    )

    # IMPORTANT: the initial BART model's decoding is penalized by no_repeat_ngram_size, and thus
    # we should disable it here to avoid problematic generation
    config.no_repeat_ngram_size = 0
    config.max_length = 1024
    config.early_stopping = False

    # load tapex tokenizer
    tokenizer = TapexTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
        add_prefix_space=True,
    )

    # load Bart based Tapex model (default tapex-large)
    model = BartForConditionalGeneration.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
    )

    if model.config.decoder_start_token_id is None:
        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")

    # Preprocessing the datasets.
    # We need to tokenize inputs and targets.
    if training_args.do_train:
        column_names = datasets["train"].column_names
    elif training_args.do_eval:
        column_names = datasets["validation"].column_names
    elif training_args.do_predict:
        column_names = datasets["test"].column_names
    else:
        logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.")
        return

    # Temporarily set max_target_length for training.
    max_target_length = data_args.max_target_length
    padding = "max_length" if data_args.pad_to_max_length else False

    if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"):
        logger.warning(
            "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for "
            f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
        )

    def preprocess_tableqa_function(examples, is_training=False):
        """
        The is_training FLAG is used to identify if we could use the supervision
        to truncate the table content if it is required.
        """

        # this function is specific for WikiSQL since the util function need the data structure
        # to retrieve the WikiSQL answer for each question
        def _convert_table_types(_table):
            """Runs the type converter over the table cells."""
            ret_table = deepcopy(_table)
            types = ret_table["types"]
            ret_table["real_rows"] = ret_table["rows"]
            typed_rows = []
            for row in ret_table["rows"]:
                typed_row = []
                for column, cell_value in enumerate(row):
                    typed_row.append(_TYPE_CONVERTER[types[column]](cell_value))
                typed_rows.append(typed_row)
            ret_table["rows"] = typed_rows
            return ret_table

        questions = [question.lower() for question in examples["question"]]
        example_tables = examples["table"]
        example_sqls = examples["sql"]
        tables = [
            pd.DataFrame.from_records(example_table["rows"], columns=example_table["header"])
            for example_table in example_tables
        ]

        # using tapas utils to obtain wikisql answer
        answers = []
        for example_sql, example_table in zip(example_sqls, example_tables):
            tapas_table = _convert_table_types(example_table)
            answer_list: List[str] = retrieve_wikisql_query_answer_tapas(tapas_table, example_sql)
            # you can choose other delimiters to split each answer
            answers.append(answer_list)

        # IMPORTANT: we cannot pass by answers during evaluation, answers passed during training are used to
        # truncate large tables in the train set!
        if is_training:
            model_inputs = tokenizer(
                table=tables,
                query=questions,
                answer=answers,
                max_length=data_args.max_source_length,
                padding=padding,
                truncation=True,
            )
        else:
            model_inputs = tokenizer(
                table=tables, query=questions, max_length=data_args.max_source_length, padding=padding, truncation=True
            )

        labels = tokenizer(
            answer=[", ".join(answer) for answer in answers],
            max_length=max_target_length,
            padding=padding,
            truncation=True,
        )

        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
        # padding in the loss.
        if padding == "max_length" and data_args.ignore_pad_token_for_loss:
            labels["input_ids"] = [
                [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
            ]

        model_inputs["labels"] = labels["input_ids"]

        return model_inputs

    # in training, we can use the answer as extra information to truncate large tables
    preprocess_tableqa_function_training = partial(preprocess_tableqa_function, is_training=True)

    if training_args.do_train:
        if "train" not in datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = datasets["train"]
        if data_args.max_train_samples is not None:
            train_dataset = train_dataset.select(range(data_args.max_train_samples))
        train_dataset = train_dataset.map(
            preprocess_tableqa_function_training,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
        )

    if training_args.do_eval:
        max_target_length = data_args.val_max_target_length
        if "validation" not in datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = datasets["validation"]
        if data_args.max_eval_samples is not None:
            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
        eval_dataset = eval_dataset.map(
            preprocess_tableqa_function,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
        )

    if training_args.do_predict:
        max_target_length = data_args.val_max_target_length
        if "test" not in datasets:
            raise ValueError("--do_predict requires a test dataset")
        predict_dataset = datasets["test"]
        if data_args.max_predict_samples is not None:
            predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
        predict_dataset = predict_dataset.map(
            preprocess_tableqa_function,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
        )

    # Data collator
    label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=label_pad_token_id,
        pad_to_multiple_of=8 if training_args.fp16 else None,
    )

    def postprocess_text(preds, labels):
        preds = [pred.strip() for pred in preds]
        labels = [label.strip() for label in labels]

        return preds, labels

    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        if data_args.ignore_pad_token_for_loss:
            # Replace -100 in the labels as we can't decode them.
            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

        # Some simple post-processing
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

        delimiter = ", "

        # define example evaluation
        def evaluate_example(predict_str: str, ground_str: str):
            predict_spans = predict_str.split(delimiter)
            ground_spans = ground_str.split(delimiter)
            predict_values = defaultdict(lambda: 0)
            ground_values = defaultdict(lambda: 0)
            for span in predict_spans:
                try:
                    predict_values[float(span)] += 1
                except ValueError:
                    predict_values[span.strip()] += 1
            for span in ground_spans:
                try:
                    ground_values[float(span)] += 1
                except ValueError:
                    ground_values[span.strip()] += 1
            is_correct = predict_values == ground_values
            return is_correct

        def get_denotation_accuracy(predictions: List[str], references: List[str]):
            assert len(predictions) == len(references)
            correct_num = 0
            for predict_str, ground_str in zip(predictions, references):
                is_correct = evaluate_example(predict_str.lower(), ground_str.lower())
                if is_correct:
                    correct_num += 1
            return correct_num / len(predictions)

        accuracy = get_denotation_accuracy(decoded_preds, decoded_labels)
        result = {"denotation_accuracy": accuracy}

        return result

    # Initialize our Trainer
    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if training_args.predict_with_generate else None,
    )

    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload

        metrics = train_result.metrics
        max_train_samples = (
            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        )
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # Evaluation
    results = {}
    if training_args.do_eval:
        logger.info("*** Evaluate ***")

        metrics = trainer.evaluate(
            max_length=data_args.val_max_target_length, num_beams=data_args.num_beams, metric_key_prefix="eval"
        )
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))

        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    if training_args.do_predict:
        logger.info("*** Predict ***")

        predict_results = trainer.predict(
            predict_dataset,
            metric_key_prefix="predict",
            max_length=data_args.val_max_target_length,
            num_beams=data_args.num_beams,
        )
        metrics = predict_results.metrics
        max_predict_samples = (
            data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
        )
        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))

        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)

        if trainer.is_world_process_zero():
            if training_args.predict_with_generate:
                predictions = tokenizer.batch_decode(
                    predict_results.predictions, skip_special_tokens=True, clean_up_tokenization_spaces=True
                )
                predictions = [pred.strip() for pred in predictions]
                output_prediction_file = os.path.join(training_args.output_dir, "tapex_predictions.txt")
                with open(output_prediction_file, "w") as writer:
                    writer.write("\n".join(predictions))



SUMMARY:
This code is for training and testing the Seq2Seq model using Hugging Face's Transformers library. The code trains a sequence-to-sequence model on a dataset, evaluates its performance on a test set, and generates predictions.

The main components of this code are:

1. Initialization: It initializes the model, tokenizer, arguments for training, evaluation, and prediction.
2. Data loading: It loads the training and testing datasets.
3. Training: The `Trainer` class is used to train the model on the training dataset.
4. Evaluation: After training, it evaluates the model's performance on the test set using a specific metric.
5. Prediction: If requested, it generates predictions from the trained model.

The output of this code includes:

* The metrics (denotation accuracy) for the entire training process
* A report of the precision, recall, and other evaluation metrics
* A file containing the generated predictions.

This code uses Hugging Face's Transformers library to implement a sequence-to-sequence model using the pre-trained Tapex model.


key: chroma:document
string_value:
CODE:

def _mp_fn(index):
    # For xla_spawn (TPUs)


SUMMARY:
**Summary:**

* **Purpose:** This is an internal function `_mp_fn` that appears to be part of a job management system, utilized when running on TPUs (Tensor Processing Units) via `xla_spawn`.
* **Key Implementation Detail:** The function calls the `main()` function after execution.
* **Notable Dependencies/Requirements:**
	+ None explicitly mentioned; however, it relies on the `TPU` execution environment and `xla_spawn`.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/tapex/run_wikitablequestions_with_tapex.py


key: chroma:document
string_value:
CODE:

def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    # Set the verbosity to info of the Transformers logger (on main process only):
    if is_main_process(training_args.local_rank):
        transformers.utils.logging.set_verbosity_info()
    logger.info(f"Training/evaluation parameters {training_args}")

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)
    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
    # (the dataset will be downloaded automatically from the datasets Hub).
    #
    # For JSON files, this script will use the `question` column for the input question and `table` column for the corresponding table.
    #
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
    if data_args.dataset_name is not None:
        # Downloading and loading a dataset from the hub.
        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir)
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
            extension = data_args.train_file.split(".")[-1]
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
            extension = data_args.validation_file.split(".")[-1]
        if data_args.test_file is not None:
            data_files["test"] = data_args.test_file
            extension = data_args.test_file.split(".")[-1]
        datasets = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)

    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
    # https://huggingface.co/docs/datasets/loading_datasets.

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
    )

    # IMPORTANT: the initial BART model's decoding is penalized by no_repeat_ngram_size, and thus
    # we should disable it here to avoid problematic generation
    config.no_repeat_ngram_size = 0
    config.max_length = 1024
    config.early_stopping = False

    # load tapex tokenizer
    tokenizer = TapexTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
        add_prefix_space=True,
    )

    # load Bart based Tapex model (default tapex-large)
    model = BartForConditionalGeneration.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
    )

    if model.config.decoder_start_token_id is None:
        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")

    # Preprocessing the datasets.
    # We need to tokenize inputs and targets.
    if training_args.do_train:
        column_names = datasets["train"].column_names
    elif training_args.do_eval:
        column_names = datasets["validation"].column_names
    elif training_args.do_predict:
        column_names = datasets["test"].column_names
    else:
        logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.")
        return

    # Temporarily set max_target_length for training.
    max_target_length = data_args.max_target_length
    padding = "max_length" if data_args.pad_to_max_length else False

    if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"):
        logger.warning(
            "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for "
            f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
        )

    def preprocess_tableqa_function(examples, is_training=False):
        """
        The is_training FLAG is used to identify if we could use the supervision
        to truncate the table content if it is required.
        """

        questions = [question.lower() for question in examples["question"]]
        example_tables = examples["table"]
        tables = [
            pd.DataFrame.from_records(example_table["rows"], columns=example_table["header"])
            for example_table in example_tables
        ]

        # using wikitablequestion's answer set
        answers = examples["answers"]

        # IMPORTANT: we cannot pass by answers during evaluation, answers passed during training are used to
        # truncate large tables in the train set!
        if is_training:
            model_inputs = tokenizer(
                table=tables,
                query=questions,
                answer=answers,
                max_length=data_args.max_source_length,
                padding=padding,
                truncation=True,
            )
        else:
            model_inputs = tokenizer(
                table=tables, query=questions, max_length=data_args.max_source_length, padding=padding, truncation=True
            )

        labels = tokenizer(
            answer=[", ".join(answer) for answer in answers],
            max_length=max_target_length,
            padding=padding,
            truncation=True,
        )

        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
        # padding in the loss.
        if padding == "max_length" and data_args.ignore_pad_token_for_loss:
            labels["input_ids"] = [
                [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
            ]

        model_inputs["labels"] = labels["input_ids"]

        return model_inputs

    # in training, we can use the answer as extra information to truncate large tables
    preprocess_tableqa_function_training = partial(preprocess_tableqa_function, is_training=True)

    if training_args.do_train:
        if "train" not in datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = datasets["train"]
        if data_args.max_train_samples is not None:
            train_dataset = train_dataset.select(range(data_args.max_train_samples))
        train_dataset = train_dataset.map(
            preprocess_tableqa_function_training,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
        )

    if training_args.do_eval:
        max_target_length = data_args.val_max_target_length
        if "validation" not in datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = datasets["validation"]
        if data_args.max_eval_samples is not None:
            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
        eval_dataset = eval_dataset.map(
            preprocess_tableqa_function,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
        )

    if training_args.do_predict:
        max_target_length = data_args.val_max_target_length
        if "test" not in datasets:
            raise ValueError("--do_predict requires a test dataset")
        predict_dataset = datasets["test"]
        if data_args.max_predict_samples is not None:
            predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
        predict_dataset = predict_dataset.map(
            preprocess_tableqa_function,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            remove_columns=column_names,
            load_from_cache_file=not data_args.overwrite_cache,
        )

    # Data collator
    label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=label_pad_token_id,
        pad_to_multiple_of=8 if training_args.fp16 else None,
    )

    def postprocess_text(preds, labels):
        preds = [pred.strip() for pred in preds]
        labels = [label.strip() for label in labels]

        return preds, labels

    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        if data_args.ignore_pad_token_for_loss:
            # Replace -100 in the labels as we can't decode them.
            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

        # Some simple post-processing
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

        delimiter = ", "

        # define example evaluation
        def evaluate_example(predict_str: str, ground_str: str):
            predict_spans = predict_str.split(delimiter)
            ground_spans = ground_str.split(delimiter)
            predict_values = defaultdict(lambda: 0)
            ground_values = defaultdict(lambda: 0)
            for span in predict_spans:
                try:
                    predict_values[float(span)] += 1
                except ValueError:
                    predict_values[span.strip()] += 1
            for span in ground_spans:
                try:
                    ground_values[float(span)] += 1
                except ValueError:
                    ground_values[span.strip()] += 1
            _is_correct = predict_values == ground_values
            return _is_correct

        def get_denotation_accuracy(predictions: List[str], references: List[str]):
            assert len(predictions) == len(references)
            correct_num = 0
            for predict_str, ground_str in zip(predictions, references):
                is_correct = evaluate_example(predict_str.lower(), ground_str.lower())
                if is_correct:
                    correct_num += 1
            return correct_num / len(predictions)

        accuracy = get_denotation_accuracy(decoded_preds, decoded_labels)
        result = {"denotation_accuracy": accuracy}

        return result

    # Initialize our Trainer
    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if training_args.predict_with_generate else None,
    )

    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        trainer.save_model()  # Saves the tokenizer too for easy upload

        metrics = train_result.metrics
        max_train_samples = (
            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        )
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # Evaluation
    results = {}
    if training_args.do_eval:
        logger.info("*** Evaluate ***")

        metrics = trainer.evaluate(
            max_length=data_args.val_max_target_length, num_beams=data_args.num_beams, metric_key_prefix="eval"
        )
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))

        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    if training_args.do_predict:
        logger.info("*** Predict ***")

        predict_results = trainer.predict(
            predict_dataset,
            metric_key_prefix="predict",
            max_length=data_args.val_max_target_length,
            num_beams=data_args.num_beams,
        )
        metrics = predict_results.metrics
        max_predict_samples = (
            data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
        )
        metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))

        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)

        if trainer.is_world_process_zero():
            if training_args.predict_with_generate:
                predictions = tokenizer.batch_decode(
                    predict_results.predictions, skip_special_tokens=True, clean_up_tokenization_spaces=True
                )
                predictions = [pred.strip() for pred in predictions]
                output_prediction_file = os.path.join(training_args.output_dir, "tapex_predictions.txt")
                with open(output_prediction_file, "w") as writer:
                    writer.write("\n".join(predictions))



SUMMARY:
This is a training script for the Hugging Face Transformers library to train and evaluate a sequence-to-sequence model for data denotation tasks. The script trains a model on a given dataset, evaluates its performance on an evaluation set, and generates predictions on a prediction set.

Key components:

* Model training: The script uses the `Seq2SeqTrainer` class from Hugging Face to train a model on the training dataset.
* Evaluation: The script evaluates the trained model's performance on the evaluation set using metrics such as denotation accuracy.
* Prediction: The script generates predictions on the prediction set and logs them along with other metrics.


key: chroma:document
string_value:
CODE:

def _mp_fn(index):
    # For xla_spawn (TPUs)


SUMMARY:
**Function Summary: `_mp_fn`**

The `_mp_fn` function takes an `index` parameter and calls the `main()` function as part of the XLA spawn execution. 

**Implementation Details:**

* Calls `main()` without any parameters.

**Notable Dependencies/Requirements:**

* The `main()` function is executed in the context of XLA spawn, indicating compatibility with TPUs (Tensor Processing Units).


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/tapex/run_tabfact_with_tapex.py


key: chroma:document
string_value:
CODE:

def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    # Set seed before initializing model.
    set_seed(training_args.seed)

    # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)
    # or specify a GLUE benchmark task (the dataset will be downloaded automatically from the datasets Hub).
    #
    # For JSON files, this script will use the `question` column for the input question and `table` column for the corresponding table.
    #
    # If the CSVs/JSONs contain only one non-label column, the script does single sentence classification on this
    # single column. You can easily tweak this behavior (see below)
    #
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
    if data_args.dataset_name is not None:
        # Downloading and loading a dataset from the hub.
        raw_datasets = load_dataset(
            data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir
        )
    else:
        # Loading a dataset from your local files.
        # CSV/JSON training and evaluation files are needed.
        data_files = {"train": data_args.train_file, "validation": data_args.validation_file}

        # Get the test dataset: you can provide your own CSV/JSON test file (see below)
        # when you use `do_predict` without specifying a GLUE benchmark task.
        if training_args.do_predict:
            if data_args.test_file is not None:
                train_extension = data_args.train_file.split(".")[-1]
                test_extension = data_args.test_file.split(".")[-1]
                assert (
                    test_extension == train_extension
                ), "`test_file` should have the same extension (csv or json) as `train_file`."
                data_files["test"] = data_args.test_file
            else:
                raise ValueError("Need either a GLUE task or a test file for `do_predict`.")

        for key in data_files.keys():
            logger.info(f"load a local file for {key}: {data_files[key]}")

        if data_args.train_file.endswith(".csv"):
            # Loading a dataset from local csv files
            raw_datasets = load_dataset("csv", data_files=data_files, cache_dir=model_args.cache_dir)
        else:
            # Loading a dataset from local json files
            raw_datasets = load_dataset("json", data_files=data_files, cache_dir=model_args.cache_dir)
    # See more about loading any type of standard or custom dataset at
    # https://huggingface.co/docs/datasets/loading_datasets.

    # Labels
    label_list = raw_datasets["train"].features["label"].names
    num_labels = len(label_list)

    # Load pretrained model and tokenizer
    #
    # In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=num_labels,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
    )
    # load tapex tokenizer
    tokenizer = TapexTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
        add_prefix_space=True,
    )
    model = BartForSequenceClassification.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        token=True if model_args.use_auth_token else None,
    )

    # Padding strategy
    if data_args.pad_to_max_length:
        padding = "max_length"
    else:
        # We will pad later, dynamically at batch creation, to the max sequence length in each batch
        padding = False

    # Some models have set the order of the labels to use, so let's make sure we do use it.
    model.config.label2id = {"Refused": 0, "Entailed": 1}
    model.config.id2label = {0: "Refused", 1: "Entailed"}

    if data_args.max_seq_length > tokenizer.model_max_length:
        logger.warning(
            f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the "
            f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
        )
    max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)

    def preprocess_tabfact_function(examples):
        # Tokenize the texts
        def _convert_table_text_to_pandas(_table_text):
            """Runs the structured pandas table object for _table_text.
            An example _table_text can be: round#clubs remaining\nfirst round#156\n
            """
            _table_content = [_table_row.split("#") for _table_row in _table_text.strip("\n").split("\n")]
            _table_pd = pd.DataFrame.from_records(_table_content[1:], columns=_table_content[0])
            return _table_pd

        questions = examples["statement"]
        tables = list(map(_convert_table_text_to_pandas, examples["table_text"]))
        result = tokenizer(tables, questions, padding=padding, max_length=max_seq_length, truncation=True)

        result["label"] = examples["label"]
        return result

    with training_args.main_process_first(desc="dataset map pre-processing"):
        raw_datasets = raw_datasets.map(
            preprocess_tabfact_function,
            batched=True,
            load_from_cache_file=not data_args.overwrite_cache,
            desc="Running tokenizer on dataset",
        )
    if training_args.do_train:
        if "train" not in raw_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = raw_datasets["train"]
        if data_args.max_train_samples is not None:
            train_dataset = train_dataset.select(range(data_args.max_train_samples))

    if training_args.do_eval:
        if "validation" not in raw_datasets and "validation_matched" not in raw_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = raw_datasets["validation"]
        if data_args.max_eval_samples is not None:
            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))

    if training_args.do_predict or data_args.test_file is not None:
        if "test" not in raw_datasets and "test_matched" not in raw_datasets:
            raise ValueError("--do_predict requires a test dataset")
        predict_dataset = raw_datasets["test"]
        if data_args.max_predict_samples is not None:
            predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))

    # Log a few random samples from the training set:
    if training_args.do_train:
        for index in random.sample(range(len(train_dataset)), 3):
            logger.info(f"Sample {index} of the training set: {train_dataset[index]}.")

    # You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a
    # predictions and label_ids field) and has to return a dictionary string to float.
    def compute_metrics(p: EvalPrediction):
        preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
        preds = np.argmax(preds, axis=1)
        return {"accuracy": (preds == p.label_ids).astype(np.float32).mean().item()}

    # Data collator will default to DataCollatorWithPadding, so we change it if we already did the padding.
    if data_args.pad_to_max_length:
        data_collator = default_data_collator
    elif training_args.fp16:
        data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
    else:
        data_collator = None

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        compute_metrics=compute_metrics,
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
        metrics = train_result.metrics
        max_train_samples = (
            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        )
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

        trainer.save_model()  # Saves the tokenizer too for easy upload

        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()

    # Evaluation
    if training_args.do_eval:
        logger.info("*** Evaluate ***")

        metrics = trainer.evaluate(eval_dataset=eval_dataset)
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))

        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)

    if training_args.do_predict:
        logger.info("*** Predict ***")

        # Removing the `label` columns because it contains -1 and Trainer won't like that.
        predict_dataset = predict_dataset.remove_columns("label")
        predictions = trainer.predict(predict_dataset, metric_key_prefix="predict").predictions
        predictions = np.argmax(predictions, axis=1)

        output_predict_file = os.path.join(training_args.output_dir, "predict_results_tabfact.txt")
        if trainer.is_world_process_zero():
            with open(output_predict_file, "w") as writer:
                logger.info("***** Predict Results *****")
                writer.write("index\tprediction\n")
                for index, item in enumerate(predictions):
                    item = label_list[item]
                    writer.write(f"{index}\t{item}\n")

    kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "text-classification"}

    if training_args.push_to_hub:
        trainer.push_to_hub(**kwargs)
    else:


SUMMARY:
Here is a concise summary of the code:

**Overview**

This code defines a text classification model that predicts one of two labels from a dataset. The model uses a pre-trained tokenizer and combines it with a custom compute function to make predictions.

**Key Components**

1. **Data Loader**: Loads the training, validation, and test datasets.
2. **Tokenization**: Pre-processes the text data for tokenization using a custom tokenizer.
3. **Model Training**: Trains the model on the training dataset using the Trainer class from Hugging Face.
4. **Model Evaluation**: Evaluates the model on the validation dataset to measure its accuracy.
5. **Prediction**: Uses the model to make predictions on the test dataset and saves the results.

**Customizations**

1. **Data Collator**: Defines a custom data collator using `DataCollatorWithPadding` to pad sequences to multiple of 8 tokens.
2. **Compute Metrics**: Defines a custom compute metrics function to calculate accuracy from predicted labels.

**Output**

The code saves the model, evaluation metrics, and prediction results in various locations (e.g., output directory, model card).

Overall, this code provides a well-structured framework for training, evaluating, and predicting text classification models using Hugging Face's Trainer class.


key: chroma:document
string_value:
CODE:

def _mp_fn(index):
    # For xla_spawn (TPUs)


SUMMARY:
**Function Summary: `_mp_fn`**

* **Purpose**: Calls the `main` function using `xla_spawn`, a TPUs execution framework.

**Detailed Breakdown:**
 
- This is a private method `_mp_fn` which takes an `index` parameter but does not utilize it.
- The primary purpose is to execute the `main` function.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/jax-projects/big_bird/prepare_natural_questions.py


key: chroma:document
string_value:
CODE:

def _get_single_answer(example):
    def choose_first(answer, is_long_answer=False):
        assert isinstance(answer, list)
        if len(answer) == 1:
            answer = answer[0]
            return {k: [answer[k]] for k in answer} if is_long_answer else answer
        for a in answer:
            if is_long_answer:
                a = {k: [a[k]] for k in a}
            if len(a["start_token"]) > 0:
                break
        return a

    answer = {"id": example["id"]}
    annotation = example["annotations"]
    yes_no_answer = annotation["yes_no_answer"]
    if 0 in yes_no_answer or 1 in yes_no_answer:
        answer["category"] = ["yes"] if 1 in yes_no_answer else ["no"]
        answer["start_token"] = answer["end_token"] = []
        answer["start_byte"] = answer["end_byte"] = []
        answer["text"] = ["<cls>"]
    else:
        answer["category"] = ["short"]
        out = choose_first(annotation["short_answers"])
        if len(out["start_token"]) == 0:
            # answer will be long if short is not available
            answer["category"] = ["long"]
            out = choose_first(annotation["long_answer"], is_long_answer=True)
            out["text"] = []
        answer.update(out)

    # disregard some samples
    if len(answer["start_token"]) > 1 or answer["start_token"] == answer["end_token"]:
        answer["remove_it"] = True
    else:
        answer["remove_it"] = False

    cols = ["start_token", "end_token", "start_byte", "end_byte", "text"]
    if not all(isinstance(answer[k], list) for k in cols):
        raise ValueError("Issue in ID", example["id"])



SUMMARY:
**Summary: `_get_single_answer` Function**

**Main Purpose and Functionality:**
The `_get_single_answer` function takes an `example` dictionary as input and returns a single answer with relevant information (category, start/end tokens/base64 byte pairs, text). It categorizes the answer into "yes", "no", or "short" based on annotation values.

**Key Implementation Details:**

* Uses `choose_first` helper function to select the first matching short-answer token.
* Categorizes answer as "long" if no short-answer tokens are available and at least one long-answer exists.
* Disregards samples with multiple start tokens or incorrect boundary conditions.

**Notable Dependencies and Requirements:**

* `annotation["yes_no_answer"]` contains boolean values indicating the overall answer.
* `annotation["short_answers"]` and `annotation["long_answer"]` contain lists of token dictionaries.
* External dependencies:
	+ No external modules are required.


key: chroma:document
string_value:
CODE:

def get_context_and_ans(example, assertion=False):
    """Gives new context after removing <html> & new answer tokens as per new context"""
    answer = _get_single_answer(example)
    # bytes are of no use
    del answer["start_byte"]
    del answer["end_byte"]

    # handle yes_no answers explicitly
    if answer["category"][0] in ["yes", "no"]:  # category is list with one element
        doc = example["document"]["tokens"]
        context = []
        for i in range(len(doc["token"])):
            if not doc["is_html"][i]:
                context.append(doc["token"][i])
        return {
            "context": " ".join(context),
            "answer": {
                "start_token": -100,  # ignore index in cross-entropy
                "end_token": -100,  # ignore index in cross-entropy
                "category": answer["category"],
                "span": answer["category"],  # extra
            },
        }

    # later, help in removing all no answers
    if answer["start_token"] == [-1]:
        return {
            "context": "None",
            "answer": {
                "start_token": -1,
                "end_token": -1,
                "category": "null",
                "span": "None",  # extra
            },
        }

    # handling normal samples

    cols = ["start_token", "end_token"]
    answer.update({k: answer[k][0] if len(answer[k]) > 0 else answer[k] for k in cols})  # e.g. [10] == 10

    doc = example["document"]["tokens"]
    start_token = answer["start_token"]
    end_token = answer["end_token"]

    context = []
    for i in range(len(doc["token"])):
        if not doc["is_html"][i]:
            context.append(doc["token"][i])
        else:
            if answer["start_token"] > i:
                start_token -= 1
            if answer["end_token"] > i:
                end_token -= 1
    new = " ".join(context[start_token:end_token])

    # checking above code
    if assertion:
        """checking if above code is working as expected for all the samples"""
        is_html = doc["is_html"][answer["start_token"] : answer["end_token"]]
        old = doc["token"][answer["start_token"] : answer["end_token"]]
        old = " ".join([old[i] for i in range(len(old)) if not is_html[i]])
        if new != old:
            print("ID:", example["id"])
            print("New:", new, end="\n")
            print("Old:", old, end="\n\n")

    return {
        "context": " ".join(context),
        "answer": {
            "start_token": start_token,
            "end_token": end_token - 1,  # this makes it inclusive
            "category": answer["category"],  # either long or short
            "span": new,  # extra
        },


SUMMARY:
**Function Summary: get_context_and_ans**

**Purpose and Functionality**
The `get_context_and_ans` function generates new context and answers for given examples. It removes HTML tokens and provides yes/no answers explicitly. The function also handles normal samples, updating answer tokens to include or exclude certain range of tokens.

### Key Implementation Details

*   Uses `_get_single_answer` to retrieve the single answer from an example.
*   Handles yes/no answers explicitly by excluding relevant tokens.
*   Updates answer tokens (`start_token` and `end_token`) for handling non-yes/no answers.
*   Verifies functionality through assertions, checking if new context generation matches expected results.

### Notable Dependencies or Requirements

*   Access to the `_get_single_answer` function.
*   Example data structure with "document" containing "tokens".
*   External dependencies (not applicable in this snippet).

### Output Format

The function returns a dictionary with two main keys: `"context"` and `"answer"`. The `"context"` key contains the new context as plain text, whereas the `"answer"` key includes additional information like the `start_token`, `end_token`, and answer category.


key: chroma:document
string_value:
CODE:

def get_strided_contexts_and_ans(example, tokenizer, doc_stride=2048, max_length=4096, assertion=True):
    # overlap will be of doc_stride - q_len

    out = get_context_and_ans(example, assertion=assertion)
    answer = out["answer"]

    # later, removing these samples
    if answer["start_token"] == -1:
        return {
            "example_id": example["id"],
            "input_ids": [[-1]],
            "labels": {
                "start_token": [-1],
                "end_token": [-1],
                "category": ["null"],
            },
        }

    input_ids = tokenizer(example["question"]["text"], out["context"]).input_ids
    q_len = input_ids.index(tokenizer.sep_token_id) + 1

    # return yes/no
    if answer["category"][0] in ["yes", "no"]:  # category is list with one element
        inputs = []
        category = []
        q_indices = input_ids[:q_len]
        doc_start_indices = range(q_len, len(input_ids), max_length - doc_stride)
        for i in doc_start_indices:
            end_index = i + max_length - q_len
            slice = input_ids[i:end_index]
            inputs.append(q_indices + slice)
            category.append(answer["category"][0])
            if slice[-1] == tokenizer.sep_token_id:
                break

        return {
            "example_id": example["id"],
            "input_ids": inputs,
            "labels": {
                "start_token": [-100] * len(category),
                "end_token": [-100] * len(category),
                "category": category,
            },
        }

    splitted_context = out["context"].split()
    complete_end_token = splitted_context[answer["end_token"]]
    answer["start_token"] = len(
        tokenizer(
            " ".join(splitted_context[: answer["start_token"]]),
            add_special_tokens=False,
        ).input_ids
    )
    answer["end_token"] = len(
        tokenizer(" ".join(splitted_context[: answer["end_token"]]), add_special_tokens=False).input_ids
    )

    answer["start_token"] += q_len
    answer["end_token"] += q_len

    # fixing end token
    num_sub_tokens = len(tokenizer(complete_end_token, add_special_tokens=False).input_ids)
    if num_sub_tokens > 1:
        answer["end_token"] += num_sub_tokens - 1

    old = input_ids[answer["start_token"] : answer["end_token"] + 1]  # right & left are inclusive
    start_token = answer["start_token"]
    end_token = answer["end_token"]

    if assertion:
        """This won't match exactly because of extra gaps => visaully inspect everything"""
        new = tokenizer.decode(old)
        if answer["span"] != new:
            print("ISSUE IN TOKENIZATION")
            print("OLD:", answer["span"])
            print("NEW:", new, end="\n\n")

    if len(input_ids) <= max_length:
        return {
            "example_id": example["id"],
            "input_ids": [input_ids],
            "labels": {
                "start_token": [answer["start_token"]],
                "end_token": [answer["end_token"]],
                "category": answer["category"],
            },
        }

    q_indices = input_ids[:q_len]
    doc_start_indices = range(q_len, len(input_ids), max_length - doc_stride)

    inputs = []
    answers_start_token = []
    answers_end_token = []
    answers_category = []  # null, yes, no, long, short
    for i in doc_start_indices:
        end_index = i + max_length - q_len
        slice = input_ids[i:end_index]
        inputs.append(q_indices + slice)
        assert len(inputs[-1]) <= max_length, "Issue in truncating length"

        if start_token >= i and end_token <= end_index - 1:
            start_token = start_token - i + q_len
            end_token = end_token - i + q_len
            answers_category.append(answer["category"][0])  # ["short"] -> "short"
        else:
            start_token = -100
            end_token = -100
            answers_category.append("null")
        new = inputs[-1][start_token : end_token + 1]

        answers_start_token.append(start_token)
        answers_end_token.append(end_token)
        if assertion:
            """checking if above code is working as expected for all the samples"""
            if new != old and new != [tokenizer.cls_token_id]:
                print("ISSUE in strided for ID:", example["id"])
                print("New:", tokenizer.decode(new))
                print("Old:", tokenizer.decode(old), end="\n\n")
        if slice[-1] == tokenizer.sep_token_id:
            break

    return {
        "example_id": example["id"],
        "input_ids": inputs,
        "labels": {
            "start_token": answers_start_token,
            "end_token": answers_end_token,
            "category": answers_category,
        },


SUMMARY:
**Summary: get_strided_contexts_and_ans function**

This function calculates context and answer pairs for a given example, splitting the context into sections of predefined length (`doc_stride`) using a query token index (`q_len`). The function follows these primary purposes:

* Calculate input IDs for the `input_ids` field in the output dictionary
* Extract category information from the `answer` object

**Key implementation details:**

* Uses tokenizer to split context and extract input IDs
* Calculates start and end token positions based on the query token index (`q_len`)
* Checks for special cases (e.g. category "yes" or "no", context splitting errors)

**Notable dependencies and requirements:**

* `tokenizer` object with methods `(example["question"]["text"], out["context"])`
* External dependency on `input_ids` field in the output dictionary
* Assumption that `answer` object contains key fields (`start_token`, `end_token`, `category`)
 

### Code snippet extraction based on provided comments


key: chroma:document
string_value:
CODE:

def prepare_inputs(example, tokenizer, doc_stride=2048, max_length=4096, assertion=False):
    example = get_strided_contexts_and_ans(
        example,
        tokenizer,
        doc_stride=doc_stride,
        max_length=max_length,
        assertion=assertion,
    )



SUMMARY:
**Summary:**

The `prepare_inputs` function prepares input data for a model, utilizing the provided `tokenizer` and document-related settings (`doc_stride`, `max_length`). It retrieves strided contexts and answers from `example` using `get_strided_contexts_and_ans`, returning the resulting `example` with the formatted data.

**Key Implementation Details:**

* Utilizes `get_strided_contexts_and_ans` function to retrieve strided contexts and answers.
* Takes optional parameters:
	+ `doc_stride`: default value of 2048
	+ `max_length`: default value of 4096
	+ `assertion`: default value of False

**Notable Dependencies/Requirements:**

* The `get_strided_contexts_and_ans` function, which is not shown in the provided code snippet.
* The `tokenizer`, with exact name specified as a parameter.


key: chroma:document
string_value:
CODE:

def save_to_disk(hf_data, file_name):
    with jsonlines.open(file_name, "a") as writer:
        for example in tqdm(hf_data, total=len(hf_data), desc="Saving samples ... "):
            labels = example["labels"]
            for ids, start, end, cat in zip(
                example["input_ids"],
                labels["start_token"],
                labels["end_token"],
                labels["category"],
            ):
                if start == -1 and end == -1:
                    continue  # leave waste samples with no answer
                if cat == "null" and np.random.rand() < 0.6:
                    continue  # removing 50 % samples
                writer.write(
                    {
                        "input_ids": ids,
                        "start_token": start,
                        "end_token": end,
                        "category": CATEGORY_MAPPING[cat],
                    }


SUMMARY:
**Concise Summary:**

### Function Purpose:

`save_to_disk` saves a subset of Hugging Face dataset data (`hf_data`) to disk in JSON line format. The function filters out samples with no answer and removes additional samples with a 50% chance of occurring.

### Key Implementation Details:

* Uses `jsonlines` library for efficient disk writing.
* Iterates over examples using `tqdm`, adding a progress bar description ("Saving samples ...").
* Filters out waste samples (no answer) and removes additional samples with a probability (CATERY_MAPPING["category"] == "null" and np.random.rand() < 0.6).
* Writes filtered data to disk in the following format:
`{"input_ids": IDs, "start_token": START, "end_token": END, "category": CATEGORY}`.

### Notable Dependencies/Requirements:

* `jsonlines`: library for efficient JSON line file writing.
* `huggingface-datasets`: library providing Hugging Face dataset data.
* `numpy`: for random number generation (`np.random.rand()`).
* `tqdm`: library for progress bar functionality.
  
### External Dependencies:

* `CATEGORY_MAPPING`: external dictionary referencing category IDs (not shown in the code snippet).


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/jax-projects/big_bird/bigbird_flax.py


key: chroma:document
string_value:
CODE:

class FlaxBigBirdForNaturalQuestionsModule(FlaxBigBirdForQuestionAnsweringModule):
    """
    BigBirdForQuestionAnswering with CLS Head over the top for predicting category

    This way we can load its weights with FlaxBigBirdForQuestionAnswering
    """

    config: BigBirdConfig
    dtype: jnp.dtype = jnp.float32
    add_pooling_layer: bool = True

    def setup(self):
        super().setup()
        self.cls = nn.Dense(5, dtype=self.dtype)

    def __call__(self, *args, **kwargs):
        outputs = super().__call__(*args, **kwargs)
        cls_out = self.cls(outputs[2])


SUMMARY:
**Class Summary: FlaxBigBirdForNaturalQuestionsModule**

The `FlaxBigBirdForNaturalQuestionsModule` is a variant of the `FlaxBigBirdForQuestionAnsweringModule`. It features a Custom Language Model (CLSM) head with a dense layer on top for predicting categories.

**Key Features:**

* Inherited configuration from `BigBirdConfig`
* Uses `jnp.float32` data type
* Adds a pooling layer (`add_pooling_layer: bool = True`)
* Custom `cls` Dense layer is added during setup

**Implementation Details:**

* Calls parent's `setup` method and creates the custom `cls` layer
* In the call method, extracts outputs from super class, applies `cls` layer to outputs[2], and returns combined output with the new category prediction.

**Dependencies:**

* **Inherited Interfaces:** `BigBirdForQuestionAnsweringModule`
* External Dependencies:
    * `flax`
    * `jnp`
    * `nn` (Neural Network)


key: chroma:document
string_value:
CODE:

class FlaxBigBirdForNaturalQuestions(FlaxBigBirdForQuestionAnswering):


SUMMARY:
**Summary:**

The `FlaxBigBirdForNaturalQuestions` class extends the `FlaxBigBirdForQuestionAnswering` base class, creating a variant for natural language question answering tasks. It inherits the `module_class` from `FlaxBigBirdForQuestionAnsweringModule`, and reuses the same underlying architecture.

**Key Details:**

* Extends `FlaxBigBirdForQuestionAnswering`
* Reuses `FlaxBigBirdForQuestionAnsweringModule` under the hood
* Primary purpose is to solve natural language question answering tasks

**Notable Dependencies/Requirements:**

* Inherit`s FlaxBigBirdForQuestionAnsweringModule`

**Design Choices:**

* Extends existing base class architecture


key: chroma:document
string_value:
CODE:

def calculate_loss_for_nq(start_logits, start_labels, end_logits, end_labels, pooled_logits, pooler_labels):
    def cross_entropy(logits, labels, reduction=None):
        """
        Args:
            logits: bsz, seqlen, vocab_size
            labels: bsz, seqlen
        """
        vocab_size = logits.shape[-1]
        labels = (labels[..., None] == jnp.arange(vocab_size)[None]).astype("f4")
        logits = jax.nn.log_softmax(logits, axis=-1)
        loss = -jnp.sum(labels * logits, axis=-1)
        if reduction is not None:
            loss = reduction(loss)
        return loss

    cross_entropy = partial(cross_entropy, reduction=jnp.mean)
    start_loss = cross_entropy(start_logits, start_labels)
    end_loss = cross_entropy(end_logits, end_labels)
    pooled_loss = cross_entropy(pooled_logits, pooler_labels)


SUMMARY:
**Summary:**
The `calculate_loss_for_nq` function calculates the average loss of three types of logits: start, end, and pooled. It utilizes a custom `cross_entropy` function that applies logsoftmax on logits and computes the cross-entropy loss between predicted and true labels.

**Key Implementation Details:**

* Custom `cross_entropy` function with partial application for mean reduction.
* Logsoftmax applied to individual logits using `jax.nn.log_softmax`.
* Loss calculation involves summing logits and labels element-wise, followed by reduction if specified.

**Notable Dependencies:**

* JAX library
* NumPy
* Scipy (for partial application)

**Function Signature:**
`calculate_loss_for_nq(start_logits, start_labels, end_logits, end_labels, pooled_logits, pooler_labels)`

**Primary Purpose:**
Calculate the average loss for a sequence-to-sequence task.


key: chroma:document
string_value:
CODE:

def get_batched_dataset(dataset, batch_size, seed=None):
    if seed is not None:
        dataset = dataset.shuffle(seed=seed)
    for i in range(len(dataset) // batch_size):
        batch = dataset[i * batch_size : (i + 1) * batch_size]


SUMMARY:
**get_batched_dataset Function Summary**

The `get_batched_dataset` function batches a given dataset into smaller chunks of size `batch_size`. It also optionally shuffle the dataset before batching.

**Key Implementation Details:**

* The function uses the `shuffle` method to randomly rearrange the dataset elements, if a seed is provided.
* It yields dictionary representations of each batch using a generator expression that slices the dataset into batches of specified size.

**Primary Purpose:**
To generate batches of data from an input dataset.

**External Dependencies:**

* The dataset is assumed to be from a library that supports shuffling (e.g., TensorFlow Datasets).
* No external library dependencies are listed explicitly, but it's worth noting that NumPy is often used for array operations.


key: chroma:document
string_value:
CODE:

class TrainState(train_state.TrainState):


SUMMARY:
**Summary**

This is an extension of the `TrainState` class from a parent module (`train_state`) and defines a subclass named `TrainState`. 

**Key Implementation Details**

* Specifies a `loss_fn` field, which is a callable function defined using `struct.field`, with a `pytree_node=False` parameter.

**Notable Dependencies or Requirements**

* Extends `train_state.TrainState`.
* Uses the `struct` module for defining fields.
* Assumes knowledge of type hints and callables with the `Callable` type.


key: chroma:document
string_value:
CODE:

def restore_checkpoint(save_dir, state):
    print(f"RESTORING CHECKPOINT FROM {save_dir}", end=" ... ")
    with open(os.path.join(save_dir, "flax_model.msgpack"), "rb") as f:
        params = from_bytes(state.params, f.read())

    with open(os.path.join(save_dir, "opt_state.msgpack"), "rb") as f:
        opt_state = from_bytes(state.opt_state, f.read())

    args = joblib.load(os.path.join(save_dir, "args.joblib"))
    data_collator = joblib.load(os.path.join(save_dir, "data_collator.joblib"))

    with open(os.path.join(save_dir, "training_state.json"), "r") as f:
        training_state = json.load(f)
    step = training_state["step"]

    print("DONE")


SUMMARY:
**Function Summary: `restore_checkpoint`**

**Purpose and Functionality:** Restores a Flax model checkpoint from a given directory.

**Key Implementation Details:**

* Loads Flax model parameters (`params`) and optimizer state (`opt_state`) from `.msgpack` files
* Loads additional models (`args` and `data_collator`) using joblib
* Reloads training state from a JSON file to retrieve the current step

**Notable Dependencies or Requirements:**

* External dependencies:
	+ `flax`
	+ `joblib`
	+ `os`
	+ `json`

**Function Signature:** `restore_checkpoint(save_dir, state)`


key: chroma:document
string_value:
CODE:

def scheduler_fn(lr, init_lr, warmup_steps, num_train_steps):
    decay_steps = num_train_steps - warmup_steps
    warmup_fn = optax.linear_schedule(init_value=init_lr, end_value=lr, transition_steps=warmup_steps)
    decay_fn = optax.linear_schedule(init_value=lr, end_value=1e-7, transition_steps=decay_steps)
    lr = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[warmup_steps])


SUMMARY:
**Summary**

`scheduler_fn` function calculates adaptive learning rate for training.

**Main Purpose and Functionality**
The function takes four parameters: learning rate (`lr`), initial learning rate (`init_lr`), warm-up steps, and number of training steps. It returns an optimal learning rate based on a combined linear schedule with warm up and decay phases.

**Key Implementation Details**

* A combination of two `optax.linear_schedule`s is created:
	+ Warm up schedule (with 4 parameters: `init_value`, `end_value`, `transition_steps`, `boundaries`)
	- Decay schedule (with 3 parameters: `init_value`, `end_value`, `transition_steps`)
* The combined schedule is returned as a single value.

**Notable Dependencies or Requirements**

* `optax.linear_schedule` (built-in)
* None

**Example Usage**
```python
lr = scheduler_fn(0.01, 1e-2, 100, 500)
```


key: chroma:document
string_value:
CODE:

def build_tx(lr, init_lr, warmup_steps, num_train_steps, weight_decay):
    def weight_decay_mask(params):
        params = traverse_util.flatten_dict(params)
        mask = {k: (v[-1] != "bias" and v[-2:] != ("LayerNorm", "scale")) for k, v in params.items()}
        return traverse_util.unflatten_dict(mask)

    lr = scheduler_fn(lr, init_lr, warmup_steps, num_train_steps)

    tx = optax.adamw(learning_rate=lr, weight_decay=weight_decay, mask=weight_decay_mask)


SUMMARY:
**Function Summary: `build_tx`**

Builds and configures the AdamWeightDecay optimizer with schedule settings.

**Main Purpose and Functionality:**
The function creates an instance of the AdamW optimizer, adapting to weight decay and schedule requirements.

**Key Implementation Details:**

* **Schedule Determination:** Uses the `scheduler_fn` function, which takes the initial learning rate (`lr`), initial value (`init_lr`), warmup steps, and total training steps as inputs.
* **Optimization Settings:** Configures the optimizer with weight decay (`weight_decay`) and a custom mask function (`weight_decay_mask`) to filter out specific weights.

**Notable Dependencies and Requirements:**

* `schedular_fn`: A function not provided in this snippet, assumed to be part of the framework or toolkit.
* `traverse_util.flatten_dict` and `unflatten_dict`: Utility functions likely part of the same framework or toolkit.
* External dependencies may include TensorFlow, PyTorch, or other similar libraries.

**Method Signatures and Interfaces:**

* The function returns a tuple containing an instance of the AdamW optimizer (`tx`) and the updated learning rate (`lr`).

This summary provides a concise overview of the `build_tx` function's main purpose, key implementation details, and notable dependencies, allowing developers to quickly understand its functionality and requirements.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/jax-projects/big_bird/evaluate.py


key: chroma:document
string_value:
CODE:

def get_sub_answers(answers, begin=0, end=None):


SUMMARY:
**Summary:**

The `get_sub_answers` function extracts and formats sub-answers from input lists. It takes advantage of list slicing (`x.split(" ")[begin:end]`) to extract specific parts of text entries.

**Implementation Details:**

* Slices text strings using `list comprehension`
* Filters out empty strings (less than 2 words) from the input

**Notable Dependencies and Requirements:**

* None, as this code snippet is self-contained.


key: chroma:document
string_value:
CODE:

def expand_to_aliases(given_answers, make_sub_answers=False):
    if make_sub_answers:
        # if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
        # *e.g.* if the correct answer contains a prefix such as "the", or "a"
        given_answers = (
            given_answers + get_sub_answers(given_answers, begin=1) + get_sub_answers(given_answers, end=-1)
        )
    answers = []
    for answer in given_answers:
        alias = answer.replace("_", " ").lower()
        alias = "".join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else " " for c in alias)
        answers.append(" ".join(alias.split()).strip())


SUMMARY:
**Function Summary**

`expand_to_aliases`: Expands given answers into their corresponding aliases by removing underscores, converting to lower case, and excluding punctuation.

**Primary Purpose**: Expand input answers into their variant aliases for potential match detection.

**Key Implementation Details**:

* Prepends and appends sub-answers from `get_sub_answers()` method to original list.
* Removes underscores, converts to lower case and excludes punctuation from each answer.
* Normalizes full words by splitting on whitespace and trimming.

**Dependencies and Requirements**:

* `PUNCTUATION_SET_TO_EXCLUDE` (external set of excluded punctuation characters).
* `get_sub_answers()` function (not included in this summary).

`expand_to_aliases(given_answers, make_sub_answers=False)`


key: chroma:document
string_value:
CODE:

def get_best_valid_start_end_idx(start_scores, end_scores, top_k=1, max_size=100):
    best_start_scores, best_start_idx = jax.lax.top_k(start_scores, top_k)
    best_end_scores, best_end_idx = jax.lax.top_k(end_scores, top_k)

    widths = best_end_idx[:, None] - best_start_idx[None, :]
    mask = jnp.logical_or(widths < 0, widths > max_size)
    scores = (best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
    best_score = jnp.argmax(scores).item()



SUMMARY:
**Summary:**

The `get_best_valid_start_end_idx` function finds the optimal starting and ending indices for sequences based on their score predictions. It takes in:

* `start_scores` and `end_scores`: predicted scores for sequences
* `top_k`: number of top-scoring indices to consider (default 1)
* `max_size`: maximum size of a sequence

The function returns the best starting and ending indices, ensuring valid widths within `max_size`.

**Key Implementation Details:**

* Uses JAX's `top_k` reduction to select top-scoring start and end indices
* Computes valid widths by subtracting start indices from end indices
* Applies a mask to invalid widths and adjusts scores accordingly

**Notable Dependencies/Requirements:**

* `jax.lax.top_k`
* `jnp.logical_or`
* External dependency on JAX library for `top_k` reduction


key: chroma:document
string_value:
CODE:

def format_dataset(sample):
    question = sample["question"]["text"]
    context = sample["document"]["tokens"]["token"]
    is_html = sample["document"]["tokens"]["is_html"]
    long_answers = sample["annotations"]["long_answer"]
    short_answers = sample["annotations"]["short_answers"]

    context_string = " ".join([context[i] for i in range(len(context)) if not is_html[i]])

    # 0 - No ; 1 - Yes
    for answer in sample["annotations"]["yes_no_answer"]:
        if answer == 0 or answer == 1:
            return {
                "question": question,
                "context": context_string,
                "short": [],
                "long": [],
                "category": "no" if answer == 0 else "yes",
            }

    short_targets = []
    for s in short_answers:
        short_targets.extend(s["text"])
    short_targets = list(set(short_targets))

    long_targets = []
    for s in long_answers:
        if s["start_token"] == -1:
            continue
        answer = context[s["start_token"] : s["end_token"]]
        html = is_html[s["start_token"] : s["end_token"]]
        new_answer = " ".join([answer[i] for i in range(len(answer)) if not html[i]])
        if new_answer not in long_targets:
            long_targets.append(new_answer)

    category = "long_short" if len(short_targets + long_targets) > 0 else "null"

    return {
        "question": question,
        "context": context_string,
        "short": short_targets,
        "long": long_targets,
        "category": category,


SUMMARY:
**Function Summary: format_dataset**

### Purpose and Functionality

Extracts annotation data from a sample dataset to create a formatted output. The function takes in a `sample` dictionary containing question, context, and annotation data.

### Key Implementation Details

- Extracts relevant data from the input `sample`.
- Preprocesses the `context` string by removing HTML tokens.
- Handles yes/no answers, defaulting to "no" if no answer is provided.
- Identifies short targets from `short_answers` and long targets from `long_answers`, ensuring unique answers.
- Determines the category based on the presence of either short or long targets.

### Notable Dependencies and Requirements

* `sample` dictionary containing:
	+ "question": question metadata
	+ "document": context metadata (tokens and is_html)
	+ "annotations": yes/no answer metadata, short_answer metadata, long_answer metadata
* External dependencies are not explicitly listed; however, the function assumes the existence of related classes/functions for tokenization and HTML parsing.

### Method Signatures and Interfaces

- `format_dataset(sample)`: main method signature
- Related classes/functions:
	+ Tokenization functions (not explicitly stated)
	+ HTML parsing functions (not explicitly stated)


key: chroma:document
string_value:
CODE:

def main():
    dataset = load_from_disk("natural-questions-validation")
    dataset = dataset.map(format_dataset).remove_columns(["annotations", "document", "id"])
    print(dataset)

    short_validation_dataset = dataset.filter(lambda x: (len(x["question"]) + len(x["context"])) < 4 * 4096)
    short_validation_dataset = short_validation_dataset.filter(lambda x: x["category"] != "null")

    model_id = "vasudevgupta/flax-bigbird-natural-questions"
    model = FlaxBigBirdForNaturalQuestions.from_pretrained(model_id)
    tokenizer = BigBirdTokenizerFast.from_pretrained(model_id)

    @jax.jit
    def forward(*args, **kwargs):
        start_logits, end_logits, pooled_logits = model(*args, **kwargs)
        return start_logits, end_logits, jnp.argmax(pooled_logits, axis=-1)

    def evaluate(example):
        # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
        inputs = tokenizer(
            example["question"],
            example["context"],
            return_tensors="np",
            max_length=4096,
            padding="max_length",
            truncation=True,
        )

        start_scores, end_scores, category = forward(**inputs)

        predicted_category = CATEGORY_MAPPING[category.item()]

        example["targets"] = example["long"] + example["short"]
        if example["category"] in ["yes", "no", "null"]:
            example["targets"] = [example["category"]]
        example["has_tgt"] = example["category"] != "null"
        # Now target can be: "yes", "no", "null", "list of long & short answers"

        if predicted_category in ["yes", "no", "null"]:
            example["output"] = [predicted_category]
            example["match"] = example["output"] == example["targets"]
            example["has_pred"] = predicted_category != "null"
            return example

        max_size = 38 if predicted_category == "short" else 1024
        start_score, end_score = get_best_valid_start_end_idx(
            start_scores[0], end_scores[0], top_k=8, max_size=max_size
        )

        input_ids = inputs["input_ids"][0].tolist()
        example["output"] = [tokenizer.decode(input_ids[start_score : end_score + 1])]

        answers = expand_to_aliases(example["targets"], make_sub_answers=True)
        predictions = expand_to_aliases(example["output"])

        # some preprocessing to both prediction and answer
        answers = {"".join(a.split()) for a in answers}
        predictions = {"".join(p.split()) for p in predictions}
        predictions = {s for s in predictions if s not in ["``", "''", "`", "'"]}

        # if there is a common element, it's a exact match
        example["match"] = len(list(answers & predictions)) > 0
        example["has_pred"] = predicted_category != "null" and len(predictions) > 0

        return example

    short_validation_dataset = short_validation_dataset.map(evaluate)

    total = len(short_validation_dataset)
    matched = len(short_validation_dataset.filter(lambda x: x["match"] == 1))


SUMMARY:
**Function Summary**

The `main` function is the entry point for training or evaluation of a Natural Questions dataset using the Flax BigBird model. It:

1. Loads the dataset from disk and formats it.
2. Filters out invalid categories and rows with fewer than 8 tokens in question/context.
3. Presets the model and tokenizer IDs.
4. Defines a `forward` function to compute start, end, and pooled logits using the Flax BigBird model.
5. Defines an evaluation `evaluate` function that:
	* Processes example data for each row.
	* Computes predictions and checks for exact matches between predicted answers and target answers.
	* Evaluates the dataset by counting exact matches against the total rows.
6. Trains or evaluates the model using the `short_validation_dataset`.

**Key Implementation Details**

* Uses JAX libraries for performance and concurrency.
* Preserves original input data via `return_tensors="np"` to use NumPy arrays.
* Utilizes tokenizer's `sep_token` and `max_length` parameters for text encoding and decoding.
* Compares prediction answers with target answers using the `==` operator.

**Notable Dependencies or Requirements**

* JAX library
* Flax BigBird model
* BigBirdTokenizerFast
*CATEGORY_MAPPING dictionary (referenced but not shown)
* Additional external dependencies may include:
	+ `numpy`
	+ `huggingface` libraries


key: docstring
string_value:
# We use torchvision for faster image pre-processing.
# We need to ensure faster processing speed as it can become a bottleneck on TPU


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/jax-projects/hybrid_clip/run_hybrid_clip.py


key: chroma:document
string_value:
CODE:
# We need to ensure faster processing speed as it can become a bottleneck on TPU
class Transform(torch.nn.Module):
    def __init__(self, image_size):
        super().__init__()
        self.transforms = torch.nn.Sequential(
            Resize([image_size], interpolation=InterpolationMode.BICUBIC),
            CenterCrop(image_size),
            ConvertImageDtype(torch.float),
            Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            x = self.transforms(x)


SUMMARY:
**Summary**

This is a PyTorch module class named `Transform` which inherits from `torch.nn.Module`. Its primary purpose is to preprocess an input image by applying a sequence of transformations.

### Main Purpose and Functionality

* Prepares images for downstream processing by resizing, cropping, converting data type, and normalizing pixel values.
* Implementation details are encapsulated within the `self.transforms` attribute, which is an instance of `torch.nn.Sequential`.

### Key Implementation Details

* Used torchvision modules:
	+ `Resize`
	+ `CenterCrop`
	+ `ConvertImageDtype`
	+ `Normalize`

### Notable Dependencies and Requirements

* PyTorch
* torchvision (library)
* TPU (TPU, Tensor Processing Unit) for potential deployment context


key: chroma:document
string_value:
CODE:

class ImageTextDataset(VisionDataset):
    """
    Dtaset for loading image-text data for tasks like CLIP training, Image Captioning.

    Args:
        root: (string): The root path where the dataset is stored
        file_path: (string): Path to the file containing the image_paths and associated captions.
            The expected format is jsonlines where each line is a json object containing to keys.
            `image_path`: The path to the image.
            `captions`: An `array` of captions.
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.ToTensor``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
        transforms (callable, optional): A function/transform that takes input sample and its target as entry
            and returns a transformed version.
    """

    def __init__(
        self,
        root: str,
        file_path: str,
        captions_per_image=2,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        transforms: Optional[Callable] = None,
    ):
        super().__init__(root, transforms, transform, target_transform)

        with open(file_path, "r") as f:
            examples = [json.loads(line) for line in f.readlines()]

        self.captions = []
        self.image_paths = []

        for example in examples:
            captions_subset = example["captions"][:captions_per_image]
            self.captions.extend(captions_subset)
            self.image_paths.extend([example["image_path"]] * len(captions_subset))

    def _load_image(self, idx: int):
        path = self.image_paths[idx]
        return read_image(path, mode=ImageReadMode.RGB)

    def _load_target(self, idx):
        return self.captions[idx]

    def __getitem__(self, index: int):
        image = self._load_image(index)
        target = self._load_target(index)

        if self.transforms is not None:
            image, target = self.transforms(image, target)

        return image, target

    def __len__(self) -> int:


SUMMARY:
**Class Summary: `ImageTextDataset`**

**Purpose:** A dataset class for loading image-text data, suitable for tasks like CLIP training and Image Captioning.

**Functionality:**

* Loads JSONL file containing image paths and captions
* Transforms images using a customizable transformation function (`transforms`)
* Returns transformed images and targets

**Key Implementation Details:**

* `_load_image` method loads an image from the given path
* `_load_target` method returns the corresponding caption for the loaded image
* `__getitem__` method transforms the image and target using a customizable function
* `transforms` parameter allows for optional transformation functions

**Notable Dependencies or Requirements:**

* `read_image` function (not shown in this code snippet)
* `ImageReadMode` enum (not shown in this code snippet)
* Customizable transformation functions (`transform`, `target_transform`, and `transforms`)


key: chroma:document
string_value:
CODE:

class TrainState(train_state.TrainState):
    dropout_rng: jnp.ndarray

    def replicate(self):


SUMMARY:
**Class Summary**

`TrainState` Class with Replication Functionality

### Purpose and Functionality

The `TrainState` class extends the parent class's `TrainState` and adds a `dropout_rng` attribute. The main purpose of this class is to provide a replicated state for gradient clipping and dropout regularization.

### Key Implementation Details

* The `replicate` method creates a new replica of the current state by replacing the `dropout_rng` attribute with a new random key generated using `shard_prng_key`.
* The `jax_utils.replicate` function creates a shallow copy of the original state, and `replace` method is used to update the `dropout_rng` attribute.

### Notable Dependencies

* `train_state.TrainState`: Parent class
* `jax_utils`: Provides the `replicate` function
* `shard_prng_key`: Generates a new random key for dropout regularization (exact type not specified)


key: chroma:document
string_value:
CODE:

def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
    summary_writer.scalar("train_time", train_time, step)

    train_metrics = get_metrics(train_metrics)
    for key, vals in train_metrics.items():
        tag = f"train_{key}"
        for i, val in enumerate(vals):
            summary_writer.scalar(tag, val, step - len(vals) + i + 1)

    for metric_name, value in eval_metrics.items():


SUMMARY:
**Function Summary:**

`write_metric` is a monitoring function that writes training and evaluation metrics to a scalar exporter.

**Main Purpose and Functionality:**

Writes key performance indicators (KPIs) to summarize training and evaluation processes at regular intervals.

**Key Implementation Details:**

- Uses `summary_writer.scalar` to collect KPIs from training metrics (`train_metrics`) and evaluates them at each timestamp.
- Writes `evalMetrics` as separate scalars with the `eval_` prefix.

**Notable Dependencies or Requirements:**

- `get_metrics(train_metrics)`: External function to process metric values
- `summary_writer.scalar`: Exporter interface for scalar values


key: chroma:document
string_value:
CODE:

def create_learning_rate_fn(
    train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
) -> Callable[[int], jnp.ndarray]:
    """Returns a linear warmup, linear_decay learning rate function."""
    steps_per_epoch = train_ds_size // train_batch_size
    num_train_steps = steps_per_epoch * num_train_epochs
    warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
    decay_fn = optax.linear_schedule(
        init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
    )
    schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])


SUMMARY:
**Summary**

The `create_learning_rate_fn` function creates and returns a callable learning rate scheduling function using OptaX. The schedules a linear warm-up from 0 to the initial learning rate, followed by a linear decay from the initial learning rate to 0.

**Key Implementation Details:**

* Utilizes OptaX's `linear_schedule` for warm-up and decay scheduling.
* Schedules warm-up at `num_warmup_steps` boundaries.
* Schedules decay from `num_train_steps - num_warmup_steps` to boundary of `num_warmup_steps`.

**Function Signatures:**

* Returns a `Callable[[int], jnp.ndarray)` function scheduling learning rate.

**Dependencies and Requirements:**

* OptaX library (`import optax`)
* JAX library (`import jnp`)


key: chroma:document
string_value:
CODE:

def main():
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if (
        os.path.exists(training_args.output_dir)
        and os.listdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({training_args.output_dir}) already exists and is not empty. "
            "Use --overwrite_output_dir to overcome."
        )

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    # Setup logging, we only want one process per machine to log things on the screen.
    logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
    if jax.process_index() == 0:
        transformers.utils.logging.set_verbosity_info()
    else:
        transformers.utils.logging.set_verbosity_error()

    # Set the verbosity to info of the Transformers logger (on main process only):
    logger.info(f"Training/evaluation parameters {training_args}")

    if model_args.tokenizer_name:
        tokenizer = AutoTokenizer.from_pretrained(
            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
        )
    elif model_args.text_model_name_or_path:
        tokenizer = AutoTokenizer.from_pretrained(
            model_args.text_model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
        )
    else:
        raise ValueError(
            "You are instantiating a new tokenizer from scratch. This is not supported by this script. "
            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
        )

    model = FlaxHybridCLIP.from_text_vision_pretrained(
        model_args.text_model_name_or_path,
        model_args.vision_model_name_or_path,
        seed=training_args.seed,
        dtype=getattr(jnp, model_args.dtype),
        text_from_pt=model_args.from_pt,
        vision_from_pt=model_args.from_pt,
    )
    config = model.config
    # set seed for torch dataloaders
    set_seed(training_args.seed)

    # Initialize torchvision transforms and jit them for faster processing
    preprocess = Transform(config.vision_config.image_size)
    preprocess = torch.jit.script(preprocess)

    # Initialize the image-text dataset
    train_dataset = ImageTextDataset(
        data_args.data_dir,
        data_args.train_file,
        captions_per_image=2,
        transform=preprocess,
    )

    eval_dataset = ImageTextDataset(
        data_args.data_dir,
        data_args.validation_file,
        captions_per_image=1,
        transform=preprocess,
    )

    # Store some constant
    num_epochs = int(training_args.num_train_epochs)
    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
    steps_per_epoch = len(train_dataset) // train_batch_size
    total_train_steps = steps_per_epoch * num_epochs

    # Use collate function to tokenizer the text and convert the processed images to numpy
    def collate_fn(examples):
        pixel_values = torch.stack([example[0] for example in examples]).permute(0, 2, 3, 1).numpy()
        captions = [example[1] for example in examples]
        inputs = tokenizer(
            captions, max_length=data_args.max_seq_length, padding="max_length", truncation=True, return_tensors="np"
        )

        batch = {
            "pixel_values": pixel_values,
            "input_ids": inputs["input_ids"],
            "attention_mask": inputs["attention_mask"],
        }

        return batch

    # Create data loaders
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=train_batch_size,
        shuffle=True,
        num_workers=data_args.preprocessing_num_workers,
        persistent_workers=True,
        drop_last=True,
        collate_fn=collate_fn,
    )

    eval_loader = torch.utils.data.DataLoader(
        eval_dataset,
        batch_size=eval_batch_size,
        shuffle=False,
        num_workers=data_args.preprocessing_num_workers,
        persistent_workers=True,
        drop_last=True,
        collate_fn=collate_fn,
    )

    # Enable tensorboard only on the master node
    if has_tensorboard and jax.process_index() == 0:
        summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir).joinpath("logs").as_posix())

    # Initialize our training
    rng = jax.random.PRNGKey(training_args.seed)
    rng, dropout_rng = jax.random.split(rng)

    # Create learning rate schedule
    linear_decay_lr_schedule_fn = create_learning_rate_fn(
        len(train_dataset),
        train_batch_size,
        training_args.num_train_epochs,
        training_args.warmup_steps,
        training_args.learning_rate,
    )

    # create adam optimizer
    adamw = optax.adamw(
        learning_rate=linear_decay_lr_schedule_fn,
        b1=training_args.adam_beta1,
        b2=training_args.adam_beta2,
        eps=training_args.adam_epsilon,
        weight_decay=training_args.weight_decay,
    )

    # Setup train state
    state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng)

    def cross_entropy(logits, axis):
        logprobs = jax.nn.log_softmax(logits, axis=axis)
        nll = jnp.diag(logprobs)
        ce = -jnp.mean(nll)
        return ce

    def clip_loss(similarity):
        loss = (cross_entropy(similarity, axis=0) + cross_entropy(similarity, axis=1)) / 2
        return loss

    # Define gradient update step fn
    def train_step(state, batch):
        dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)

        def compute_loss(params):
            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
            loss = clip_loss(logits)
            return loss

        grad_fn = jax.value_and_grad(compute_loss)
        loss, grad = grad_fn(state.params)
        grad = jax.lax.pmean(grad, "batch")

        new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)

        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
        metrics = jax.lax.pmean(metrics, axis_name="batch")

        return new_state, metrics

    # Define eval fn
    def eval_step(params, batch):
        logits = model(**batch, params=params, train=False)[0]
        loss = clip_loss(logits)

        # summarize metrics
        metrics = {"loss": loss}
        metrics = jax.lax.pmean(metrics, axis_name="batch")
        return metrics

    # Create parallel version of the train and eval step
    p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))
    p_eval_step = jax.pmap(eval_step, "batch")

    # Replicate the train state on each device
    state = state.replicate()

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel & distributed) = {train_batch_size}")
    logger.info(f"  Total optimization steps = {total_train_steps}")

    train_time = 0
    # Create sampling rng
    rng, input_rng = jax.random.split(rng)

    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
    for epoch in epochs:
        # ======================== Training ================================
        train_start = time.time()

        # Create sampling rng
        rng, input_rng = jax.random.split(rng)
        train_metrics = []

        steps_per_epoch = len(train_dataset) // train_batch_size
        train_step_progress_bar = tqdm(total=steps_per_epoch, desc="Training...", position=1, leave=False)
        # train
        for batch in train_loader:
            batch = shard(batch)
            state, train_metric = p_train_step(state, batch)
            train_metrics.append(train_metric)

            train_step_progress_bar.update(1)

        train_time += time.time() - train_start

        train_metric = unreplicate(train_metric)

        train_step_progress_bar.close()
        epochs.write(
            f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {train_metric['loss']}, Learning Rate:"
            f" {train_metric['learning_rate']})"
        )

        # ======================== Evaluating ==============================
        eval_metrics = []
        eval_steps = len(eval_dataset) // eval_batch_size
        eval_step_progress_bar = tqdm(total=eval_steps, desc="Evaluating...", position=2, leave=False)
        for batch in eval_loader:
            # Model forward
            batch = shard(batch)
            metrics = p_eval_step(state.params, batch)
            eval_metrics.append(metrics)

            eval_step_progress_bar.update(1)

        # normalize eval metrics
        eval_metrics = get_metrics(eval_metrics)

        eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)

        # Print metrics and update progress bar
        eval_step_progress_bar.close()
        desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']})"
        epochs.write(desc)
        epochs.desc = desc

        # Save metrics
        if has_tensorboard and jax.process_index() == 0:
            cur_step = epoch * (len(train_dataset) // train_batch_size)
            write_metric(summary_writer, train_metrics, eval_metrics, train_time, cur_step)

        # save checkpoint after each epoch and push checkpoint to the hub
        if jax.process_index() == 0:
            params = jax.device_get(unreplicate(state.params))
            model.save_pretrained(
                training_args.output_dir,
                params=params,
                push_to_hub=training_args.push_to_hub,
                commit_message=f"Saving weights and logs of epoch {epoch+1}",


SUMMARY:
The provided code is an implementation of the Hugging Face Transformers library for training and evaluating a model using the CLIP (Contrastive Learning of Visual Representations) algorithm. It appears to be a script written in Python for training a specific task on Vision Transformer models.

The main function of this script is to train and evaluate a Vision Transformer model over multiple epochs, accumulating metrics at each step. The script also saves the model's weights and logs after each epoch and pushes them to a hub.

Here are some key features of this script:

1.  **Model Training:** This script uses PyTorch for building and training a Vision Transformer model. It defines two functions: `train_step` for updating model parameters during training, and `evalStep` for evaluating the trained model's performance.
2.  **Data Preparation:** The data is prepared by creating batches and sharding them across devices (GPU or TPUs), enabling parallelization of batch processing.
3.  **Evaluation Metrics:** This script uses various evaluation metrics (loss) to assess the performance of the trained model at each evaluation step throughout training.

**Key sections:**

*   Training Loop: The main part of this script, where epochs range from a start value (`start = 0`) up until `num_epochs`, performing training and evaluating operations in that loop.
*   Data Loading and Training: At epoch iteration, it gets data through some "shard function".
*   Evaluation Metrics: Uses evaluation metrics like loss during validation periods.

**Improvement opportunities:**

1.  Error Checking for Invalid Input Dimensions
2.  Consider Implementing Batch Normalization Regularly, to Improve Speed.
3.  More Informed Regularization Strategies.
4.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/jax-projects/hybrid_clip/configuration_hybrid_clip.py


key: chroma:document
string_value:
CODE:

class HybridCLIPConfig(PretrainedConfig):
    r"""
    :class:`HybridCLIPConfig` is the configuration class to store the configuration of a
    :class:`~HybridCLIPModel`. It is used to instantiate HybridCLIPModel model according to the specified arguments,
    defining the text model and vision model configs.

    Configuration objects inherit from :class:`~transformers.PretrainedConfig` and can be used to control the model
    outputs. Read the documentation from :class:`~transformers.PretrainedConfig` for more information.

    Args:
        text_config_dict (:obj:`dict`):
            Dictionary of configuration options that defines text model config.
        vision_config_dict (:obj:`dict`):
            Dictionary of configuration options that defines vison model config.
        projection_dim (:obj:`int`, `optional`, defaults to 512):
            Dimentionality of text and vision projection layers.
        kwargs (`optional`):
            Dictionary of keyword arguments.

    Examples::

        >>> from transformers import BertConfig, CLIPConfig, HybridCLIPConfig, FlaxHybridCLIP

        >>> # Initializing a BERT and CLIP configuration
        >>> config_text = BertConfig()
        >>> config_vision = CLIPConfig()

        >>> config = HybridCLIPConfig.from_text_vision_configs(config_text, config_vision, projection_dim=512)

        >>> # Initializing a BERT and CLIPVision model
        >>> model = EncoderDecoderModel(config=config)

        >>> # Accessing the model configuration
        >>> config_text = model.config.text_config
        >>> config_vision  = model.config.vision_config

        >>> # Saving the model, including its configuration
        >>> model.save_pretrained('my-model')

        >>> # loading model and config from pretrained folder
        >>> encoder_decoder_config = HybridCLIPConfig.from_pretrained('my-model')
        >>> model = FlaxHybridCLIP.from_pretrained('my-model', config=encoder_decoder_config)
    """

    model_type = "hybrid-clip"
    is_composition = True

    def __init__(self, projection_dim=512, **kwargs):
        super().__init__(**kwargs)

        if "text_config" not in kwargs:
            raise ValueError("`text_config` can not be `None`.")

        if "vision_config" not in kwargs:
            raise ValueError("`vision_config` can not be `None`.")

        text_config = kwargs.pop("text_config")
        vision_config = kwargs.pop("vision_config")

        text_model_type = text_config.pop("model_type")
        vision_model_type = vision_config.pop("model_type")

        from transformers import AutoConfig

        self.text_config = AutoConfig.for_model(text_model_type, **text_config)

        if vision_model_type == "clip":
            self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config).vision_config
        elif vision_model_type == "clip_vision_model":
            from transformers import CLIPVisionConfig

            self.vision_config = CLIPVisionConfig(**vision_config)
        else:
            self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config)

        self.projection_dim = projection_dim
        self.initializer_factor = 1.0

    @classmethod
    def from_text_vision_configs(cls, text_config: PretrainedConfig, vision_config: PretrainedConfig, **kwargs):
        r"""
        Instantiate a :class:`HybridCLIPConfig` (or a derived class) from text model configuration and
        vision model configuration.

        Returns:
            :class:`HybridCLIPConfig`: An instance of a configuration object
        """

        return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs)

    def to_dict(self):
        """
        Serializes this instance to a Python dictionary. Override the default
        :meth:`~transformers.PretrainedConfig.to_dict`.

        Returns:
            :obj:`Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
        """
        output = copy.deepcopy(self.__dict__)
        output["text_config"] = self.text_config.to_dict()
        output["vision_config"] = self.vision_config.to_dict()
        output["model_type"] = self.__class__.model_type


SUMMARY:
**HybridCLIPConfig Class Summary**

*   **Purpose**: Configurational class for HybridCLIPModel, defining text and vision model configurations.

### Key Implementation Details

*   Initializes with `text_config`, `vision_config`, and optional `projection_dim` parameters.
*   Derives `text_config` and `vision_config` from specific model types (e.g., BERT, CLIP).
*   Handles projection dimensionality and initializes `initializer_factor`.

### Notable Dependencies and Requirements

*   Imports AutoConfig, CLIPVisionConfig, and other transformer modules.

### Interfaces and Inheritance

*   Extends PretrainedConfig.
*   Uses to_dict() method to serialize to Python dictionary.

### Design Choices

*   Allows composition of text-vision models using separate configurations.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/jax-projects/hybrid_clip/modeling_hybrid_clip.py


key: chroma:document
string_value:
CODE:

class FlaxHybridCLIPModule(nn.Module):
    config: HybridCLIPConfig
    dtype: jnp.dtype = jnp.float32

    def setup(self):
        text_config = self.config.text_config
        vision_config = self.config.vision_config

        self.projection_dim = self.config.projection_dim
        self.text_embed_dim = text_config.hidden_size
        self.vision_embed_dim = vision_config.hidden_size

        text_module = FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
        vision_module = FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__, FlaxCLIPVisionModel).module_class

        self.text_model = text_module(text_config, dtype=self.dtype)
        self.vision_model = vision_module(vision_config, dtype=self.dtype)

        self.visual_projection = nn.Dense(
            self.projection_dim,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(0.02),
            use_bias=False,
        )
        self.text_projection = nn.Dense(
            self.projection_dim,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(0.02),
            use_bias=False,
        )
        self.logit_scale = self.param("logit_scale", jax.nn.initializers.ones, [])

    def __call__(
        self,
        input_ids=None,
        pixel_values=None,
        attention_mask=None,
        position_ids=None,
        token_type_ids=None,
        deterministic: bool = True,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        vision_outputs = self.vision_model(
            pixel_values=pixel_values,
            deterministic=deterministic,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        text_outputs = self.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            deterministic=deterministic,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        image_embeds = vision_outputs[1]
        image_embeds = self.visual_projection(image_embeds)

        text_embeds = text_outputs[1]
        text_embeds = self.text_projection(text_embeds)

        # normalized features
        image_embeds = image_embeds / jnp.linalg.norm(image_embeds, axis=-1, keepdims=True)
        text_embeds = text_embeds / jnp.linalg.norm(text_embeds, axis=-1, keepdims=True)

        # cosine similarity as logits
        logit_scale = jnp.exp(self.logit_scale)
        logits_per_text = jnp.matmul(text_embeds, image_embeds.T) * logit_scale
        logits_per_image = logits_per_text.T

        if not return_dict:
            return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

        return FlaxCLIPOutput(
            logits_per_image=logits_per_image,
            logits_per_text=logits_per_text,
            text_embeds=text_embeds,
            image_embeds=image_embeds,
            text_model_output=text_outputs,
            vision_model_output=vision_outputs,


SUMMARY:
**Class Summary**

`FlaxHybridCLIPModule`
------------------------

### Purpose and Functionality

A PyTorch Hugging Face module for training a hybrid BERT + Vision Transformer (ViT) architecture using the CLIP prompt conditioning technique.

This class extends [PyTorch's `nn.Module`](https://pytorch.org/docs/stable/introduction/tensor.html) with specialized functionality for text and vision inputs, projection and cosine similarity calculations, and output processing.

### Key Implementation Details

*   **Model Initialization**: Initializes the `text_model`, `vision_model`, `visual_projection`, `text_projection`, and `logit_scale` modules.
*   **Input Processing**: Takes in input IDs, pixel values, attention masks, position IDs, token type IDs, and outputs dictionary flags. Processes text and vision inputs through their respective models, followed by projections and normalization.
*   **Logits Calculations**: Calculates image Embeds (`image_embeds`) and Text Embeds (`text_embeds`), computes cosine similarity as logits between them.

### Notable Dependencies and Requirements

*   `jnp`: JAX library for numerical computations
*   `jax.nn.initializers`: Module for initializing neural network weights
*   `FlaxModelMapping`: A dictionary mapping model classes to their corresponding module classes, implemented elsewhere in the codebase.
*   `FlaxCLIPVisionModel`: The Vision Transformer (ViT) implementation.

### Interfaces and Inheritance

Inherits from PyTorch's [nn.Module](https://pytorch.org/docs/stable/introduction/autograd.html), implements a custom method (`__call__`).

The following classes are referenced or inherited within the `FlaxHybridCLIPModule` class:

*   `FLAX_MODEL_MAPPING`: A Python dictionary implementing model mapping logic.

Example output from this module (`FlaxCLIPOutput`) can be further processed by the user.


key: chroma:document
string_value:
CODE:

class FlaxHybridCLIP(FlaxPreTrainedModel):
    config_class = HybridCLIPConfig
    module_class = FlaxHybridCLIPModule

    def __init__(
        self,
        config: HybridCLIPConfig,
        input_shape: Optional[Tuple] = None,
        seed: int = 0,
        dtype: jnp.dtype = jnp.float32,
        **kwargs,
    ):
        if input_shape is None:
            input_shape = ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

        module = self.module_class(config=config, dtype=dtype, **kwargs)
        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)

    def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> FrozenDict:
        # init input tensor
        input_ids = jnp.zeros(input_shape[0], dtype="i4")
        position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape[0])
        token_type_ids = jnp.ones_like(input_ids)
        attention_mask = jnp.ones_like(input_ids)

        pixel_values = jax.random.normal(rng, input_shape[1])

        params_rng, dropout_rng = jax.random.split(rng)
        rngs = {"params": params_rng, "dropout": dropout_rng}

        return self.module.init(rngs, input_ids, pixel_values, attention_mask, position_ids, token_type_ids)["params"]

    def __call__(
        self,
        input_ids,
        pixel_values,
        attention_mask=None,
        position_ids=None,
        token_type_ids=None,
        params: dict = None,
        dropout_rng: jax.random.PRNGKey = None,
        train: bool = False,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        if position_ids is None:
            position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)

        if token_type_ids is None:
            token_type_ids = jnp.zeros_like(input_ids)

        if attention_mask is None:
            attention_mask = jnp.ones_like(input_ids)

        # Handle any PRNG if needed
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        return self.module.apply(
            {"params": params or self.params},
            jnp.array(input_ids, dtype="i4"),
            jnp.array(pixel_values, dtype=jnp.float32),
            jnp.array(attention_mask, dtype="i4"),
            jnp.array(position_ids, dtype="i4"),
            jnp.array(token_type_ids, dtype="i4"),
            not train,
            output_attentions,
            output_hidden_states,
            return_dict,
            rngs=rngs,
        )

    def get_text_features(
        self,
        input_ids,
        attention_mask=None,
        position_ids=None,
        token_type_ids=None,
        params: dict = None,
        dropout_rng: jax.random.PRNGKey = None,
        train=False,
    ):
        r"""
        Args:
            input_ids (:obj:`numpy.ndarray` of shape :obj:`(batch_size, sequence_length)`):
                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you
                provide it.

                Indices can be obtained using :class:`~transformers.PreTrainedTokenizer`. See
                :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__`
                for details.

                `What are input IDs? <../glossary.html#input-ids>`__

        Returns:
            text_features (:obj:`jnp.ndarray` of shape :obj:`(batch_size, output_dim`): The text embeddings
            obtained by applying the projection layer to the pooled output of text model.
        """
        if position_ids is None:
            position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)

        if token_type_ids is None:
            token_type_ids = jnp.zeros_like(input_ids)

        if attention_mask is None:
            attention_mask = jnp.ones_like(input_ids)

        # Handle any PRNG if needed
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        def _get_features(module, input_ids, attention_mask, position_ids, token_type_ids, deterministic):
            text_outputs = module.text_model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                token_type_ids=token_type_ids,
                deterministic=deterministic,
            )
            pooled_output = text_outputs[1]
            text_features = module.text_projection(pooled_output)
            return text_features

        return self.module.apply(
            {"params": params or self.params},
            jnp.array(input_ids, dtype="i4"),
            jnp.array(attention_mask, dtype="i4"),
            jnp.array(position_ids, dtype="i4"),
            jnp.array(token_type_ids, dtype="i4"),
            not train,
            method=_get_features,
            rngs=rngs,
        )

    def get_image_features(
        self, pixel_values, params: dict = None, dropout_rng: jax.random.PRNGKey = None, train=False
    ):
        r"""
        Args:
            pixel_values (:obj:`numpy.ndarray` of shape :obj:`(batch_size, num_channels, height, width)`):
                Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained
                using :class:`~transformers.ImageFeatureExtractionMixin`. See
                :meth:`transformers.ImageFeatureExtractionMixin.__call__` for details.

        Returns:
            image_features (:obj:`jnp.ndarray` of shape :obj:`(batch_size, output_dim`): The image embeddings
            obtained by applying the projection layer to the pooled output of vision model.
        """

        # Handle any PRNG if needed
        rngs = {}
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

        def _get_features(module, pixel_values, deterministic):
            vision_outputs = module.vision_model(pixel_values=pixel_values, deterministic=deterministic)
            pooled_output = vision_outputs[1]  # pooled_output
            image_features = module.visual_projection(pooled_output)
            return image_features

        return self.module.apply(
            {"params": params or self.params},
            jnp.array(pixel_values, dtype=jnp.float32),
            not train,
            method=_get_features,
            rngs=rngs,
        )

    @classmethod
    def from_text_vision_pretrained(
        cls,
        text_model_name_or_path: str = None,
        vision_model_name_or_path: str = None,
        *model_args,
        **kwargs,
    ) -> FlaxPreTrainedModel:
        """
        Params:
            text_model_name_or_path (:obj: `str`, `optional`):
                Information necessary to initiate the text model. Can be either:

                    - A string, the `model id` of a pretrained model hosted inside a model repo on huggingface.co.
                    - A path to a `directory` containing model weights saved using
                      :func:`~transformers.FlaxPreTrainedModel.save_pretrained`, e.g., ``./my_model_directory/``.
                    - A path or url to a `PyTorch checkpoint folder` (e.g, ``./pt_model``). In
                      this case, ``from_pt`` should be set to :obj:`True` and a configuration object should be provided
                      as ``config`` argument. This loading path is slower than converting the PyTorch checkpoint in
                      a Flax model using the provided conversion scripts and loading the Flax model afterwards.

            vision_model_name_or_path (:obj: `str`, `optional`, defaults to `None`):
                Information necessary to initiate the vision model. Can be either:

                    - A string, the `model id` of a pretrained model hosted inside a model repo on huggingface.co.
                    - A path to a `directory` containing model weights saved using
                      :func:`~transformers.FlaxPreTrainedModel.save_pretrained`, e.g., ``./my_model_directory/``.
                    - A path or url to a `PyTorch checkpoint folder` (e.g, ``./pt_model``). In
                      this case, ``from_pt`` should be set to :obj:`True` and a configuration object should be provided
                      as ``config`` argument. This loading path is slower than converting the PyTorch checkpoint in
                      a Flax model using the provided conversion scripts and loading the Flax model afterwards.

            model_args (remaining positional arguments, `optional`):
                All remaning positional arguments will be passed to the underlying model's ``__init__`` method.

            kwargs (remaining dictionary of keyword arguments, `optional`):
                Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
                :obj:`output_attentions=True`).

                - To update the text configuration, use the prefix `text_` for each configuration parameter.
                - To update the vision configuration, use the prefix `vision_` for each configuration parameter.
                - To update the parent model configuration, do not use a prefix for each configuration parameter.

                Behaves differently depending on whether a :obj:`config` is provided or automatically loaded.

        Example::

            >>> from transformers import FlaxHybridCLIP
            >>> # initialize a model from pretrained BERT and CLIP models. Note that the projection layers will be randomly initialized.
            >>> # If using CLIP's vision model the vision projection layer will be initialized using pre-trained weights
            >>> model = FlaxHybridCLIP.from_text_vision_pretrained('google-bert/bert-base-uncased', 'openai/clip-vit-base-patch32')
            >>> # saving model after fine-tuning
            >>> model.save_pretrained("./bert-clip")
            >>> # load fine-tuned model
            >>> model = FlaxHybridCLIP.from_pretrained("./bert-clip")
        """

        kwargs_text = {
            argument[len("text_") :]: value for argument, value in kwargs.items() if argument.startswith("text_")
        }

        kwargs_vision = {
            argument[len("vision_") :]: value for argument, value in kwargs.items() if argument.startswith("vision_")
        }

        # remove text, vision kwargs from kwargs
        for key in kwargs_text.keys():
            del kwargs["text_" + key]
        for key in kwargs_vision.keys():
            del kwargs["vision_" + key]

        # Load and initialize the text and vision model
        text_model = kwargs_text.pop("model", None)
        if text_model is None:
            assert (
                text_model_name_or_path is not None
            ), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
            from transformers import FlaxAutoModel

            if "config" not in kwargs_text:
                from transformers import AutoConfig

                text_config = AutoConfig.from_pretrained(text_model_name_or_path)
                kwargs_text["config"] = text_config

            text_model = FlaxAutoModel.from_pretrained(text_model_name_or_path, *model_args, **kwargs_text)

        vision_model = kwargs_vision.pop("model", None)
        if vision_model is None:
            assert (
                vision_model_name_or_path is not None
            ), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
            from transformers import FlaxAutoModel

            if "config" not in kwargs_vision:
                from transformers import AutoConfig

                vision_config = AutoConfig.from_pretrained(vision_model_name_or_path)
                kwargs_vision["config"] = vision_config

            vision_model = FlaxAutoModel.from_pretrained(vision_model_name_or_path, *model_args, **kwargs_vision)

        # instantiate config with corresponding kwargs
        dtype = kwargs.pop("dtype", jnp.float32)
        config = HybridCLIPConfig.from_text_vision_configs(text_model.config, vision_model.config, **kwargs)

        # init model
        model = cls(config, *model_args, dtype=dtype, **kwargs)

        if vision_config.model_type == "clip":
            model.params["vision_model"]["vision_model"] = vision_model.params["vision_model"]
            model.params["visual_projection"]["kernel"] = vision_model.params["visual_projection"]["kernel"]
        else:
            model.params["vision_model"] = vision_model.params

        model.params["text_model"] = text_model.params



SUMMARY:
This is the implementation of the Hybrid CLIP model in JAX and the TFP library, which is a multi-task learning architecture for visual question answering (VQA) tasks. The code defines the model's configuration, initializes the parameters, and provides methods for computing forward pass outputs.

Here is a concise summary:

- **Model Configuration**: The `HybridCLIPConfig` class loads text and vision configurations from the specified models.
- **Model Initialization**: The `FlaxHybridCLIP` class initializes the model with a given configuration and dtype (data type).
- **Forward Pass Computing**: The code provides methods for computing forward pass outputs, including:
  - `text_embedding`: computes the text embeddings using the `text_model`.
  - `vision_embedding`: computes the vision embeddings using the `vision_model`.
  - `cross_text_vision_norm`: computes the cross-text and cross-vision norms.
- **Initialization and Training**: The model is initialized with the specified configuration, dtype, and optional keyword arguments. There is no training mechanism provided in this implementation.

Overall, this code provides a foundation for implementing the Hybrid CLIP model using JAX and TFP libraries, allowing for efficient computation of forward pass outputs.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/jax-projects/dataset-streaming/run_mlm_flax_stream.py


key: chroma:document
string_value:
CODE:

def generate_batch_splits(samples_idx: np.ndarray, batch_size: int) -> np.ndarray:
    num_samples = len(samples_idx)
    samples_to_remove = num_samples % batch_size

    if samples_to_remove != 0:
        samples_idx = samples_idx[:-samples_to_remove]
    sections_split = num_samples // batch_size
    batch_idx = np.split(samples_idx, sections_split)


SUMMARY:
**Function Summary: `generate_batch_splits`**

Generates batch splits for the given sample indices by splitting them into multiple batches.

**Purpose:** Reorders sample indices into equally sized batches with remaining samples omitted if necessary.

**Key Implementation Details:**

* Removes last `<batch_size>` samples to make up a multiple of the batch size.
* Splits the reordered array into `num_samples // batch_size` equal-length batches.

**Parameters/Interfaces:**

* `samples_idx`: Input sample indices (np.ndarray)
* `batch_size`: Desired batch size (int)

**External Dependencies:**

* NumPy (`np`) for numerical computations


key: chroma:document
string_value:
CODE:

def advance_iter_and_group_samples(train_iterator, num_samples, max_seq_length):
    """
    The training iterator is advanced so that after groupifying the samples,
    `num_samples` of length `max_seq_length` are returned.
    """
    num_total_tokens = max_seq_length * num_samples
    samples = defaultdict(list)

    i = 0
    while i < num_total_tokens:
        tokenized_samples = next(train_iterator)
        i += len(tokenized_samples["input_ids"])

        # concatenate tokenized samples to list (excluding "id" and "text")
        samples = {
            k: samples[k] + tokenized_samples[k] for k in ["input_ids", "attention_mask", "special_tokens_mask"]
        }

    # Concatenated tokens are split to lists of length `max_seq_length`.
    # Note that remainedr of % max_seq_length are thrown away.
    def group_texts(examples):
        result = {
            k: [t[i : i + max_seq_length] for i in range(0, num_total_tokens, max_seq_length)]
            for k, t in examples.items()
        }
        return result

    grouped_samples = group_texts(samples)


SUMMARY:
**Concise Summary**

`advance_iter_and_group_samples` function:

*   **Purpose**: Advances training iterator to return `num_samples` of length `max_seq_length` samples, groupifying them along the way.

*   **Key Algorithm**: Iteratively tokenizes and concatenates sample data from `train_iterator`, then groups the resulting tokens into sequences of fixed length using a custom function `group_texts`.

*   **Implementation Details**:
    *   Utilizes `defaultdict` to store concatenated samples.
    *   Applies string slicing to split tokens into valid sequence length.

*   **Dependencies and Requirements**:
    *   `train_iterator`: Must be an iterable providing tokenized sample data.
    *   `num_samples`: Integer specifying number of samples to return.
    *   `max_seq_length`: Integer specifying maximum sequence length.


key: chroma:document
string_value:
CODE:

def write_train_metric(summary_writer, train_metrics, train_time, step):
    summary_writer.scalar("train_time", train_time, step)

    train_metrics = get_metrics(train_metrics)
    for key, vals in train_metrics.items():
        tag = f"train_{key}"
        for i, val in enumerate(vals):


SUMMARY:
**Function Summary: `write_train_metric`**

**Purpose:** Writes training metrics to a summary writer.

* **Primary purpose:** Tracks and logs training time and metrics during optimization steps.

**Implementation Details:**

* Retrieves validated training metrics using `get_metrics(train_metrics)`.
* Loops through each metric key-value pair:
	+ Creates a scalar tag by prefixing with "train_".
	+ Writes scalar metrics to writer for each value in the paired list, incrementing step numbers.

**Notable Dependencies/Requirements:**

* `summary_writer`: object responsible for writing summaries to a format (e.g., tensorboard).
* External dependency:
	+ `get_metrics(train_metrics)` function (exact implementation not provided).


key: chroma:document
string_value:
CODE:

def write_eval_metric(summary_writer, eval_metrics, step):
    for metric_name, value in eval_metrics.items():


SUMMARY:
**Summary:**

The `write_eval_metric` function takes three inputs:

- `summary_writer`: an object responsible for writing metrics to a summary
- `eval_metrics`: a dictionary of evaluation metric names and their corresponding values
- `step`: a step number in the training process

**Main Purpose:** To write evaluation metrics to a summary.

**Key Implementation Details:**

* Utilizes `scalar` method on `summary_writer` to write each metric as a single scalar value.
* Iterates over `eval_metrics` dictionary using.items() and writes each pair of key-value to summaryWriter in the format "eval_X_key".

**Notable Dependencies/Affiliations:**

* `summary_writer`: an object of unspecified type, but expected to have a `scalar` method.
* `eval_metrics`: a dictionary containing evaluation metrics with string keys and numeric values. 
* No external dependencies are referenced.

This summary provides a clear understanding of the function's purpose and key implementation details, while omitting unnecessary information.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/jax-projects/wav2vec2/run_wav2vec2_pretrain_flax.py


key: chroma:document
string_value:
CODE:

def configure_logger(model_args: ModelArguments, training_args: TrainingArguments):
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )
    logging_level = logging.WARNING
    if model_args.verbose_logging:
        logging_level = logging.DEBUG


SUMMARY:
**Code Summary: Configure Logger Function**

### Main Purpose and Functionality

Configures the logger for a PyTorch model based on provided `Model` and `Training` arguments.

### Key Implementation Details

* Sets up basic logging configuration with timestamp, log level, and output to standard output.
* Logs at either WARNING or DEBUG level depending on `model_args.verbose_logging`.

### Notable Dependencies and Requirements

None explicitly mentioned, but utilizes:
- `logging` module
- `sys.stdout` for handling output

Note: No external library dependencies are listed.


key: chroma:document
string_value:
CODE:

def write_train_metric(summary_writer, train_metrics, train_time, step):
    summary_writer.scalar("train_time", train_time, step)

    train_metrics = get_metrics(train_metrics)
    for key, vals in train_metrics.items():
        tag = f"train_{key}"
        for i, val in enumerate(vals):


SUMMARY:
**Function Summary: `write_train_metric`**

writes training metrics and timestamp to a summary writer.

**Primary Purpose:** Visualize scalar metrics during training using TensorBoard-compatible summary writer, tracking time and individual step values for each metric.

**Key Implementation Details:**

*   Updates `train_metrics` with transformed data using `get_metrics`.
*   Iterates over each metric (`key`) and its corresponding values.
*   Creates a tag `tag`) to uniquely identify metrics in TensorBoard.
*   Writes scalar metrics to summary writer, referencing step number.

**Notable Dependencies/Requirements:**

*   `summary_writer`: TensorFlow summary writer or compatible implementation.
*   `get_metrics` function (not shown): applies transformation to `train_metrics`.
*   External dependencies not specified; assumes compatibility with TensorFlow library.


key: chroma:document
string_value:
CODE:

def write_eval_metric(summary_writer, eval_metrics, step):
    for metric_name, value in eval_metrics.items():


SUMMARY:
**Function Summary: `write_eval_metric`**

```python
Write evaluation metrics to the summary writer.

Purpose:
  - Writes scalar values for various evaluation metrics (e.g., accuracy, loss) at each step.
  - Utilizes a `summary_writer` object to record these values.

Key Implementation Details:

- Iterates through key-value pairs in `eval_metrics`, where keys are metric names and values represent their corresponding evaluation scores.
- Utilizes the `scalar` method of `summary_writer` to log each metric with its name, value, and step number.
- No aggregation or transformation is performed on metrics; raw values are written as-is.

Notable Dependencies:

- `summary_writer`: an object providing a scalar logging interface.
- `eval_metrics`: a dictionary containing evaluation metric names as keys and their scores as values.
```
Note: The code maintains simplicity and clarity, following efficient data structures (dictionary) for storing metrics and using the provided `summary_writer` to minimize complexity.


key: chroma:document
string_value:
CODE:

def generate_batch_splits(samples_idx: np.ndarray, batch_size: int) -> np.ndarray:
    num_samples = len(samples_idx)
    samples_to_remove = num_samples % batch_size

    if samples_to_remove != 0:
        samples_idx = samples_idx[:-samples_to_remove]
    sections_split = num_samples // batch_size
    batch_idx = np.split(samples_idx, sections_split)


SUMMARY:
**Function Summary: `generate_batch_splits`**

* **Purpose**: Splits a sequence of sample indices into batches based on a specified batch size.

* **Key Implementation Details**:
	+ Removes samples to ensure batch size is consistent across all sections.
	+ Uses the `//` operator for integer division to determine the number of sections and subsequent split of the indices array.

* **Method Signatures and Interfaces**:
	+ Input: `samples_idx` (`np.ndarray`)
	+ Output: `batch_idx` (`list` of `np.ndarray` objects)

* **Dependencies and Requirements**: 
  - No external dependencies.
  - Uses NumPy (`np`) for array manipulation.


key: chroma:document
string_value:
CODE:

def compute_contrastive_loss(
    quantized_features, transformer_features, negative_indices, mask_time_indices, logits_temp, num_negatives
):
    batch_size, sequence_length, hidden_size = quantized_features.shape

    # take negative vectors from sampled indices
    quantized_negatives = quantized_features.reshape(-1, hidden_size)[negative_indices.reshape(-1)]
    quantized_negatives = quantized_negatives.reshape(
        batch_size, sequence_length, num_negatives, hidden_size
    ).transpose(2, 0, 1, 3)

    target_features = jnp.concatenate([quantized_features[None, :], quantized_negatives], axis=0)
    loss_logits = optax.cosine_similarity(transformer_features, target_features)
    loss_logits = loss_logits / logits_temp

    neg_is_pos = (quantized_features == quantized_negatives).all(-1)
    neg_is_pos = jnp.concatenate([jnp.full((1,) + loss_logits.shape[1:], False), neg_is_pos], axis=0)

    # make sure incorrectly sampled vectors don't contribute to loss
    loss_logits = jnp.where(neg_is_pos, -1e9, loss_logits)

    predictions = loss_logits.transpose(2, 1, 0).reshape(-1, loss_logits.shape[0])
    targets = ((1 - mask_time_indices) * -100).transpose(1, 0).flatten()

    target_mask = jnp.where(targets >= 0, 1.0, 0.0)
    contrastive_loss = optax.softmax_cross_entropy(predictions, onehot(targets, predictions.shape[-1])) * target_mask

    contrastive_loss = contrastive_loss.sum()



SUMMARY:
**Compute Contrastive Loss Function**
=====================================

### Functionality Overview

This function computes the contrastive loss between quantized and transformer features using cosine similarity.

### Main Purpose

The primary purpose is to calculate the negative logistic loss between positive pairs of input features with sampled negative indices.

### Key Implementation Details

*   Computes the loss for each sample pair between `quantized_features` and `transformer_features`.
*   Applies batch normalization on both features.
*   Calculates the loss using cosine similarity via `optax.cosine_similarity`.

```python
def compute_contrastive_loss(
    quantized_features: jnp.ndarray, 
    transformer_features: jnp.ndarray, 
    negative_indices: jnp.ndarray, 
    mask_time_indices: jnp.ndarray, 
    logits_temp: float,
    num_negatives: int
):
```

### Notable Dependencies and Requirements

*   `optax`: library for optimization algorithms.
*   `jnp`: library for numerical computation (NumPy-like operations).

**Summary in Bullet Points**

•  Computes contrastive loss between positive pairs of input features with sampled negative indices.
•  Uses cosine similarity to calculate the inner product of both feature sets.
•  Applies batch normalization and calculates negative logistic loss using softmax cross-entropy.
•  Requires `optax` and `jnp` libraries for computations.


key: chroma:document
string_value:
CODE:

def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))

    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    configure_logger(model_args, training_args)

    # Downloading and loading a dataset from the hub.
    datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir)

    if "validation" not in datasets.keys():
        # make sure only "validation" and "train" keys remain"
        datasets = DatasetDict()
        datasets["validation"] = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            split=f"{data_args.train_split_name}[:{data_args.validation_split_percentage}%]",
            cache_dir=model_args.cache_dir,
        )
        datasets["train"] = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            split=f"{data_args.train_split_name}[{data_args.validation_split_percentage}%:]",
            cache_dir=model_args.cache_dir,
        )
    else:
        # make sure only "validation" and "train" keys remain"
        datasets = DatasetDict()
        datasets["validation"] = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            split="validation",
            cache_dir=model_args.cache_dir,
        )
        datasets["train"] = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            split=f"{data_args.train_split_name}",
            cache_dir=model_args.cache_dir,
        )

    # only normalized-inputs-training is supported
    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
        model_args.model_name_or_path, cache_dir=model_args.cache_dir, do_normalize=True
    )

    def prepare_dataset(batch):
        # check that all files have the correct sampling rate
        batch["speech"], _ = librosa.load(batch[data_args.speech_file_column], sr=feature_extractor.sampling_rate)
        return batch

    # load audio files into numpy arrays
    vectorized_datasets = datasets.map(
        prepare_dataset, num_proc=data_args.preprocessing_num_workers, remove_columns=datasets["train"].column_names
    )

    # filter audio files that are too long
    vectorized_datasets = vectorized_datasets.filter(
        lambda data: len(data["speech"]) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate)
    )

    def normalize(batch):
        return feature_extractor(batch["speech"], sampling_rate=feature_extractor.sampling_rate)

    # normalize and transform to `BatchFeatures`
    vectorized_datasets = vectorized_datasets.map(
        normalize,
        batched=True,
        num_proc=data_args.preprocessing_num_workers,
        load_from_cache_file=not data_args.overwrite_cache,
        remove_columns=vectorized_datasets["train"].column_names,
    )

    # pretraining is only supported for "newer" stable layer norm architecture
    # apply_spec_augment has to be True, mask_feature_prob has to be 0.0
    config = Wav2Vec2Config.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
    )

    if not config.do_stable_layer_norm or config.feat_extract_norm != "layer":
        raise ValueError(
            "PreTraining is only supported for ``config.do_stable_layer_norm=True`` and"
            " ``config.feat_extract_norm='layer'"
        )

    model = FlaxWav2Vec2ForPreTraining(config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype))

    # Activate gradient checkpointing if needed
    if training_args.gradient_checkpointing:
        model.gradient_checkpointing_enable()

    data_collator = FlaxDataCollatorForWav2Vec2Pretraining(
        model=model, feature_extractor=feature_extractor, pad_to_multiple_of=data_args.pad_to_multiple_of
    )

    # Enable tensorboard only on the master node
    has_tensorboard = is_tensorboard_available()
    if has_tensorboard and jax.process_index() == 0:
        try:
            from flax.metrics.tensorboard import SummaryWriter

            summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
        except ImportError as ie:
            has_tensorboard = False
            logger.warning(
                f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
            )
    else:
        logger.warning(
            "Unable to display metrics through TensorBoard because the package is not installed: "
            "Please run pip install tensorboard to enable."
        )

    # Initialize our training
    rng = jax.random.PRNGKey(training_args.seed)
    dropout_rngs = jax.random.split(rng, jax.local_device_count())
    gumbel_rngs = jax.random.split(rng, jax.local_device_count())

    num_epochs = int(training_args.num_train_epochs)
    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()

    num_train_steps = len(vectorized_datasets["train"]) // train_batch_size * num_epochs

    # Create learning rate schedule
    warmup_fn = optax.linear_schedule(
        init_value=0.0, end_value=training_args.learning_rate, transition_steps=training_args.warmup_steps
    )
    decay_fn = optax.linear_schedule(
        init_value=training_args.learning_rate,
        end_value=0,
        transition_steps=num_train_steps - training_args.warmup_steps,
    )
    linear_decay_lr_schedule_fn = optax.join_schedules(
        schedules=[warmup_fn, decay_fn], boundaries=[training_args.warmup_steps]
    )

    # We use Optax's "masking" functionality to not apply weight decay
    # to bias and LayerNorm scale parameters. decay_mask_fn returns a
    # mask boolean with the same structure as the parameters.
    # The mask is True for parameters that should be decayed.
    def decay_mask_fn(params):
        flat_params = traverse_util.flatten_dict(params)
        flat_mask = {
            path: (path[-1] != "bias" and path[-2:] not in [("layer_norm", "scale"), ("final_layer_norm", "scale")])
            for path in flat_params
        }
        return traverse_util.unflatten_dict(flat_mask)

    # create adam optimizer
    adamw = optax.adamw(
        learning_rate=linear_decay_lr_schedule_fn,
        b1=training_args.adam_beta1,
        b2=training_args.adam_beta2,
        eps=training_args.adam_epsilon,
        weight_decay=training_args.weight_decay,
        mask=decay_mask_fn,
    )

    # Setup train state and define training hyper-parameters
    state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw)
    num_negatives = model.config.num_negatives
    contrastive_logits_temperature = model.config.contrastive_logits_temperature
    num_codevectors = model.config.num_codevectors_per_group * model.config.num_codevector_groups
    diversity_loss_weight = model.config.diversity_loss_weight

    # Define gradient update step fn
    def train_step(state, batch, dropout_rng, gumbel_rng):
        dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
        gumbel_rng, new_gumbel_rng = jax.random.split(gumbel_rng)

        def loss_fn(params):
            negative_indices = batch.pop("sampled_negative_indices")

            gumbel_temperature = jnp.clip(
                model_args.max_gumbel_temperature * model_args.gumbel_temperature_decay**state.step,
                a_min=model_args.min_gumbel_temperature,
            )

            outputs = state.apply_fn(
                **batch,
                gumbel_temperature=gumbel_temperature,
                params=params,
                dropout_rng=dropout_rng,
                gumbel_rng=gumbel_rng,
                train=True,
            )

            contrastive_loss = compute_contrastive_loss(
                outputs.projected_quantized_states,
                outputs.projected_states,
                negative_indices,
                batch["mask_time_indices"],
                contrastive_logits_temperature,
                num_negatives,
            )

            diversity_loss = (num_codevectors - outputs.codevector_perplexity) / num_codevectors
            loss = contrastive_loss + diversity_loss_weight * diversity_loss

            return loss

        grad_fn = jax.value_and_grad(loss_fn)
        loss, grad = grad_fn(state.params)
        grad = jax.lax.pmean(grad, "batch")
        new_state = state.apply_gradients(grads=grad)

        metrics = jax.lax.pmean(
            {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}, axis_name="batch"
        )

        return new_state, metrics, new_dropout_rng, new_gumbel_rng

    # Create parallel version of the train step
    p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))

    # Define eval fn
    def eval_step(params, batch):
        negative_indices = batch.pop("sampled_negative_indices")

        outputs = model(**batch, params=params, train=False)

        contrastive_loss = compute_contrastive_loss(
            outputs.projected_quantized_states,
            outputs.projected_states,
            negative_indices,
            batch["mask_time_indices"],
            contrastive_logits_temperature,
            num_negatives,
        )

        diversity_loss = (num_codevectors - outputs.codevector_perplexity) / num_codevectors
        loss = contrastive_loss + diversity_loss_weight * diversity_loss

        # summarize metrics
        metrics = {"loss": loss.mean(), "codevector_perplexity": outputs.codevector_perplexity}
        metrics = jax.lax.pmean(metrics, axis_name="batch")

        return metrics

    p_eval_step = jax.pmap(eval_step, "batch", donate_argnums=(0,))

    # Replicate the train state on each device
    state = jax_utils.replicate(state)

    train_time = 0
    train_metrics = []
    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
    for epoch in epochs:
        # ======================== Training ================================
        train_start = time.time()

        # Create sampling rng
        rng, input_rng = jax.random.split(rng)

        # Generate an epoch by shuffling sampling indices from the train dataset
        num_train_samples = len(vectorized_datasets["train"])
        # Avoid using jax.numpy here in case of TPU training
        train_samples_idx = np.random.permutation(np.arange(num_train_samples))
        train_batch_idx = generate_batch_splits(train_samples_idx, train_batch_size)

        # Gather the indexes for creating the batch and do a training step
        for step, batch_idx in enumerate(tqdm(train_batch_idx, desc="Training...", position=1)):
            samples = [vectorized_datasets["train"][int(idx)] for idx in batch_idx]
            model_inputs = data_collator(samples)
            model_inputs = shard(model_inputs.data)

            # Model forward
            state, train_metric, dropout_rngs, gumbel_rngs = p_train_step(
                state, model_inputs, dropout_rngs, gumbel_rngs
            )
            train_metrics.append(train_metric)

            cur_step = epoch * (num_train_samples // train_batch_size) + step

            if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                # Save metrics
                train_metric = jax_utils.unreplicate(train_metric)
                train_time += time.time() - train_start
                if has_tensorboard and jax.process_index() == 0:
                    write_train_metric(summary_writer, train_metrics, train_time, cur_step)

                epochs.write(
                    f"Step... ({cur_step} | Loss: {train_metric['loss'].mean()}, Learning Rate:"
                    f" {train_metric['learning_rate'].mean()})"
                )

                train_metrics = []

        # ======================== Evaluating ==============================
        num_eval_samples = len(vectorized_datasets["validation"])
        # Avoid using jax.numpy here in case of TPU training
        eval_samples_idx = np.arange(num_eval_samples)
        eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)

        eval_metrics = []
        for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
            samples = [vectorized_datasets["validation"][int(idx)] for idx in batch_idx]
            model_inputs = data_collator(samples)

            # Model forward
            model_inputs = shard(model_inputs.data)
            metrics = p_eval_step(state.params, model_inputs)
            eval_metrics.append(metrics)

        # get eval metrics
        eval_metrics = get_metrics(eval_metrics)
        eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)

        # Update progress bar
        epochs.write(
            f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {eval_metrics['loss']}, Perplexity:"
            f" {eval_metrics['codevector_perplexity']})"
        )

        # Save metrics
        if has_tensorboard and jax.process_index() == 0:
            cur_step = epoch * (len(vectorized_datasets["train"]) // train_batch_size)
            write_eval_metric(summary_writer, eval_metrics, cur_step)

        # save checkpoint after each epoch and push checkpoint to the hub
        if jax.process_index() == 0:
            params = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state.params))


SUMMARY:
This is a TPU (Tensor Processing Unit) training script for a language modeling task. The script trains a pre-trained model on a dataset and performs evaluation at each epoch step.

Key points:

* The script uses `TPU` to accelerate the training process
* It loads a pre-trained model from the `hub`
* The training dataset is split into batches, and each batch is processed sequentially
* Evaluation is done after each epoch, using a separate batch of samples
* The script prints loss values at each step
* It writes loss and learning rate metrics to a TensorBoard summary writer

Note that this script uses JAX, Hugging Face, and other libraries to simplify the training process.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/jax-projects/model_parallel/run_clm_mp.py


key: chroma:document
string_value:
CODE:

def data_loader(rng: jax.random.PRNGKey, dataset: Dataset, batch_size: int, shuffle: bool = False):
    """
    Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
    Shuffle batches if `shuffle` is `True`.
    """
    steps_per_epoch = len(dataset) // batch_size

    if shuffle:
        batch_idx = jax.random.permutation(rng, len(dataset))
    else:
        batch_idx = jnp.arange(len(dataset))

    batch_idx = batch_idx[: steps_per_epoch * batch_size]  # Skip incomplete batch.
    batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))

    for idx in batch_idx:
        batch = dataset[idx]
        batch = {k: jnp.array(v) for k, v in batch.items()}


SUMMARY:
**Summary**

`data_loader` Function:

* Purpose: Shards and yields batches of size `batch_size` from truncated `dataset`.
* Primary purpose: Data augmentation for model training.

**Key Implementation Details**

* Random shuffling enabled using JAX's `jax.random.permutation` with optional `shuffle` flag.
* Batch indices generated using a buffer (`jnp.arange`) if shuffling is disabled.
* Batches created by reshaping batch indices and iterating over them, yielding each batch as a dictionary.

**Dependencies/Requirements**

* `Dataset`: The input dataset to be shredded.
* `jax.random.PRNGKey` (rng): Random number generator key for shuffling batches.
* `jnp`: Necessary JAX numerical library functions.

Overall, this function efficiently batches data while enabling easy adaptation of shuffling enabled or disabled depending on the requirement.


key: chroma:document
string_value:
CODE:

def write_train_metric(summary_writer, train_metrics, train_time, step):
    summary_writer.scalar("train_time", train_time, step)

    train_metrics = stack_forest(train_metrics)
    for key, vals in train_metrics.items():
        tag = f"train_{key}"
        for i, val in enumerate(vals):


SUMMARY:
**Function Summary: `write_train_metric`**

**Purpose:** Writes training metrics and time to a summary writer.

**Main Functionality:**

1. Writes train time scalar to the summary writer.
2. Applies a `stack_forest` transformation to the input `train_metrics`.
3. Iterates through transformed metrics, writing each metric value with increasing step values to the summary writer.

**Key Implementation Details:**

* Applies `stack_forest` transformation to `train_metrics`.
* Writes scalar values using `summary_writer.scalar()` method.

**Notable Dependencies/Requirements:**

* `summary_writer` object (not shown in the provided code snippet).
* `train_time` variable.
* `train_metrics` list of dictionaries (transformed by `stack_forest` function).
* External dependency: The `stack_forest` transformation function is assumed to be implemented elsewhere.


key: chroma:document
string_value:
CODE:

def write_eval_metric(summary_writer, eval_metrics, step):
    for metric_name, value in eval_metrics.items():


SUMMARY:
**Function Summary: `write_eval_metric`**

The `write_eval_metric` function takes three parameters:

*   `summary_writer`: Used to write scalar metrics
*   `eval_metrics`: A dictionary containing evaluation metric names and their corresponding values
*   `step`: The current iteration step

**Primary Purpose:** Writes eval metric data to a summary writer.

**Key Implementation Details:**

*   Uses a dictionary (`eval_metrics`) to store metric names and values.
*   Utilizes `summary_writer.scalar` method to write scalar metrics, appending the metric name and value at each step.

**Notable Dependencies/Requirements:**

*   None (uses built-in Python functionality)
*   No external dependencies

**Interfaces:**
- Implemented from a base class or interface not specified in this function.


key: chroma:document
string_value:
CODE:

def create_learning_rate_fn(
    train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
) -> Callable[[int], jnp.ndarray]:
    """Returns a linear warmup, linear_decay learning rate function."""
    steps_per_epoch = train_ds_size // train_batch_size
    num_train_steps = steps_per_epoch * num_train_epochs
    warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
    decay_fn = optax.linear_schedule(
        init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
    )
    schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])


SUMMARY:
**Summary:**

The `create_learning_rate_fn` function creates a learning rate schedule that combines linear warmup and decay components. It takes in parameters for training data size, batch size, number of epochs, warmup steps, and the initial learning rate.

*   Sets up a linear warmup schedule with `optax.linear_schedule` over `num_warmup_steps`.
*   Creates a linear decay schedule with `optax.linear_schedule` over the remaining training steps.
*   Joins the two schedules using `optax.join_schedules` and returns the resulting learning rate schedule.

**Key Implementation Details:**

*   Linear warmup schedule with `warmup_fn`: `init_value=0.0`, `end_value=learning_rate`.
*   Linear decay schedule with `decay_fn`: `init_value=learning_rate`, `end_value=0`.

**Notable Dependencies/Requirements:**

*   `optax` ( dependency for linear schedules and functions).
*   `jnp` (dependency for numerical computations).
*   `Callable[[int], jnp.ndarray]` (return type indicating a function that takes an integer as input and returns a JAX NumPy array).


key: chroma:document
string_value:
CODE:

def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if (
        os.path.exists(training_args.output_dir)
        and os.listdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({training_args.output_dir}) already exists and is not empty. "
            "Use --overwrite_output_dir to overcome."
        )

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    # Setup logging, we only want one process per machine to log things on the screen.
    logger.setLevel(logging.INFO if jax.process_index() == 0 else logging.ERROR)
    if jax.process_index() == 0:
        datasets.utils.logging.set_verbosity_warning()
        transformers.utils.logging.set_verbosity_info()
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

    # Set the verbosity to info of the Transformers logger (on main process only):
    logger.info(f"Training/evaluation parameters {training_args}")

    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
    # (the dataset will be downloaded automatically from the datasets Hub).
    #
    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
    # 'text' is found. You can easily tweak this behavior (see below).
    if data_args.dataset_name is not None:
        # Downloading and loading a dataset from the hub.
        dataset = load_dataset(
            data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir, keep_in_memory=False
        )

        if "validation" not in dataset.keys():
            dataset["validation"] = load_dataset(
                data_args.dataset_name,
                data_args.dataset_config_name,
                split=f"train[:{data_args.validation_split_percentage}%]",
                cache_dir=model_args.cache_dir,
            )
            dataset["train"] = load_dataset(
                data_args.dataset_name,
                data_args.dataset_config_name,
                split=f"train[{data_args.validation_split_percentage}%:]",
                cache_dir=model_args.cache_dir,
            )
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
            extension = data_args.train_file.split(".")[-1]
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
            extension = data_args.validation_file.split(".")[-1]
        if extension == "txt":
            extension = "text"
        dataset = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)
    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
    # https://huggingface.co/docs/datasets/loading_datasets.

    # Load pretrained config and tokenizer
    if model_args.config_name:
        config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
    elif model_args.model_name_or_path:
        config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
    else:
        config = CONFIG_MAPPING[model_args.model_type]()
        logger.warning("You are instantiating a new config instance from scratch.")

    if model_args.tokenizer_name:
        tokenizer = AutoTokenizer.from_pretrained(
            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
        )
    elif model_args.model_name_or_path:
        tokenizer = AutoTokenizer.from_pretrained(
            model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
        )
    else:
        raise ValueError(
            "You are instantiating a new tokenizer from scratch. This is not supported by this script. "
            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
        )

    if training_args.do_train:
        column_names = dataset["train"].column_names
    else:
        column_names = dataset["validation"].column_names
    text_column_name = "text" if "text" in column_names else column_names[0]

    # since this will be pickled to avoid _LazyModule error in Hasher force logger loading before tokenize_function
    tok_logger = transformers.utils.logging.get_logger("transformers.tokenization_utils_base")

    def tokenize_function(examples):
        with CaptureLogger(tok_logger) as cl:
            output = tokenizer(examples[text_column_name])
        # clm input could be much much longer than block_size
        if "Token indices sequence length is longer than the" in cl.out:
            tok_logger.warning(
                "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits"
                " before being passed to the model."
            )
        return output

    tokenized_datasets = dataset.map(
        tokenize_function,
        batched=True,
        num_proc=data_args.preprocessing_num_workers,
        remove_columns=column_names,
        load_from_cache_file=not data_args.overwrite_cache,
    )

    if data_args.block_size is None:
        block_size = tokenizer.model_max_length
        if block_size > config.max_position_embeddings:
            logger.warning(
                f"The tokenizer picked seems to have a very large `model_max_length` ({tokenizer.model_max_length}). "
                f"Using block_size={min(1024, config.max_position_embeddings)} instead. You can change that default value by passing --block_size xxx."
            )
            block_size = min(1024, config.max_position_embeddings)
    else:
        if data_args.block_size > tokenizer.model_max_length:
            logger.warning(
                f"The block_size passed ({data_args.block_size}) is larger than the maximum length for the model "
                f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
            )
        block_size = min(data_args.block_size, tokenizer.model_max_length)

    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
    def group_texts(examples):
        # Concatenate all texts.
        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
        total_length = len(concatenated_examples[list(examples.keys())[0]])
        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
        # customize this part to your needs.
        if total_length >= block_size:
            total_length = (total_length // block_size) * block_size
        # Split by chunks of max_len.
        result = {
            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
            for k, t in concatenated_examples.items()
        }
        result["labels"] = result["input_ids"].copy()
        return result

    # Note that with `batched=True`, this map processes 1,000 texts together, so group_texts throws away a remainder
    # for each of those groups of 1,000 texts. You can adjust that batch_size here but a higher value might be slower
    # to preprocess.
    #
    # To speed up this part, we use multiprocessing. See the documentation of the map method for more information:
    # https://huggingface.co/docs/datasets/process#map

    lm_datasets = tokenized_datasets.map(
        group_texts,
        batched=True,
        num_proc=data_args.preprocessing_num_workers,
        load_from_cache_file=not data_args.overwrite_cache,
    )

    if training_args.do_train:
        if "train" not in tokenized_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = lm_datasets["train"]
        if data_args.max_train_samples is not None:
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))

    if training_args.do_eval:
        if "validation" not in tokenized_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = lm_datasets["validation"]
        if data_args.max_eval_samples is not None:
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))

    # Enable tensorboard only on the master node
    has_tensorboard = is_tensorboard_available()
    if has_tensorboard and jax.process_index() == 0:
        try:
            from flax.metrics.tensorboard import SummaryWriter

            summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))
        except ImportError as ie:
            has_tensorboard = False
            logger.warning(
                f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
            )
    else:
        logger.warning(
            "Unable to display metrics through TensorBoard because the package is not installed: "
            "Please run pip install tensorboard to enable."
        )

    # Initialize our training
    rng = jax.random.PRNGKey(training_args.seed)
    rng, dropout_rng = jax.random.split(rng)

    # Store some constant
    num_epochs = int(training_args.num_train_epochs)
    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
    steps_per_epoch = len(train_dataset) // train_batch_size
    total_train_steps = steps_per_epoch * num_epochs

    # TODO: weights should be initialized in pjitted fun, this won't work for REALLY large models
    # TODO: when loading from pre-trained model we need to make sure the vocab is divisible by num_partitions
    # GPT2's vocab is odd, we need to resize it for fine-tuning
    model = FlaxAutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
    )

    # Create learning rate schedule
    linear_decay_lr_schedule_fn = create_learning_rate_fn(
        len(train_dataset),
        train_batch_size,
        training_args.num_train_epochs,
        training_args.warmup_steps,
        training_args.learning_rate,
    )

    optimizer = optax.adamw(
        learning_rate=linear_decay_lr_schedule_fn,
        b1=training_args.adam_beta1,
        b2=training_args.adam_beta2,
        eps=training_args.adam_epsilon,
        weight_decay=training_args.weight_decay,
    )

    def get_initial_state(params):
        state = optimizer.init(params)
        return tuple(state), params

    # Get PartitionSpec for model params
    param_spec = set_partitions(unfreeze(model.params))

    # Get the PyTree for opt_state, we don't actually initialize the opt_state yet.
    params_shapes = jax.tree_util.tree_map(lambda x: x.shape, model.params)
    state_shapes = jax.eval_shape(get_initial_state, params_shapes)

    # get PartitionSpec for opt_state, this is very specific to adamw
    # TODO: optax returns different state for different optimizers, how can we handle this generically ?
    # or maybe we don't since in our examples we just use adamw or adafactor
    def get_opt_spec(x):
        if isinstance(x, dict):
            return param_spec
        return None

    opt_state_spec, param_spec = jax.tree_util.tree_map(
        get_opt_spec, state_shapes, is_leaf=lambda x: isinstance(x, (dict, optax.EmptyState))
    )

    # pjit the get_initial_state function to shard params and init
    # optimizer state in sharded way
    p_get_initial_state = pjit(
        get_initial_state,
        in_axis_resources=None,
        out_axis_resources=(opt_state_spec, param_spec),
    )

    # hack: move the inital params to CPU to free up device memory
    # TODO: allow loading weights on CPU in pre-trained model
    model.params = jax.tree_util.tree_map(lambda x: np.asarray(x), model.params)

    # mesh defination
    mesh_devices = np.array(jax.devices()).reshape(1, jax.local_device_count())

    # actually initialize the opt_state
    with mesh(mesh_devices, ("dp", "mp")):
        opt_state, params = p_get_initial_state(freeze(model.params))

    # cross-entropy with z loss
    def loss_fn(logits, labels, z_loss=0):
        shift_logits = logits[..., :-1, :]
        shift_labels = labels[..., 1:]

        shift_labels = onehot(shift_labels, shift_logits.shape[-1])

        shift_logits = shift_logits - jax.lax.stop_gradient(shift_logits.max(axis=-1, keepdims=True))
        log_z = jnp.log(jnp.sum(jnp.exp(shift_logits), axis=-1, keepdims=True))
        log_softmax = shift_logits - log_z
        loss = -jnp.sum(shift_labels * log_softmax, axis=-1)

        loss += (1e-4 * jnp.square(log_z.squeeze(-1))) * z_loss

        return loss.mean()

    # Define gradient update step fn
    # TODO: try to use TrainState instead of passing params and opt_state individually
    def train_step(params, opt_state, dropout_rng, batch, step):
        dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)

        def compute_loss(params):
            labels = batch.pop("labels")
            logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
            loss = loss_fn(logits, labels, z_loss=1.0)
            return loss

        grad_fn = jax.value_and_grad(compute_loss)
        loss, grads = grad_fn(params)

        updates, new_opt_state = optimizer.update(grads, opt_state, params)
        new_params = optax.apply_updates(params, updates)

        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(step)}
        return new_params, tuple(new_opt_state), new_dropout_rng, metrics, step + 1

    # Define eval fn
    def eval_step(input_ids, labels, params):
        logits = model(input_ids=input_ids, params=params, train=False)[0]
        loss = loss_fn(logits, labels)
        # metrics
        return {"loss": loss}

    p_train_step = pjit(
        train_step,
        in_axis_resources=(param_spec, opt_state_spec, None, None, None),
        out_axis_resources=(param_spec, opt_state_spec, None, None, None),
        donate_argnums=(0, 1),
    )

    p_eval_step = pjit(
        eval_step,
        in_axis_resources=(None, None, param_spec),
        out_axis_resources=None,
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel & distributed) = {train_batch_size}")
    logger.info(f"  Total optimization steps = {total_train_steps}")

    train_time = 0
    train_metrics = []
    epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
    global_step = 0
    # we are not doing 2D parallelism (yet!), this just does model parallelism
    with mesh(mesh_devices, ("dp", "mp")):
        for _ in epochs:
            # ======================== Training ================================
            train_start = time.time()

            # Create sampling rng
            rng, input_rng = jax.random.split(rng)

            # Generate an epoch by shuffling sampling indices from the train dataset
            train_metrics = []
            train_loader = data_loader(input_rng, train_dataset, train_batch_size, shuffle=True)
            steps_per_epoch = len(train_dataset) // train_batch_size

            # train
            for _ in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
                batch = next(train_loader)
                params, opt_state, dropout_rng, train_metric, global_step = p_train_step(
                    params,
                    opt_state,
                    dropout_rng,
                    batch,
                    global_step,
                )
                train_metrics.append(train_metric)

                cur_step = global_step

                if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                    # Save metrics
                    train_time += time.time() - train_start
                    if has_tensorboard and jax.process_index() == 0:
                        write_train_metric(summary_writer, train_metrics, train_time, cur_step)

                    epochs.write(
                        f"Step... ({cur_step} | Loss: {train_metric['loss']}, Learning Rate:"
                        f" {train_metric['learning_rate']})"
                    )

                    train_metrics = []

                if cur_step % training_args.eval_steps == 0 and cur_step > 0:
                    # ======================== Evaluating ==============================
                    eval_metrics = []
                    eval_loader = data_loader(input_rng, eval_dataset, eval_batch_size)
                    eval_steps = len(eval_dataset) // eval_batch_size

                    for _ in tqdm(range(eval_steps), desc="Evaluating...", position=2, leave=False):
                        batch = next(eval_loader)
                        metrics = p_eval_step(batch["input_ids"], batch["labels"], params)
                        eval_metrics.append(metrics)

                    # normalize eval metrics
                    eval_metrics = stack_forest(eval_metrics)
                    eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)

                    try:
                        eval_metrics["perplexity"] = math.exp(eval_metrics["loss"])
                    except OverflowError:
                        eval_metrics["perplexity"] = float("inf")

                    logger.info(
                        f"Step... ({cur_step} | Eval loss: {eval_metrics['loss']} | Eval Perplexity:"
                        f" {eval_metrics['perplexity']}"
                    )

                if cur_step % training_args.save_steps == 0 and cur_step > 0:
                    # save checkpoint after each epoch and push checkpoint to the hub
                    if jax.process_index() == 0:
                        params = jax.device_get(params)
                        model.save_pretrained(
                            training_args.output_dir,
                            params=params,
                            push_to_hub=training_args.push_to_hub,
                            commit_message=f"Saving weights and logs of step {cur_step}",


SUMMARY:
A code for training a machine learning model using the JAX library. The code consists of several components:

1. **Model Definition**: A neural network architecture is defined, which outputs a prediction based on input data.
2. **Data Loaders**: Two data loaders are created to load training and evaluation datasets.
3. **Training Loop**: A loop iterates over epochs, batches, and steps. It:
	* Generates random mini-batches from the train dataset.
	* Updates model parameters using the Adam optimizer and loss function.
	* Evaluates the model on the eval dataset.
4. **Evaluation**: The model's performance is measured by various metrics (e.g., accuracy, perplexity).
5. **Logging and Saving**: Log metrics during training and save checkpoint after each epoch.

**Concise Summary**

Using JAX, this code trains a neural network model on a custom dataset with Adam optimizer and batch processing. It logs train metrics during epochs and saves checkpoints frequently. The training loop includes evaluation on the eval dataset, which provides an independent estimate of the model's performance.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/jax-projects/model_parallel/partitions.py


key: chroma:document
string_value:
CODE:

def _match(qs, ks):
    """Return True if regexes in qs match any window of strings in tuple ks."""
    # compile regexes and force complete match
    qts = tuple((re.compile(x + "$") for x in qs))
    for i in range(len(ks) - len(qs) + 1):
        matches = [x.match(y) for x, y in zip(qts, ks[i:])]
        if matches and all(matches):
            return True


SUMMARY:
**Summary:** `_match` function verifies if any complete regex pattern exists within a sliding window of `n` strings from the input tuple `ks`.

**Main Purpose and Functionality:**

This function takes two inputs, `qs` (regex patterns) and `ks` (strings), and returns `True` if any regex pattern in `qs` matches a whole string within a window of size `len(qs)` from `ks`.
 
The key implementation detail is the use of a sliding window approach to efficiently compare each regex with multiple subsequent strings in `ks`.

**Key Implementation Details:**

- It compiles all regex patterns at once and stores them in `qts`, ensuring complete matches are forced.
- The function employs a sliding window pattern, considering each string in `ks` as part of a partial match against the compiled regexes.

**Notable Dependencies or Requirements:**

- `re.compile(x + "$")`: Compiles regex patterns with the `$` anchor (end-of-string) for complete matches.
- `zip(qts, ks[i:])`: Zips up the compiled regexes and windows in `ks`.
- `all(matches)`: Verifies if all regex matches within a window are successful.


key: chroma:document
string_value:
CODE:

def _replacement_rules(rules):
    def replace(key, val):
        for rule, replacement in rules:
            if _match(rule, key):
                return replacement
        return val



SUMMARY:
**Code Summary**

`_replacement_rules` function creates a closure to apply replacement rules to input `key`. It takes an iterable of tuples (`rules`) as input, where each tuple contains a match key `_rule` and its corresponding replacement value. The function defines a private helper function `_match` (not shown in this snippet) to determine if the input `key` matches any rule.

**Key Implementation Details**

* Applies replacement rules using a closure
* Uses a single `_match` function to validate rule matches
* Returns the replacement value or original value for unmatched keys

**Function Signature**
```python
def _replacement_rules(rules):
    def replace(key, val):
        # ...
```
**Notable Dependencies**

None. This function is self-sufficient and does not depend on external libraries.

**Interfaces Implemented (Implicit)**
The `_replace` closure implicitly implements a simple mapping interface to replace keys with their corresponding replacement values.


key: chroma:document
string_value:
CODE:
# replicate the hidden dim and shard feed-forward and head dim
def _get_partition_rules():
    return [
        # embeddings
        (("transformer", "wpe", "embedding"), P("mp", None)),
        (("transformer", "wte", "embedding"), P("mp", None)),
        # atention
        (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(None, "mp")),
        (("attention", "out_proj", "kernel"), P("mp", None)),
        (("attention", "out_proj", "bias"), None),
        # mlp
        (("mlp", "c_fc", "kernel"), P(None, "mp")),
        (("mlp", "c_fc", "bias"), P("mp")),
        (("mlp", "c_proj", "kernel"), P("mp", None)),
        (("mlp", "c_proj", "bias"), None),
        # layer norms
        ((r"ln_\d+", "bias"), None),
        ((r"\d+", r"ln_\d+", "scale"), None),
        (("ln_f", "bias"), None),
        (("ln_f", "scale"), None),


SUMMARY:
**Function Summary: `_get_partition_rules`**

This function returns a list of partition rules configuration for `GPTNeo`, which is used to optimize memory usage while maintaining performance.

**Key Functionality:**

The function defines various partition patterns for different components of the model, including:

* Embeddings (transformer and wte)
* Attention
* MLP (feed-forward and projection modules)
* Layer norms

**Notable Implementation Details:**

* Each pattern consists of a tuple containing the prefix expression and optional replacement expressions.
* `P(expr1, expr2)` is a placeholder that allows for dynamic generation of partitions based on input patterns.

**External Dependencies and Interfaces:**

None notable dependencies are mentioned explicitly in this code chunk. However, the function likely relies on internal mechanisms not exposed here, such as `P` with its placeholder syntax.


key: chroma:document
string_value:
CODE:

def set_partitions(in_dict):
    rules = _get_partition_rules()
    replace = _replacement_rules(rules)
    initd = {k: _unmatched for k in flatten_dict(in_dict)}
    result = {k: replace(k, v) for k, v in initd.items()}
    assert _unmatched not in result.values(), "Incomplete partition spec."


SUMMARY:
**Function Summary**

`set_partitions(in_dict)` sets partitions based on provided input dictionary.

**Key Features**

- Retrieves and applies pre-defined partition rules (`_get_partition_rules()`)
- Initializes partition data structure by inverting the input dictionary
- Replaces unmatched values using `_replacement_rules()` function
- Asserts that all slots are assigned to avoid incomplete partition specifications

**Implemented Functions/Interfaces**

* `_set_partitions()`: parent function
* `_get_partition_rules()`
* `_replacement_rules()`
* `flatten_dict(in_dict)`
* `unflatten_dict(result)`

**External Dependencies**

None. This snippet does not reference any external libraries or frameworks.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/examples/research_projects/movement-pruning/masked_run_glue.py


key: chroma:document
string_value:
CODE:

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:


SUMMARY:
**Summary:**

Sets the seed for various Random Number Generators (RNGs) and PyTorch. Specifically:

* Sets seeds for NumPy (`np.random`), Python's built-in `random`, and PyTorch's manual tensor randomization (`torch.manual_seed`).
* If multiple GPUs are used (`args.n_gpu > 0`), also sets the seed for PyTorch's CUDA tensor randomization.

**Key Details:**

* Applies seed to:
	+ NumPy's random number generator (`np.random.seed`)
	+ Python's built-in `random` module (`random.seed`)
	+ PyTorch manual tensor randomization (`torch.manual_seed`)
* Conditional seeding based on multiple GPUs using PyTorch's CUDA tensor randomization (`torch.cuda.manual_seed_all`) if `args.n_gpu > 0`.


key: chroma:document
string_value:
CODE:

def schedule_threshold(
    step: int,
    total_step: int,
    warmup_steps: int,
    initial_threshold: float,
    final_threshold: float,
    initial_warmup: int,
    final_warmup: int,
    final_lambda: float,
):
    if step <= initial_warmup * warmup_steps:
        threshold = initial_threshold
    elif step > (total_step - final_warmup * warmup_steps):
        threshold = final_threshold
    else:
        spars_warmup_steps = initial_warmup * warmup_steps
        spars_schedu_steps = (final_warmup + initial_warmup) * warmup_steps
        mul_coeff = 1 - (step - spars_warmup_steps) / (total_step - spars_schedu_steps)
        threshold = final_threshold + (initial_threshold - final_threshold) * (mul_coeff**3)
    regu_lambda = final_lambda * threshold / final_threshold


SUMMARY:
**Function Summary: `schedule_threshold`**

**Purpose:** Calculates the learning rate (`threshold`) and regularization lambda based on the current step.

**Main Functionality:**

* Schedules learning rate (`threshold`) based on current step and total steps, with warm-up phases.
* Returns:
	+ `threshold`: scheduled learning rate
	+ `regu_lambda`: regulated regularization lambda

**Key Details:**

* Warm-up periods are divided into sparse warm-up phase (`spars_warmup_steps`) and schedule phase (`spars_schedu_steps`).
* Learning rate adjustment is done using a cubic interpolation formula.
* Regularization lambda is calculated based on the scheduled threshold.

**Notable Dependencies/Requirements:**

* None

**Method Signatures and Interfaces:**

* No external interfaces or dependencies are used. 

The function uses only built-in Python functions and operators for calculation purposes.


key: chroma:document
string_value:
CODE:

def regularization(model: nn.Module, mode: str):
    regu, counter = 0, 0
    for name, param in model.named_parameters():
        if "mask_scores" in name:
            if mode == "l1":
                regu += torch.norm(torch.sigmoid(param), p=1) / param.numel()
            elif mode == "l0":
                regu += torch.sigmoid(param - 2 / 3 * np.log(0.1 / 1.1)).sum() / param.numel()
            else:
                raise ValueError("Don't know this mode.")
            counter += 1


SUMMARY:
**Function Summary**

`regularization(model: nn.Module, mode: str)`: Computes regularization term for a given neural network model in one of three Lasso (L1), L0 (L1 with shrinkage), or unknown modes.

### Parameters:

* `model`: Neural Network Module
* `mode`: String indicating Lasso mode ("l1") or L0 mode ("l0")

### Returns:

Regularization term as a scalar value, normalized by the number of trainable parameters.

### Implementation Details:

* Calculates regularization term using Lasso/L0 penalty functions applied to masked score parameters.
* Handles two modes: "l1" (Lasso) and "l0" (Lasso with shrinkage).

### Notes:

* No explicit dependencies listed, assuming neural network module dependencies are handled within the model namespace (`nn.Module`).
* Inherited or implemented interfaces not explicitly mentioned.


key: chroma:document
string_value:
CODE:

def train(args, train_dataset, model, tokenizer, teacher=None):
    """Train the model"""
    if args.local_rank in [-1, 0]:
        tb_writer = SummaryWriter(log_dir=args.output_dir)

    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)

    if args.max_steps > 0:
        t_total = args.max_steps
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
    else:
        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs

    # Prepare optimizer and schedule (linear warmup and decay)
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if "mask_score" in n and p.requires_grad],
            "lr": args.mask_scores_learning_rate,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if "mask_score" not in n and p.requires_grad and not any(nd in n for nd in no_decay)
            ],
            "lr": args.learning_rate,
            "weight_decay": args.weight_decay,
        },
        {
            "params": [
                p
                for n, p in model.named_parameters()
                if "mask_score" not in n and p.requires_grad and any(nd in n for nd in no_decay)
            ],
            "lr": args.learning_rate,
            "weight_decay": 0.0,
        },
    ]

    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
    )

    # Check if saved optimizer or scheduler states exist
    if os.path.isfile(os.path.join(args.model_name_or_path, "optimizer.pt")) and os.path.isfile(
        os.path.join(args.model_name_or_path, "scheduler.pt")
    ):
        # Load in optimizer and scheduler states
        optimizer.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "optimizer.pt")))
        scheduler.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "scheduler.pt")))

    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = nn.DataParallel(model)

    # Distributed training (should be after apex fp16 initialization)
    if args.local_rank != -1:
        model = nn.parallel.DistributedDataParallel(
            model,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            find_unused_parameters=True,
        )

    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
    logger.info(
        "  Total train batch size (w. parallel, distributed & accumulation) = %d",
        args.train_batch_size
        * args.gradient_accumulation_steps
        * (torch.distributed.get_world_size() if args.local_rank != -1 else 1),
    )
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)
    # Distillation
    if teacher is not None:
        logger.info("  Training with distillation")

    global_step = 0
    # Global TopK
    if args.global_topk:
        threshold_mem = None
    epochs_trained = 0
    steps_trained_in_current_epoch = 0
    # Check if continuing training from a checkpoint
    if os.path.exists(args.model_name_or_path):
        # set global_step to global_step of last saved checkpoint from model path
        try:
            global_step = int(args.model_name_or_path.split("-")[-1].split("/")[0])
        except ValueError:
            global_step = 0
        epochs_trained = global_step // (len(train_dataloader) // args.gradient_accumulation_steps)
        steps_trained_in_current_epoch = global_step % (len(train_dataloader) // args.gradient_accumulation_steps)

        logger.info("  Continuing training from checkpoint, will skip to saved global_step")
        logger.info("  Continuing training from epoch %d", epochs_trained)
        logger.info("  Continuing training from global step %d", global_step)
        logger.info("  Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch)

    tr_loss, logging_loss = 0.0, 0.0
    model.zero_grad()
    train_iterator = trange(
        epochs_trained,
        int(args.num_train_epochs),
        desc="Epoch",
        disable=args.local_rank not in [-1, 0],
    )
    set_seed(args)  # Added here for reproducibility
    for _ in train_iterator:
        epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
        for step, batch in enumerate(epoch_iterator):
            # Skip past any already trained steps if resuming training
            if steps_trained_in_current_epoch > 0:
                steps_trained_in_current_epoch -= 1
                continue

            model.train()
            batch = tuple(t.to(args.device) for t in batch)
            threshold, regu_lambda = schedule_threshold(
                step=global_step,
                total_step=t_total,
                warmup_steps=args.warmup_steps,
                final_threshold=args.final_threshold,
                initial_threshold=args.initial_threshold,
                final_warmup=args.final_warmup,
                initial_warmup=args.initial_warmup,
                final_lambda=args.final_lambda,
            )
            # Global TopK
            if args.global_topk:
                if threshold == 1.0:
                    threshold = -1e2  # Or an indefinitely low quantity
                else:
                    if (threshold_mem is None) or (global_step % args.global_topk_frequency_compute == 0):
                        # Sort all the values to get the global topK
                        concat = torch.cat(
                            [param.view(-1) for name, param in model.named_parameters() if "mask_scores" in name]
                        )
                        n = concat.numel()
                        kth = max(n - (int(n * threshold) + 1), 1)
                        threshold_mem = concat.kthvalue(kth).values.item()
                        threshold = threshold_mem
                    else:
                        threshold = threshold_mem
            inputs = {"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
            if args.model_type != "distilbert":
                inputs["token_type_ids"] = (
                    batch[2] if args.model_type in ["bert", "masked_bert", "xlnet", "albert"] else None
                )  # XLM, DistilBERT, RoBERTa, and XLM-RoBERTa don't use segment_ids

            if "masked" in args.model_type:
                inputs["threshold"] = threshold

            outputs = model(**inputs)
            loss, logits_stu = outputs  # model outputs are always tuple in transformers (see doc)

            # Distillation loss
            if teacher is not None:
                if "token_type_ids" not in inputs:
                    inputs["token_type_ids"] = None if args.teacher_type == "xlm" else batch[2]
                with torch.no_grad():
                    (logits_tea,) = teacher(
                        input_ids=inputs["input_ids"],
                        token_type_ids=inputs["token_type_ids"],
                        attention_mask=inputs["attention_mask"],
                    )

                loss_logits = nn.functional.kl_div(
                    input=nn.functional.log_softmax(logits_stu / args.temperature, dim=-1),
                    target=nn.functional.softmax(logits_tea / args.temperature, dim=-1),
                    reduction="batchmean",
                ) * (args.temperature**2)

                loss = args.alpha_distil * loss_logits + args.alpha_ce * loss

            # Regularization
            if args.regularization is not None:
                regu_ = regularization(model=model, mode=args.regularization)
                loss = loss + regu_lambda * regu_

            if args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

            tr_loss += loss.item()
            if (step + 1) % args.gradient_accumulation_steps == 0 or (
                # last step in epoch but step is always smaller than gradient_accumulation_steps
                len(epoch_iterator) <= args.gradient_accumulation_steps and (step + 1) == len(epoch_iterator)
            ):
                if args.fp16:
                    nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                else:
                    nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    tb_writer.add_scalar("threshold", threshold, global_step)
                    for name, param in model.named_parameters():
                        if not param.requires_grad:
                            continue
                        tb_writer.add_scalar("parameter_mean/" + name, param.data.mean(), global_step)
                        tb_writer.add_scalar("parameter_std/" + name, param.data.std(), global_step)
                        tb_writer.add_scalar("parameter_min/" + name, param.data.min(), global_step)
                        tb_writer.add_scalar("parameter_max/" + name, param.data.max(), global_step)
                        tb_writer.add_scalar("grad_mean/" + name, param.grad.data.mean(), global_step)
                        tb_writer.add_scalar("grad_std/" + name, param.grad.data.std(), global_step)
                        if args.regularization is not None and "mask_scores" in name:
                            if args.regularization == "l1":
                                perc = (torch.sigmoid(param) > threshold).sum().item() / param.numel()
                            elif args.regularization == "l0":
                                perc = (torch.sigmoid(param - 2 / 3 * np.log(0.1 / 1.1))).sum().item() / param.numel()
                            tb_writer.add_scalar("retained_weights_perc/" + name, perc, global_step)

                optimizer.step()
                scheduler.step()  # Update learning rate schedule
                model.zero_grad()
                global_step += 1

                if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    logs = {}
                    if (
                        args.local_rank == -1 and args.evaluate_during_training
                    ):  # Only evaluate when single GPU otherwise metrics may not average well
                        results = evaluate(args, model, tokenizer)
                        for key, value in results.items():
                            eval_key = "eval_{}".format(key)
                            logs[eval_key] = value

                    loss_scalar = (tr_loss - logging_loss) / args.logging_steps
                    learning_rate_scalar = scheduler.get_lr()
                    logs["learning_rate"] = learning_rate_scalar[0]
                    if len(learning_rate_scalar) > 1:
                        for idx, lr in enumerate(learning_rate_scalar[1:]):
                            logs[f"learning_rate/{idx+1}"] = lr
                    logs["loss"] = loss_scalar
                    if teacher is not None:
                        logs["loss/distil"] = loss_logits.item()
                    if args.regularization is not None:
                        logs["loss/regularization"] = regu_.item()
                    if (teacher is not None) or (args.regularization is not None):
                        if (teacher is not None) and (args.regularization is not None):
                            logs["loss/instant_ce"] = (
                                loss.item()
                                - regu_lambda * logs["loss/regularization"]
                                - args.alpha_distil * logs["loss/distil"]
                            ) / args.alpha_ce
                        elif teacher is not None:
                            logs["loss/instant_ce"] = (
                                loss.item() - args.alpha_distil * logs["loss/distil"]
                            ) / args.alpha_ce
                        else:
                            logs["loss/instant_ce"] = loss.item() - regu_lambda * logs["loss/regularization"]
                    logging_loss = tr_loss

                    for key, value in logs.items():
                        tb_writer.add_scalar(key, value, global_step)
                    print(json.dumps({**logs, **{"step": global_step}}))

                if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
                    # Save model checkpoint
                    output_dir = os.path.join(args.output_dir, "checkpoint-{}".format(global_step))
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    model_to_save = (
                        model.module if hasattr(model, "module") else model
                    )  # Take care of distributed/parallel training
                    model_to_save.save_pretrained(output_dir)
                    tokenizer.save_pretrained(output_dir)

                    torch.save(args, os.path.join(output_dir, "training_args.bin"))
                    logger.info("Saving model checkpoint to %s", output_dir)

                    torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
                    torch.save(scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
                    logger.info("Saving optimizer and scheduler states to %s", output_dir)

            if args.max_steps > 0 and global_step > args.max_steps:
                epoch_iterator.close()
                break
        if args.max_steps > 0 and global_step > args.max_steps:
            train_iterator.close()
            break

    if args.local_rank in [-1, 0]:
        tb_writer.close()



SUMMARY:
Here is a concise summary:

This code snippet appears to be the main training loop for a neural network model using Hugging Face's Transformer library. The code takes in several hyperparameters and training arguments from ` entrenamiento.py` (presumably a Python script for configuring training settings). 

The training process involves the following steps:
1. Model checkpointing: saving the trained model, tokenizer, training arguments, optimizer state, and scheduler state to disk at regular intervals.
2. Logging metrics to TensorBoard writer (`tb_writer`) including loss, learning rate, instant CE, and regularization loss.
3. Evaluating the model during training (if specified) using `evaluate` function from Hugging Face Transformers library.
4. Stopping the training process when maximum steps are reached or when the loss stops improving.

Note: This summary does not include all the details and nuances of the code but highlights its main functionality.


key: chroma:document
string_value:
CODE:

def evaluate(args, model, tokenizer, prefix=""):
    # Loop to handle MNLI double evaluation (matched, mis-matched)
    eval_task_names = ("mnli", "mnli-mm") if args.task_name == "mnli" else (args.task_name,)
    eval_outputs_dirs = (args.output_dir, args.output_dir + "/MM") if args.task_name == "mnli" else (args.output_dir,)

    results = {}
    for eval_task, eval_output_dir in zip(eval_task_names, eval_outputs_dirs):
        eval_dataset = load_and_cache_examples(args, eval_task, tokenizer, evaluate=True)

        if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
            os.makedirs(eval_output_dir)

        args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
        # Note that DistributedSampler samples randomly
        eval_sampler = SequentialSampler(eval_dataset)
        eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)

        # multi-gpu eval
        if args.n_gpu > 1 and not isinstance(model, nn.DataParallel):
            model = nn.DataParallel(model)

        # Eval!
        logger.info("***** Running evaluation {} *****".format(prefix))
        logger.info("  Num examples = %d", len(eval_dataset))
        logger.info("  Batch size = %d", args.eval_batch_size)
        eval_loss = 0.0
        nb_eval_steps = 0
        preds = None
        out_label_ids = None

        # Global TopK
        if args.global_topk:
            threshold_mem = None

        for batch in tqdm(eval_dataloader, desc="Evaluating"):
            model.eval()
            batch = tuple(t.to(args.device) for t in batch)

            with torch.no_grad():
                inputs = {"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
                if args.model_type != "distilbert":
                    inputs["token_type_ids"] = (
                        batch[2] if args.model_type in ["bert", "masked_bert", "xlnet", "albert"] else None
                    )  # XLM, DistilBERT, RoBERTa, and XLM-RoBERTa don't use segment_ids
                if "masked" in args.model_type:
                    inputs["threshold"] = args.final_threshold
                    if args.global_topk:
                        if threshold_mem is None:
                            concat = torch.cat(
                                [param.view(-1) for name, param in model.named_parameters() if "mask_scores" in name]
                            )
                            n = concat.numel()
                            kth = max(n - (int(n * args.final_threshold) + 1), 1)
                            threshold_mem = concat.kthvalue(kth).values.item()
                        inputs["threshold"] = threshold_mem
                outputs = model(**inputs)
                tmp_eval_loss, logits = outputs[:2]

                eval_loss += tmp_eval_loss.mean().item()
            nb_eval_steps += 1
            if preds is None:
                preds = logits.detach().cpu().numpy()
                out_label_ids = inputs["labels"].detach().cpu().numpy()
            else:
                preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
                out_label_ids = np.append(out_label_ids, inputs["labels"].detach().cpu().numpy(), axis=0)

        eval_loss = eval_loss / nb_eval_steps
        if args.output_mode == "classification":
            from scipy.special import softmax

            probs = softmax(preds, axis=-1)
            entropy = np.exp((-probs * np.log(probs)).sum(axis=-1).mean())
            preds = np.argmax(preds, axis=1)
        elif args.output_mode == "regression":
            preds = np.squeeze(preds)
        result = compute_metrics(eval_task, preds, out_label_ids)
        results.update(result)
        if entropy is not None:
            result["eval_avg_entropy"] = entropy

        output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt")
        with open(output_eval_file, "w") as writer:
            logger.info("***** Eval results {} *****".format(prefix))
            for key in sorted(result.keys()):
                logger.info("  %s = %s", key, str(result[key]))
                writer.write("%s = %s\n" % (key, str(result[key])))



SUMMARY:
**Summary:**

This function, `evaluate`, runs model evaluation using the provided `args`, `model`, and `tokenizer`. Its primary purpose is to calculate metrics for a given evaluation task.

**Key Functions and Variables:**

* `load_and_cache_examples`: loads examples for the evaluation task
* `DataLoader` : creates a data loader for the eval dataset
* `SequentialSampler`: uses a sequential sampler instead of distributed sampling
* `model.eval()` : sets the model to evaluation mode

**Algorithms/Design Choices:**

* Global TopK method is applied if `args.global_topk` is True
* `Threshold` input is processed based on the model type and `global_topk` flag

**Notable Dependencies/Requirements:**

* `load_and_cache_examples`: requires `args`, `tokenizer`

**Interfaces/Classes:**

* `DataLoader`
* `SequentialSampler`
* `nn.DataParallel`
* `torch.nn.Module` (model)

**External Dependencies:**

* `scipy.special`: used for softmax function
* `logger`: a logging module (not specified in the provided code)


