- model timing added

app.py

@@ -1,3 +1,5 @@
+import time
+
 import streamlit as st
 from transformers import pipeline
 from constants import tweet_generator_prompt, absa_prompt
@@ -6,10 +8,16 @@ from constants import tweet_generator_prompt, absa_prompt
 # Initialize the model and tokenizer once, to avoid reloading them on each user interaction
 @st.cache_resource
 def load_model():
+    start = time.time()
     classification_pipe = pipeline(
         "text-classification", model="tweetpie/toxic-content-detector", top_k=None)
+    print(f"Time to load the classification model: {time.time() - start:.2f}s")
+    start = time.time()
     absa_pipe = pipeline("text2text-generation", model="tweetpie/stance-aware-absa")
+    print(f"Time to load the absa model: {time.time() - start:.2f}s")
+    start = time.time()
     tweet_generation_pipe = pipeline("text2text-generation", model="tweetpie/stance-directed-tweet-generator")
+    print(f"Time to load the tweet generation model: {time.time() - start:.2f}s")
     return classification_pipe, absa_pipe, tweet_generation_pipe
 
 
@@ -48,17 +56,19 @@ classifier, absa, generator = load_model()
 if generate_button:
     with st.spinner('Generating the tweet...'):
         # Call the model with the aspects inputs
-        generated_tweet = generator(
-            tweet_generator_prompt.format(
-                ideology=model_selection.lower(),
-                pro_entities=pro_entities,
-                anti_entities=anti_entities,
-                neutral_entities=neutral_entities,
-                pro_aspects=pro_aspects,
-                anti_aspects=anti_aspects,
-                neutral_aspects=neutral_aspects
-            )
+        prompt = tweet_generator_prompt.format(
+            ideology=model_selection.lower(),
+            pro_entities=pro_entities,
+            anti_entities=anti_entities,
+            neutral_entities=neutral_entities,
+            pro_aspects=pro_aspects,
+            anti_aspects=anti_aspects,
+            neutral_aspects=neutral_aspects
         )
+        print("Prompt: ", prompt)
+        start = time.time()
+        generated_tweet = generator(prompt)
+        print(f"Time to generate the tweet: {time.time() - start:.2f}s")
 
         # Displaying the input and model's output
         st.write(f"Generated Tweet: {generated_tweet[0]['generated_text']}")
@@ -66,6 +76,7 @@ if generate_button:
     with st.spinner('Generating the Stance-Aware ABSA output...'):
         # Call the model with the aspects inputs
         absa_output = absa(absa_prompt.format(generated_tweet=generated_tweet[0]['generated_text']))
+        print("ABSA Output: ", absa_output)
         stances = [x.strip() for x in absa_output[0]['generated_text'].split(',')]
 
         st.write("Stance-Aware ABSA Output:")
@@ -79,10 +90,10 @@ if generate_button:
         st.write("Toxicity Classifier Output:")
         for i in range(3):
             if output[i]['label'] == 'LABEL_0':
-                # st.write(f"Non-Toxic Content: {output[i]['score']*100:.1f}%")
-                print(f"Non-Toxic Content: {output[i]['score']*100:.1f}%")
+                st.write(f"Non-Toxic Content: {output[i]['score']*100:.1f}%")
+                # print(f"Non-Toxic Content: {output[i]['score']*100:.1f}%")
             elif output[i]['label'] == 'LABEL_1':
-                # st.write(f"Toxic Content: {output[i]['score']*100:.1f}%")
-                print(f"Toxic Content: {output[i]['score']*100:.1f}%")
+                st.write(f"Toxic Content: {output[i]['score']*100:.1f}%")
+                # print(f"Toxic Content: {output[i]['score']*100:.1f}%")
             else:
                 continue