app.py

import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import math
import torch

model_name = "Nopphakorn/mt5-small-thaisum-512-title"
max_input_length = 512

st.header("Generate headline titles for Thai news")

st_model_load = st.text('Loading headlines summarizer model...')

@st.cache(allow_output_mutation=True)
def load_model():
    print("Loading model...")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    print("Model loaded!")
    return tokenizer, model

tokenizer, model = load_model()
st.success('Model loaded!')
st_model_load.text("")

with st.sidebar:
    st.header("Model parameters")
    if 'num_titles' not in st.session_state:
        st.session_state.num_titles = 5
    def on_change_num_titles():
        st.session_state.num_titles = num_titles
    num_titles = st.slider("Number of titles to generate", min_value=1, max_value=10, value=1, step=1, on_change=on_change_num_titles)
    if 'temperature' not in st.session_state:
        st.session_state.temperature = 0.7
    def on_change_temperatures():
        st.session_state.temperature = temperature
    temperature = st.slider("Temperature", min_value=0.1, max_value=1.5, value=0.6, step=0.05, on_change=on_change_temperatures)
    st.markdown("_High temperature means that results are more random_")

if 'text' not in st.session_state:
    st.session_state.text = '/"วรงค์ เดชกิจวิกรม/" วอน กกต. ทบทวนคะแนนปาร์ตี้ลิสต์ หลังมีผู้หวังดีแจ้งไทยภักดีมีคะแนนสูงกว่าบัญชีรายชื่อคนสุดท้ายของพรรคก้าวไกล จี้คืน 1 เก้าอี้ให้ไทยภักดี วันที่ 17 พฤษภาคม 2566 ผู้สื่อข่าวรายงานว่า นายแพทย์วรงค์ เดชกิจวิกรม หัวหน้าพรรคไทยภักดี โพสต์ข้อความผ่านเฟซบุ๊กส่วนตัว ส่งเสียงถึงคณะกรรมการการเลือกตั้ง (กกต.) ว่า อยากให้ทบทวนคะแนนของ ส.ส.บัญชีรายชื่อ โดยละเอียด เพราะมีผู้หวังดีแจ้งว่า ส.ส.บัญชีรายชื่อคนสุดท้ายของพรรคก้าวไกล ได้คะแนนคิดเป็น 0.23% (90,355 คะแนน) ขณะที่ของพรรคไทยภักดีได้ 0.26% (102,237 คะแนน) ดังนั้น ส.ส.บัญชีรายชื่อคนสุดท้ายที่ได้รับของพรรคก้าวไกล ควรเป็นของพรรคไทยภักดี ผมหวังว่าพรรคไทยภักดี จะได้รับความยุติธรรม และทุกคะแนนที่พี่น้องประชาชนเลือกพรรคไทยภักดี จะได้เป็นประโยชน์นำพาความต้องการ ของพี่น้องเข้าสู่การนำเสนอต่อไป/"'
st_text_area = st.text_area('Text to summarize the headlines for', value=st.session_state.text, height=500)

def generate_title():
    st.session_state.text = st_text_area

    # tokenize text
    inputs = ["summarize: " + st_text_area]
    inputs = tokenizer(inputs, return_tensors="pt")

    # compute span boundaries
    num_tokens = len(inputs["input_ids"][0])
    print(f"Input has {num_tokens} tokens")
    max_input_length = 500
    num_spans = math.ceil(num_tokens / max_input_length)
    print(f"Input has {num_spans} spans")
    overlap = math.ceil((num_spans * max_input_length - num_tokens) / max(num_spans - 1, 1))
    spans_boundaries = []
    start = 0
    for i in range(num_spans):
        spans_boundaries.append([start + max_input_length * i, start + max_input_length * (i + 1)])
        start -= overlap
    print(f"Span boundaries are {spans_boundaries}")
    spans_boundaries_selected = []
    j = 0
    for _ in range(num_titles):
        spans_boundaries_selected.append(spans_boundaries[j])
        j += 1
        if j == len(spans_boundaries):
            j = 0
    print(f"Selected span boundaries are {spans_boundaries_selected}")

    # transform input with spans
    tensor_ids = [inputs["input_ids"][0][boundary[0]:boundary[1]] for boundary in spans_boundaries_selected]
    tensor_masks = [inputs["attention_mask"][0][boundary[0]:boundary[1]] for boundary in spans_boundaries_selected]

    inputs = {
        "input_ids": torch.stack(tensor_ids),
        "attention_mask": torch.stack(tensor_masks)
    }

    # compute predictions
    outputs = model.generate(**inputs, do_sample=True, temperature=temperature, max_new_tokens=64)
    decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

    st.session_state.titles = decoded_outputs

# generate title button
st_generate_button = st.button('Generate headlines', on_click=generate_title)

# title generation labels
if 'titles' not in st.session_state:
    st.session_state.titles = []

if len(st.session_state.titles) > 0:
    with st.container():
        st.subheader("Generated headlines")
        for title in st.session_state.titles:
            st.markdown("__" + title + "__")