Update app.py

app.py

@@ -1,3 +1,107 @@
+# 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
@@ -47,43 +151,11 @@ 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)
-    }
+    inputs = tokenizer.encode("summarize: " + st_text_area, return_tensors="pt", max_length=max_input_length, truncation=True)
 
     # compute predictions
-    outputs = model.generate(**inputs, do_sample=True, temperature=temperature)
-    decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    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
@@ -99,4 +171,4 @@ if len(st.session_state.titles) > 0:
     with st.container():
         st.subheader("Generated titles")
         for title in st.session_state.titles:
-            st.markdown("__" + title + "__")
+            st.markdown("__" + title + "__")