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}")