def predict(self, text: Union[List[str], str], as_argilla_records: bool = True, **kwargs):
        """
        The function takes in a list of strings and returns a list of predictions

        Args:
          text (Union[List[str], str]): The text to be classified.
          as_argilla_records (bool): If True, the prediction will be returned as an Argilla record. If
        False, the prediction will be returned as a string. Defaults to True

        Returns:
          A list of predictions
        """
        import torch

        if self.trainer_model is None:
            self._logger.warning("Using model without fine-tuning.")
            self.init_model(new=False)

        str_input = False
        if isinstance(text, str):
            text = [text]
            str_input = True

        if self._record_class == TextClassificationRecord:
            inputs = self.trainer_tokenizer(text, truncation=True, padding="longest", return_tensors="pt")

            with torch.no_grad():
                logits = self.trainer_model(**inputs).logits

            if self._multi_label:
                probabilities = torch.sigmoid(logits)
            else:
                probabilities = torch.softmax(logits, dim=1)
            predictions = []
            for probability in probabilities:
                prediction = []
                for idx, score in enumerate(probability):
                    prediction.append({"label": self._id2label[idx], "score": score})
                predictions.append(prediction)
        else:
            # Tokenize the text
            inputs_with_offsets = self.trainer_tokenizer(
                text, truncation=True, padding="longest", return_offsets_mapping=True, return_tensors="pt"
            )
            inputs_with_offsets = {k: v.to(self.device) for k, v in inputs_with_offsets.items()}
            offsets = inputs_with_offsets["offset_mapping"]

            # Perform the forward pass through the model
            with torch.no_grad():
                outputs = self.trainer_model(**{k: v for k, v in inputs_with_offsets.items() if k != "offset_mapping"})
                probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1).tolist()
                predictions = outputs.logits.argmax(dim=-1).tolist()

            # Iterate over the predictions and the offsets
            batch_formatted_predictions = []
            for batch_idx, example in enumerate(text):
                formatted_predictions = []
                idx = 0
                while idx < len(predictions[batch_idx]):
                    pred = predictions[batch_idx][idx]
                    label = self.trainer_model.config.id2label[pred]
                    if label != "O":
                        # Remove the B- or I-
                        label = label[2:]
                        start, end = offsets[batch_idx][idx]

                        # Grab all the tokens labeled with I-label
                        all_scores = []
                        while (
                            idx < len(predictions[batch_idx])
                            and self.trainer_model.config.id2label[predictions[batch_idx][idx]] == f"I-{label}"
                        ):
                            all_scores.append(probabilities[batch_idx][idx][pred])
                            _, end = offsets[batch_idx][idx]
                            idx += 1

                        # If the entity is only one token, we don't need to do anything
                        if start != end:
                            # The score is the mean of all the scores of the tokens in that grouped entity
                            score = None if np.isnan(np.mean(all_scores).item()) else np.mean(all_scores).item()
                            word = example[start:end]
                            formatted_predictions.append(
                                {
                                    "entity_group": label,
                                    "score": score,
                                    "word": word,
                                    "start": int(start),
                                    "end": int(end),
                                }
                            )
                    idx += 1
                batch_formatted_predictions.append(formatted_predictions)
            predictions = batch_formatted_predictions

        if as_argilla_records:
            formatted_prediction = []

            for val, pred in zip(text, predictions):
                if self._record_class == TextClassificationRecord:
                    formatted_prediction.append(
                        self._record_class(
                            text=val,
                            prediction=[(entry["label"], entry["score"]) for entry in pred],
                            multi_label=self._multi_label,
                        )
                    )
                elif self._record_class == TokenClassificationRecord:
                    formatted_prediction.append(
                        self._record_class(
                            text=val,
                            tokens=val.split(),
                            prediction=[
                                (
                                    char_spans["entity_group"],
                                    char_spans["start"],
                                    char_spans["end"],
                                    char_spans["score"],
                                )
                                for char_spans in pred
                            ],
                        )
                    )

                else:
                    raise NotImplementedError("`rg.Text2TextRecord` is not supported yet.")
        else:
            formatted_prediction = predictions

        if str_input:
            formatted_prediction = formatted_prediction[0]

        return formatted_prediction

        The function takes in a list of strings and returns a list of predictions

        Args:
          text (Union[List[str], str]): The text to be classified.
          as_argilla_records (bool): If True, the prediction will be returned as an Argilla record. If
        False, the prediction will be returned as a string. Defaults to True

        Returns:
          A list of predictions
        
argilla-v1/src/argilla_v1/training/peft.py