Added label UI and added min karma threshold

app.py

@@ -3,8 +3,19 @@ import gradio as gr
 import praw
 from praw.models import MoreComments
 import os
+from statistics import mean
+import torch
+import time
 
-pipe = pipeline(model="facebook/bart-large-mnli")
+MIN_SCORE_THRESHOLD=50
+
+device = 0
+if not torch.cuda.is_available():
+    print("GPU isn't available. Running on CPU.")
+    device = -1
+
+
+pipe = pipeline(model="facebook/bart-large-mnli", device=device)
 reddit = praw.Reddit(
     client_id=os.environ.get("reddit_client_id"),
     client_secret=os.environ.get("reddit_client_secret"),
@@ -33,18 +44,28 @@ vibes = [
 ]
 
 
-def vibe_check(url):
+def vibe_check(url, min_karma):
     comments = get_comments(url)
-    comments.sort(key=comment_compare, reverse=True)
-    comment_bodies = [c["Comment"] for c in comments[0:10]]
+    #comments.sort(key=comment_compare, reverse=True)
+    comment_bodies = [c["Comment"] for c in comments if c["Score"] > int(min_karma)]
+
+    print("Total comments: " + str(len(comment_bodies)))
 
+    print("Starting comment classification.")
+    start = time.time()
     classes = pipe(
         comment_bodies,
         candidate_labels=vibes,
     )
+    end = time.time()
 
-    return classes
+    print("Comment classification took: " + str(end - start) + "ms")
 
+    averages = {}
+    for i in range(len(vibes)):
+        averages[vibes[i]] = mean([c["scores"][i] for c in classes])
+        
+    return averages
 
 def get_comments(url):
     submission = reddit.submission(url=url)
@@ -71,11 +92,12 @@ def comment_compare(comment):
 
 with gr.Blocks() as demo:
     url = gr.Textbox(label="Url")
-    output = gr.Textbox(label="Output")
+    min_karma = gr.Textbox(label="Minimum Karma", value=MIN_SCORE_THRESHOLD)
+    output = gr.Label(label="Output", num_top_classes=5)
     submit_button = gr.Button("Submit")
 
     submit_button.click(
-        fn=vibe_check, inputs=url, outputs=output, api_name="vibe_check"
+        fn=vibe_check, inputs=[url, min_karma], outputs=output, api_name="vibe_check"
     )
 
     gr.Examples(

requirements.txt

@@ -2,4 +2,5 @@ pylint
 transformers
 torch
 praw
-gradio
+gradio
+torch