Update app.py

app.py

@@ -14,8 +14,11 @@ def load_model_cache():
  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
+ return tokenizer_pl, model_pl, model_det_pl
 
 
 img_full = Image.open("images/vl-logo-nlp-blue.png")
@@ -25,24 +28,35 @@ max_length: int = 5000
 cache_size: int = 100
 
 st.set_page_config(
- page_title="DEMO - Reason for Contact detection",
+ page_title="DEMO - Reason for Contact generation",
  page_icon=img_favicon,
  initial_sidebar_state="expanded",
 )
 
-tokenizer_pl, model_pl = load_model_cache()
+tokenizer_pl, model_pl, model_det_pl = load_model_cache()
 
 
-def get_predictions(text):
+def get_predictions(text, mode):
  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,
- )
+ 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.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
 
@@ -55,7 +69,7 @@ def trim_length():
 if __name__ == "__main__":
  st.sidebar.image(img_short)
  st.image(img_full)
- st.title("VLT5 - RfC generation")
+ st.title("VLT5 - Reason for Contact generator")
 
  generated_keywords = ""
  user_input = st.text_area(
@@ -66,16 +80,29 @@ if __name__ == "__main__":
  key="input",
  )
 
- language = st.sidebar.title("Model settings")
- language = st.sidebar.radio(
+ mode = st.sidebar.title("Model settings")
+ mode = st.sidebar.radio(
  "Select model to test",
  [
- "Polish",
+ "Polish - RfC Generation",
+ "Polish - RfC Detection",
  ],
  )
 
  result = st.button("Find reason for contact")
+ if mode == "Polish - RfC Generation (accepts whole conversation)":
+ print("You selected RfC Generation model.")
+ print("-- 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.")
+ print("-- 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)":
+ print("You selected RfC Detection model.")
+ print("-- 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.'")
+ print("-- 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)
- st.text_area("Reason", generated_rfc)
+ 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}")