app.py

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

import nltk
nltk.download('punkt', quiet=True)
nltk.download('punkt_tab', quiet=True)

model_name = "Lifeinhockey/T5_fine_tuning"
max_input_length = 512

st.header("Generate candidate titles for articles from V. Gorsky")

st_model_load = st.text('Loading title generator model...')

@st.cache_resource
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 = ""
st_text_area = st.text_area('Text to generate the title 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, max_length=512, truncation=True, return_tensors="pt")

    # compute predictions
    outputs = model.generate(
        **inputs, 
        do_sample=True, 
        temperature=temperature,
        max_length=64,
        num_return_sequences=num_titles
    )
    
    decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    
    # Обработка результатов
    predicted_titles = []
    for decoded_output in decoded_outputs:
        decoded_output = decoded_output.strip()
        if decoded_output:  # Проверяем, что строка не пустая
            sentences = decoded_output.split('. ')
            if sentences:
                first_sentence = sentences[0]
                if not first_sentence.endswith('.'):
                    first_sentence += '.'
                predicted_titles.append(first_sentence)
            else:
                predicted_titles.append(decoded_output)
        else:
            predicted_titles.append("Не удалось сгенерировать заголовок")

    st.session_state.titles = predicted_titles

# generate title button
st_generate_button = st.button('Generate title', 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 titles")
        for title in st.session_state.titles:
            st.markdown("__" + title + "__")