app.py

import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import numpy as np
import torch


def main():

    st.set_page_config(
        layout="wide",
        initial_sidebar_state="auto",
        page_title="Title Generator!",
        page_icon=None,
    )

    st.title("Title Generator: Generate a title from the abstract of a paper")
    st.text("")
    st.text("")

    example_prompts = [
        """Neural Painters is a class of models that follows a GAN framework to generate brushstrokes, which are then composed to create paintings. GANs are great generative models for AI Art but they are known to be notoriously difficult to train. To overcome GAN's limitations and to speed up the Neural Painter training, we applied Transfer Learning to the process reducing it from days to only hours, while achieving the same level of visual aesthetics in the final paintings generated. We report our approach and results in this work.""",
        """In autonomous driving, learning a segmentation model that can adapt to various environmental conditions is crucial. In particular, copying with severe illumination changes is an impelling need, as models trained on daylight data will perform poorly at nighttime. In this paper, we study the problem of Domain Adaptive Nighttime Semantic Segmentation (DANSS), which aims to learn a discriminative nighttime model with a labeled daytime dataset and an unlabeled dataset, including coarsely aligned day-night image pairs. To this end, we propose a novel Bidirectional Mixing (Bi-Mix) framework for DANSS, which can contribute to both image translation and segmentation adaptation processes. Specifically, in the image translation stage, Bi-Mix leverages the knowledge of day-night image pairs to improve the quality of nighttime image relighting. On the other hand, in the segmentation adaptation stage, Bi-Mix effectively bridges the distribution gap between day and night domains for adapting the model to the night domain. In both processes, Bi-Mix simply operates by mixing two samples without extra hyper-parameters, thus it is easy to implement. Extensive experiments on Dark Zurich and Nighttime Driving datasets demonstrate the advantage of the proposed Bi-Mix and show that our approach obtains state-of-the-art performance in DANSS."""
        ]

    example = st.selectbox("Choose an example abstract", example_prompts)

    # Take the message which needs to be processed
    message = st.text_area("...or paste a papers abstract to generate a title", example)
    # st.title(message)
    st.text("")
    models_to_choose = [
        "AryanLala/autonlp-Scientific_Title_Generator-34558227",
        "shamikbose89/mt5-small-finetuned-arxiv-cs-finetuned-arxiv-cs-full"
    ]

    BASE_MODEL = st.selectbox("Choose a model to generate the title", models_to_choose)

    def preprocess(text):
        if ((BASE_MODEL == "AryanLala/autonlp-Scientific_Title_Generator-34558227") |
          (BASE_MODEL == "shamikbose89/mt5-small-finetuned-arxiv-cs-finetuned-arxiv-cs-full")):
            return [text]
        else:
            st.error("Please select a model first")

    @st.cache(allow_output_mutation=True, suppress_st_warning=True, show_spinner=True)
    def load_model():
        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
        model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)
        return model, tokenizer

    def get_summary(text):
        model, tokenizer = load_model()
        preprocessed = preprocess(text)
        inputs = tokenizer(
            preprocessed, truncation=True, padding="longest", return_tensors="pt"
        )
        output = model.generate(
            **inputs,
            max_length=60,
            num_beams=10,
            num_return_sequences=1,
            temperature=1.5,
        )
        target_text = tokenizer.batch_decode(output, skip_special_tokens=True)
        return target_text[0]

    # Define function to run when submit is clicked
    def submit(message):
        if len(message) > 0:
            emoji = get_summary(message)
            html_str = f"""
            <style>
            p.a {{
            font: 20px Courier;
            }}
            </style>
            <p class="a">{emoji}</p>
            """

            st.markdown(html_str, unsafe_allow_html=True)
            # st.markdown(emoji)
        else:
            st.error("The text can't be empty")

    # Run algo when submit button is clicked
    if st.button("Submit"):
        submit(message)

    with st.expander("Additional Model Info"):
        st.write("""
        The models used were fine-tuned on subset of data from the [Arxiv Dataset](https://huggingface.co/datasets/arxiv_dataset)
        The task of the models is to suggest an appropraite title from the abstract of a scientific paper.
        
        The model [AryanLala/autonlp-Scientific_Title_Generator-34558227]() was trained on data
        from the Cs.AI (Artificial Intelligence) category of papers.

        The model [shamikbose89/mt5-small-finetuned-arxiv-cs-finetuned-arxiv-cs-full](https://huggingface.co/shamikbose89/mt5-small-finetuned-arxiv-cs-finetuned-arxiv-cs-full) 
        was trained on the categories: cs.AI, cs.LG, cs.NI, cs.GR cs.CL, cs.CV (Artificial Intelligence, Machine Learning, Networking and Internet Architecture, Graphics, Computation and Language, Computer Vision and Pattern Recognition)
        """)

    st.text('')
    st.markdown(
    '''<span style="color:blue; font-size:10px">App created by [@AlekseyDorkin](https://huggingface.co/AlekseyDorkin)
     and [@akshay7](https://huggingface.co/akshay7)</span>''', 
     unsafe_allow_html=True,
    )

if __name__ == "__main__":
    main()