train_test_model.py

import argparse
import glob
import logging
import os
import os.path
import re
from pathlib import Path
from typing import Tuple

import evaluate
import numpy as np
import torch
from datasets import Dataset
from grobid_client import Client
from grobid_client.api.pdf import process_fulltext_document
from grobid_client.models import Article, ProcessForm, TextWithRefs
from grobid_client.types import TEI, File
from sklearn.model_selection import StratifiedKFold
from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer, pipeline
from transformers import AutoTokenizer
from transformers import DataCollatorWithPadding

logger = logging.getLogger(__name__)


def compute_metrics(eval_pred):
    accuracy = evaluate.load("accuracy")
    f1 = evaluate.load("f1")
    recall = evaluate.load("recall")
    precision = evaluate.load("precision")
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    acc = accuracy.compute(predictions=predictions, references=labels)
    f1_score = f1.compute(predictions=predictions, references=labels, average='binary')
    rec = recall.compute(predictions=predictions, references=labels, average='binary')
    prec = precision.compute(predictions=predictions, references=labels, average='binary')
    return {**f1_score, **prec, **rec, **acc}


def get_sentences_from_tei_section(section):
    sentences = []
    for paragraph in section.paragraphs:
        if isinstance(paragraph, TextWithRefs):
            paragraph = [paragraph]
        for sentence in paragraph:
            try:
                if not sentence.text.isdigit() and not (
                        len(section.paragraphs) == 3 and
                        section.paragraphs[0][0].text in ['\n', ' '] and
                        section.paragraphs[-1][0].text in ['\n', ' ']
                ):
                    sentences.append(re.sub('<[^<]+>', '', sentence.text))
            except Exception as e:
                logger.error(f"Error parsing sentence {str(e)}")
    sentences = [sentence if sentence.endswith(".") else f"{sentence}." for sentence in sentences]
    return sentences


def get_documents(input_docs_dir: str) -> Tuple[str, str, str]:
    client = None
    for file_path in glob.glob(input_docs_dir):
        file_obj = Path(file_path)
        if file_path.endswith(".tei") or file_path.endswith(".pdf"):
            with file_obj.open("rb") as fin:
                if file_path.endswith(".pdf"):
                    if client is None:
                        client = Client(base_url=os.environ.get("GROBID_API_URL"), timeout=1000, verify_ssl=False)
                    logger.info("Started pdf to TEI conversion")
                    form = ProcessForm(
                        segment_sentences="1",
                        input_=File(file_name=file_obj.name, payload=fin, mime_type="application/pdf"))
                    r = process_fulltext_document.sync_detailed(client=client, multipart_data=form)
                    file_stream = r.content
                else:
                    file_stream = fin
                try:
                    article: Article = TEI.parse(file_stream, figures=True)
                except Exception as e:
                    logging.error(f"Error parsing TEI file for {str(file_path)}: {str(e)}")
                    continue
                sentences = []
                for section in article.sections:
                    sentences.extend(get_sentences_from_tei_section(section))
                abstract = ""
                for section in article.sections:
                    if section.name == "ABSTRACT":
                        abstract = " ".join(get_sentences_from_tei_section(section))
                        break
                yield " ".join(sentences), article.title, abstract


def main():
    # Argument parsing
    parser = argparse.ArgumentParser(description="Train a sequence classification model.")
    parser.add_argument("-t", "--training-dir", required=True, type=str, help="Path to the csv dataset.")
    parser.add_argument("-o", "--output-dir", required=True, type=str,
                        help="Directory to save the model and tokenizer.")
    parser.add_argument("-m", "--model", required=True, type=str,
                        help="Huggingface model to finetune for classification")
    parser.add_argument("-e", "--epochs", default=4, type=int, help="Number of training epochs.")
    parser.add_argument("-s", "--summarize", default=False, action="store_true",
                        help="Summarize documents before classification")
    args = parser.parse_args()

    model_name = args.model

    positive_docs = [fulltext for fulltext, title, abstract in get_documents(os.path.join(args.training_dir,
                                                                                          "positive", "*"))
                     if fulltext]
    negative_docs = [fulltext for fulltext, title, abstract in get_documents(os.path.join(args.training_dir,
                                                                                          "negative", "*"))
                     if fulltext]

    if args.summarize:
        summarizer = pipeline("summarization", model="allenai/led-base-16384")
        positive_docs = [
            summarizer(fulltext[0:16384], max_length=512, min_length=30, do_sample=False)[0]["summary_text"]
            for fulltext in positive_docs]
        negative_docs = [
            summarizer(fulltext[0:16384], max_length=512, min_length=30, do_sample=False)[0]["summary_text"]
            for fulltext in negative_docs]

    positive_labels = [1] * len(positive_docs)
    negative_labels = [0] * len(negative_docs)

    all_fulltexts = positive_docs + negative_docs
    all_labels = positive_labels + negative_labels

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    encodings = tokenizer(all_fulltexts, truncation=True)
    encodings["labels"] = torch.tensor(all_labels)
    tokenized_dataset = Dataset.from_dict(encodings)

    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    training_args = TrainingArguments(
        output_dir=args.output_dir,
        learning_rate=2e-5,
        per_device_train_batch_size=20,
        per_device_eval_batch_size=20,
        num_train_epochs=args.epochs,
        weight_decay=0.1,
        fp16=True
    )

    precision_scores = []
    recall_scores = []
    f1_scores = []

    # Stratified K-Fold Cross-Validation
    skf = StratifiedKFold(n_splits=5)

    for train_index, test_index in skf.split(tokenized_dataset["input_ids"], tokenized_dataset["labels"]):
        train_dataset = tokenized_dataset.select(train_index)
        eval_dataset = tokenized_dataset.select(test_index)
        logger.info("Started fold for cross validation")
        model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2,
                                                                   ignore_mismatched_sizes=True)
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            tokenizer=tokenizer,
            data_collator=data_collator,
            compute_metrics=compute_metrics,
        )
        trainer.train()
        metrics = trainer.evaluate()
        print(f"Computed single k-fold validation with Precision: {metrics['eval_precision']}, Recall: "
              f"{metrics['eval_recall']}, F1: {metrics['eval_f1']}")
        precision_scores.append(metrics["eval_precision"])
        recall_scores.append(metrics["eval_recall"])
        f1_scores.append(metrics["eval_f1"])

    # Calculate average metrics
    avg_precision = np.mean(precision_scores)
    avg_recall = np.mean(recall_scores)
    avg_f1 = np.mean(f1_scores)
    std_precision = np.std(precision_scores)
    std_recall = np.std(recall_scores)
    std_f1 = np.std(f1_scores)

    print(f"Precision scores: {precision_scores}")
    print(f"Recall scores: {recall_scores}")
    print(f"F1 Scores: {f1_scores}")
    print(f"Average Precision: {avg_precision:5.4f} ± {std_precision:5.4f}")
    print(f"Average Recall: {avg_recall:5.4f} ± {std_recall:5.4f}")
    print(f"Average F1 Score: {avg_f1:5.4f} ± {std_f1:5.4f}")

    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2,
                                                               ignore_mismatched_sizes=True)
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
        tokenizer=tokenizer,
        data_collator=data_collator
    )
    trainer.train()
    trainer.save_model()


if __name__ == "__main__":
    main()