Okay so I managed to resolve the problem.
Maybe it helps someone so this is the script I wrote:
Note: utils_ner is in master/examples/token-classification. align_predictions is from run_ner in the same dir. The script basically follows the predictions in run_ner. Also, (at least here) the input data must be in CoNLL-2003 data format as well!
from transformers import AutoTokenizer, AutoModelForTokenClassification, Trainer, TrainingArguments
from utils_ner import NerDataset
import torch
import numpy as np
from torch import nn
from typing import Tuple, List

pretrained_dir = "./biobert_v1.1_pubmed"
cache_dir = "./BioBERT_ner/data"
config_dir = "./biobert_v1.1_pubmed/config.json"
model_dir = "./BioBERT_ner/output_3/pytorch_model.bin"
data_dir = "./BioBERT_ner/pred_data"
labels_dir = "./BioBERT_ner/data/labels.txt"
output_dir = "./BioBERT_ner"

label_map = {
    0: "O",
    1: "B"
    }

def align_predictions(predictions: np.ndarray, label_ids: np.ndarray) -> Tuple[List[int], List[int]]:
    preds = np.argmax(predictions, axis=2)

    batch_size, seq_len = preds.shape

    out_label_list = [[] for _ in range(batch_size)]
    preds_list = [[] for _ in range(batch_size)]

    for i in range(batch_size):
        for j in range(seq_len):
            if label_ids[i, j] != nn.CrossEntropyLoss().ignore_index:
                out_label_list[i].append(label_map[label_ids[i][j]])
                preds_list[i].append(label_map[preds[i][j]])

    return preds_list, out_label_list

tokenizer = AutoTokenizer.from_pretrained(pretrained_dir, cache_dir=cache_dir)
model = AutoModelForTokenClassification.from_pretrained(pretrained_dir, config=config_dir,cache_dir=cache_dir)
model.load_state_dict(torch.load(model_dir))
dataset = NerDataset(
    data_dir=data_dir,
    tokenizer=tokenizer,
    labels=["B", "O"],
    model_type="BertForTokenClassification",
    max_seq_length=256,
    mode="pred"
)

args = TrainingArguments(
	output_dir=output_dir
)

trainer = Trainer(
    model=model,
    args=args
)

predictions, label_ids, metrics = trainer.predict(dataset)
preds_list, _ = align_predictions(predictions, label_ids)
print(preds_list)```