I am trying to extract 'agreement date' label from a corpus of legal contracts. In the train dataset, I used pytorch-transformer model to train.

model = AutoModelForTokenClassification.from_pretrained(model_checkpoint, num_labels=len(label_list))

Here label_list is the IOB format which gives ['B-Date', 'I-Date', 'O'] and model_checkpoint is "distilbert-base-uncased" I train the dataset after defining TrainingArguements, datacollator and matrics computaitons from predictions

model_output_dir = 'C:/Python/model_output_dir'

args = TrainingArguments(
    output_dir = model_output_dir,
    evaluation_strategy = "epoch",
    run_name = model_checkpoint,
    load_best_model_at_end = True,

from transformers import DataCollatorForTokenClassification
data_collator = DataCollatorForTokenClassification(tokenizer)
from seqeval.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report

def compute_metrics(p):
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)

    # Remove ignored index (special tokens)
    true_predictions = [
        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)]
    true_labels = [
        [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)]

    # Define the metric parameters
    overall_precision = precision_score(true_labels, true_predictions, zero_division=1)
    overall_recall = recall_score(true_labels, true_predictions, zero_division=1)
    overall_f1 = f1_score(true_labels, true_predictions, zero_division=1)
    overall_accuracy = accuracy_score(true_labels, true_predictions)
    # Return a dictionary with the calculated metrics
    return {
        "precision": overall_precision,
        "recall": overall_recall,
        "f1": overall_f1,
        "accuracy": overall_accuracy,}

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

My input training data after preprocessing look like this enter image description here

My test data has agreement text only in tokenized form. When I predict I get the big tensor like below enter image description here

How I get the required date label from this output ?

  • $\begingroup$ You are trying to extract a custom label "Agreement Date", right? If yes, shouldn't you train your model to recognize "Agreement Dates" with a custom dataset in the first place? $\endgroup$ Oct 11, 2021 at 8:01
  • $\begingroup$ Yes, of course. The training is done with a dataset having 'agreement date' label for each of the given contracts $\endgroup$
    – Jay
    Oct 12, 2021 at 9:06
  • $\begingroup$ Sorry, I'm not shur if I've well understood, but here is an answer anyway. The output is actually a set of probabilities. It is a raw result and you should take the highest score with argmax that would define the most probable solution. For instance, [0.2 , 0.6 , 0.3] would have the 2nd result as solution ('I-Date'). $\endgroup$ Oct 13, 2021 at 15:36


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.