from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
import gradio as gr

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 = abstract.split(". ")  # abstract to 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."


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# SCORE Claim Identification")
    gr.Markdown("This is a demo to find scientific claims made by a given abstract.")
    abstract = gr.Textbox(label="Abstract")
    greet_btn = gr.Button("Find Claims")
    output = gr.Textbox(label="Detected Claims")
    greet_btn.click(fn=inference, inputs=abstract, outputs=output, api_name="inference")
    examples = gr.Examples(
        examples=[
            "This study adopted a person (actor) by partner perspective to examine how actor personality traits, partner personality traits, and specific actor by partner personality trait interactions predict actor's depressive symptoms across the first 2years of the transition to parenthood. Data were collected from a large sample of new parents (both partners in each couple) 6weeks before the birth of their first child, and then at 6, 12, 18, and 24months postpartum. The results revealed that higher actor neuroticism and lower partner agreeableness predicted higher levels of depressive symptoms in actors. Moreover, the specific combination of high actor neuroticism and low partner agreeableness was a particularly problematic combination, which was intensified when prepartum dysfunctional problem-solving communication and aggression existed in the relationship. These results demonstrate the importance of considering certain actor by partner disposition pairings to better understand actors' emotional well-being during major life transitions."
        ],
        inputs=[abstract]
    )

demo.launch()