biodatlab/score-claim-identification

import spacy
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification

nlp = spacy.load("en_core_web_lg")
model_name = "biodatlab/score-claim-identification"
tokenizer_name = "allenai/scibert_scivocab_uncased"

tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

def inference(abstract: str):
    """
    Split an abstract into sentences and perform claim identification.
    """
    if abstract.strip() == "":
        return "Please provide an abstract as an input."
    claims = []
    sents = [sent.text for sent in nlp(abstract).sents]  # a list of sentences
    inputs = tokenizer(
        sents,
        return_tensors="pt",
        truncation=True,
        padding="longest"
    )
    logits = model(**inputs).logits
    preds = logits.argmax(dim=1)  # convert logits to predictions
    claims = [sent for sent, pred in zip(sents, preds) if pred == 1]
    if len(claims) > 0:
        return ".\n".join(claims)
    else:
        return "No claims found from a given abstract."

claims = inference(abstract)  # string of claim joining with \n