updated to new version. using load api and improved looks. it's fast now.

app.py

@@ -4,6 +4,7 @@ import sys
 import nltk
 import praw
 import matplotlib
+from tqdm import tqdm
 import gradio as gr
 import pandas as pd
 import praw.exceptions
@@ -44,19 +45,13 @@ def preprocessData(df):
     df["text"] = df["text"].apply(lambda x: re.sub(r"http\S+", "", x, flags=re.M))
     df["text"] = df["text"].apply(lambda x: re.sub(r"^>.+", "", x, flags=re.M))
 
-    smax = df.score.max()
-
-    threshold = round(0.05 * smax)
-
-    df = df[df.score >= threshold]
-
-    # empirically, having more than 200 comments doesn't change much but slows down the summarizer.
-    if len(df.text) >= 200:
-        df = df[:200]
+    # The df is sorted by comment score
+    # Empirically, having more than ~100 comments doesn't change much but slows down the summarizer.
+    # Slowdown is not present with load api but still seems good to limit low score comments.
+    if len(df.text) >= 128:
+        df = df[:128]
 
     # chunking to handle giving the model too large of an input which crashes
-
-    # chunked = list(index_chunk(df.text))
     chunked = sentence_chunk(df.text)
 
     return chunked
@@ -123,56 +118,67 @@ def summarizer(url: str) -> str:
 
     # pushshift.io submission comments api doesn't work so have to use praw
     df = getComments(url=url)
+
+    submission_title = '# ' + df.submission_title.unique()[0]
+
     chunked_df = preprocessData(df)
-    submission_title = df.submission_title.unique()[0]
 
     text = ' '.join(chunked_df)
-    # transparent bg: background_color=None, mode='RGBA'
+    # transparent bg: background_color=None, mode='RGBA')
     wc_opts = dict(collocations=False, width=1920, height=1080)
     wcloud = WordCloud(**wc_opts).generate(text)
 
-    fig = plt.figure(figsize=(16, 10))
-    fig.patch.set_alpha(0.0)
-    plt.imshow(wcloud)
+    plt.imshow(wcloud, aspect='auto')
     plt.axis("off")
-    plt.tight_layout()
+    plt.gca().set_position([0, 0, 1, 1])
+    plt.autoscale(tight=True)
+    fig = plt.gcf()
+    fig.patch.set_alpha(0.0)
+    fig.set_size_inches((12, 7))
 
     lst_summaries = []
 
-    nlp = pipeline('summarization', model="sshleifer/distilbart-cnn-12-6")
-
-    for grp in chunked_df:
+    for grp in tqdm(chunked_df):
         # treating a group of comments as one block of text
-        result = nlp(grp, max_length=500)[0]["summary_text"]
+        result = sum_api(grp)
         lst_summaries.append(result)
 
-    joined_summaries = ' '.join(lst_summaries).replace(" .", ".")
+    long_output = ' '.join(lst_summaries).replace(" .", ".")
 
-    total_summary = nlp(joined_summaries, max_length=500)[0]["summary_text"].replace(" .", ".")
+    short_output = sum_api(long_output).replace(" .", ".")
 
-    short_output = submission_title + '\n' + '\n' + total_summary
+    sentiment = clf_api(short_output)
 
-    long_output = submission_title + '\n' + '\n' + joined_summaries
-
-    return short_output, long_output, fig
+    return submission_title, short_output, long_output, sentiment, fig
 
 
 if __name__ == "__main__":
 
+    sum_model = "models/sshleifer/distilbart-cnn-12-6"
+    clf_model = "models/finiteautomata/bertweet-base-sentiment-analysis"
+
+    hf_token = os.environ["HF_TOKEN"]
+
+    sum_api = gr.Interface.load(sum_model, api_key=hf_token)
+    clf_api = gr.Interface.load(clf_model, api_key=hf_token)
+
     with gr.Blocks(css=".gradio-container {max-width: 900px !important; width: 100%}") as demo:
         submission_url = gr.Textbox(label='Post URL')
 
         sub_btn = gr.Button("Summarize")
 
+        title = gr.Markdown("")
+
         with gr.Row():
             short_summary = gr.Textbox(label='Short Summary')
-            thread_cloud = gr.Plot(label='Word Cloud')
+            summary_sentiment = gr.Label(label='Sentiment')
 
+        thread_cloud = gr.Plot(label='Word Cloud')
         long_summary = gr.Textbox(label='Long Summary')
 
         sub_btn.click(fn=summarizer,
                       inputs=[submission_url],
-                      outputs=[short_summary, long_summary, thread_cloud])
+                      outputs=[title, short_summary, long_summary, summary_sentiment, thread_cloud])
 
     try:
         demo.launch()

requirements.txt

@@ -3,6 +3,7 @@ matplotlib==3.7.1
 nltk==3.8.1
 pandas==1.5.3
 praw==7.7.0
+tqdm==4.65.0
 transformers==4.26.1
 wordcloud==1.8.2.2
 torch==1.13.1