Update app.py

app.py

@@ -6,62 +6,42 @@ import os
 @st.cache(allow_output_mutation=True)
 def load_model_cache():
  auth_token = os.environ.get("TOKEN_FROM_SECRET") or True
- tokenizer_en = T5Tokenizer.from_pretrained(
- "Voicelab/vlt5-base-keywords-v4_3-en", use_auth_token=auth_token
- )
- model_en = T5ForConditionalGeneration.from_pretrained(
- "Voicelab/vlt5-base-keywords-v4_3-en", use_auth_token=auth_token
- )
  
  tokenizer_pl = T5Tokenizer.from_pretrained(
- "Voicelab/vlt5-base-keywords-v4_3", use_auth_token=auth_token
+ "Voicelab/vlt5-base-rfc-v1_2", use_auth_token=auth_token
  )
  model_pl = T5ForConditionalGeneration.from_pretrained(
- "Voicelab/vlt5-base-keywords-v4_3", use_auth_token=auth_token
+ "Voicelab/vlt5-base-rfc-v1_2", use_auth_token=auth_token
  )
 
- return tokenizer_en, model_en, tokenizer_pl, model_pl
+ 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 = 1000
+max_length: int = 5000
 cache_size: int = 100
 
 st.set_page_config(
- page_title="DEMO - keywords generation",
+ page_title="DEMO - Reason for Contact detection",
  page_icon=img_favicon,
  initial_sidebar_state="expanded",
 )
 
 tokenizer_en, model_en, tokenizer_pl, model_pl = load_model_cache()
 
-def get_predictions(text, language):
- if language == "Polish":
+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=2,
+ no_repeat_ngram_size=1,
  num_beams=3,
  num_beam_groups=3,
- repetition_penalty=1.5,
- diversity_penalty=2.0,
- length_penalty=2.0,
+ min_length=10,
+ max_length=100,
  )
  predicted_kw = tokenizer_pl.decode(output[0], skip_special_tokens=True)
- elif language == "English":
- input_ids = tokenizer_en(text, return_tensors="pt", truncation=True).input_ids
- output = model_en.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_kw = tokenizer_en.decode(output[0], skip_special_tokens=True)
  return predicted_kw
 
 
@@ -73,7 +53,7 @@ def trim_length():
 if __name__ == "__main__":
  st.sidebar.image(img_short)
  st.image(img_full)
- st.title("VLT5 - keywords generation")
+ st.title("VLT5 - RfC generation")
 
  generated_keywords = ""
  user_input = st.text_area(
@@ -89,12 +69,11 @@ if __name__ == "__main__":
  "Select model to test",
  [
  "Polish",
- "English",
  ],
  )
 
- result = st.button("Generate keywords")
+ result = st.button("Find reason for contact")
  if result:
- generated_keywords = get_predictions(text=user_input, language=language)
- st.text_area("Generated keywords", generated_keywords)
- print(f"Input: {user_input}---> Keywords: {generated_keywords}")
+ generated_rfc = get_predictions(text=user_input, language=language)
+ st.text_area("Reason", generated_rfc)
+ print(f"Input: {user_input} ---> Reason for contact: {generated_rfc}")