# import streamlit as st # from transformers import AutoTokenizer, AutoModelForSeq2SeqLM # import nltk # import math # import torch # # model_name = "fabiochiu/t5-base-medium-title-generation" # model_name = "momoyukki/t5-base-news-title-generation_model" # max_input_length = 512 # st.header("Generate candidate titles for news") # st_model_load = st.text('Loading title generator 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) # nltk.download('punkt') # 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, 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) # decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True) # predicted_titles = [nltk.sent_tokenize(decoded_output.strip())[0] for decoded_output in decoded_outputs] # 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 + "__") import streamlit as st from transformers import AutoTokenizer, AutoModelForSeq2SeqLM import nltk import math import torch # model_name = "fabiochiu/t5-base-medium-title-generation" model_name = "momoyukki/t5-base-news-title-generation_model" max_input_length = 512 st.header("Generate candidate titles for news") st_model_load = st.text('Loading title generator 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) nltk.download('punkt') 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 = tokenizer.encode("summarize: " + st_text_area, return_tensors="pt", max_length=max_input_length, truncation=True) # compute predictions outputs = model.generate(inputs, num_return_sequences=num_titles, temperature=temperature, max_length=150) decoded_outputs = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs] predicted_titles = [nltk.sent_tokenize(decoded_output.strip())[0] for decoded_output in decoded_outputs] 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 + "__")