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
 )

 return tokenizer_pl, model_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 detection",
 page_icon=img_favicon,
 initial_sidebar_state="expanded",
)

tokenizer_pl, model_pl = load_model_cache()


def get_predictions(text):
 input_ids = tokenizer_pl(text, return_tensors="pt", truncation=True).input_ids
 output = model_pl.generate(
 input_ids,
 no_repeat_ngram_size=1,
 num_beams=3,
 num_beam_groups=3,
 min_length=10,
 max_length=100,
 )
 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 - RfC generation")

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

 language = st.sidebar.title("Model settings")
 language = st.sidebar.radio(
 "Select model to test",
 [
 "Polish",
 ],
 )

 result = st.button("Find reason for contact")
 if result:
 generated_rfc = get_predictions(text=user_input)
 st.text_area("Reason", generated_rfc)
 print(f"Input: {user_input} ---> Reason for contact: {generated_rfc}")