app.py

from transformers import T5ForConditionalGeneration, T5Tokenizer
import streamlit as st
from PIL import Image
import os


@st.cache(allow_output_mutation=True)
def load_model_cache():
 auth_token = os.environ.get("TOKEN_FROM_SECRET") or True

 tokenizer_pl = T5Tokenizer.from_pretrained(
 "Voicelab/vlt5-base-rfc-v1_2", use_auth_token=auth_token
 )
 model_pl = T5ForConditionalGeneration.from_pretrained(
 "Voicelab/vlt5-base-rfc-v1_2", use_auth_token=auth_token
 )
 model_det_pl = T5ForConditionalGeneration.from_pretrained(
 "Voicelab/vlt5-base-rfc-detector-1.0", use_auth_token=auth_token
 )

 return tokenizer_pl, model_pl, model_det_pl


img_full = Image.open("images/vl-logo-nlp-blue.png")
img_short = Image.open("images/sVL-NLP-short.png")
img_favicon = Image.open("images/favicon_vl.png")
max_length: int = 5000
cache_size: int = 100

st.set_page_config(
 page_title="DEMO - Reason for Contact generation",
 page_icon=img_favicon,
 initial_sidebar_state="expanded",
)

tokenizer_pl, model_pl, model_det_pl = load_model_cache()


def get_predictions(text, mode):
 input_ids = tokenizer_pl(text, return_tensors="pt", truncation=True).input_ids
 if mode == "Polish - RfC Generation":
 output = model_pl.generate(
 input_ids,
 no_repeat_ngram_size=1,
 num_beams=3,
 num_beam_groups=3,
 min_length=10,
 max_length=100,
 )
 elif mode == "Polish - RfC Detection":
 output = model_det_pl.generate(
 input_ids,
 no_repeat_ngram_size=2,
 num_beams=3,
 num_beam_groups=3,
 repetition_penalty=1.5,
 diversity_penalty=2.0,
 length_penalty=2.0,
 )
 predicted_rfc = tokenizer_pl.decode(output[0], skip_special_tokens=True)
 return predicted_rfc


def trim_length():
 if len(st.session_state["input"]) > max_length:
 st.session_state["input"] = st.session_state["input"][:max_length]


if __name__ == "__main__":
 st.sidebar.image(img_short)
 st.image(img_full)
 st.title("VLT5 - Reason for Contact generator")

 generated_keywords = ""
 user_input = st.text_area(
 label=f"Input text (max {max_length} characters)",
 value="",
 height=300,
 on_change=trim_length,
 key="input",
 )

 mode = st.sidebar.title("Model settings")
 mode = st.sidebar.radio(
 "Select model to test",
 [
 "Polish - RfC Generation",
 "Polish - RfC Detection",
 ],
 )

 result = st.button("Find reason for contact")
 if mode == "Polish - RfC Generation (accepts whole conversation)":
 st.markdown("### You selected RfC Generation model.")
 st.markdown("-- *Input*: Whole conversation. Should specify roles (e.g. **AGENT: Hello, how can I help you? CLIENT: Hi, I would like to report a stolen card.**")
 st.markdown("-- *Output*: Reason for calling for the whole conversation.")
 text_area = "Put a whole conversation or full e-mail here."
 
 elif mode == "Polish - RfC Detection (accepts one turn)":
 st.markdown("### You selected RfC Detection model.")
 st.markdown("-- *Input*: A single turn from the conversation e.g. **'Hello, how can I help you?'** or **'Hi, I would like to report a stolen card.'**")
 st.markdown("-- *Output*: Model will return an empty string if a turn possibly does not includes Reason for Calling, or a sentence if the RfC is detected.")
 text_area = "Put a single turn or a few sentences here."
 
 if result:
 generated_rfc = get_predictions(text=user_input, mode=mode)
 st.text_area(text_area, generated_rfc)
 print(f"Input: {user_input} ---> Reason for contact: {generated_rfc}")