Added min karma threshold and color logic

app.py

@@ -8,6 +8,7 @@ import torch
 import time
 
 MIN_SCORE_THRESHOLD=50
+NUM_VIBE_LABELS=5
 
 device = 0
 if not torch.cuda.is_available():
@@ -22,7 +23,7 @@ reddit = praw.Reddit(
     user_agent="Hugging Face Vibe Checker",
 )
 
-vibes = [
+vibe_names = [
     "wholesome",
     "chill",
     "funny",
@@ -43,6 +44,7 @@ vibes = [
     "meme",
 ]
 
+vibe_name_color_map={vibe:"red" for vibe in vibe_names}
 
 def vibe_check(url, min_karma):
     comments = get_comments(url)
@@ -54,19 +56,26 @@ def vibe_check(url, min_karma):
     print("Starting comment classification.")
     start = time.time()
     classes = pipe(
-        comment_bodies,
-        candidate_labels=vibes,
+        comment_bodies[:1],
+        candidate_labels=vibe_names,
     )
     end = time.time()
 
     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])
+    for i in range(len(vibe_names)):
+        averages[vibe_names[i]] = mean([c["scores"][i] for c in classes])
         
     return averages
 
+def get_vibes_html(vibes):
+    return "    ".join(
+        [
+            f"<span style=\"color:{vibe_name_color_map[vibe]};font-size:{12 * i}px\">{vibe}</span>"
+            for i, vibe in enumerate(vibes)
+        ]) 
+
 def get_comments(url):
     submission = reddit.submission(url=url)
     comments = []
@@ -93,7 +102,7 @@ def comment_compare(comment):
 with gr.Blocks() as demo:
     url = gr.Textbox(label="Url")
     min_karma = gr.Textbox(label="Minimum Karma", value=MIN_SCORE_THRESHOLD)
-    output = gr.Label(label="Output", num_top_classes=5)
+    output = gr.Label(label="Output", num_top_classes=NUM_VIBE_LABELS)
     submit_button = gr.Button("Submit")
 
     submit_button.click(
@@ -111,5 +120,7 @@ with gr.Blocks() as demo:
         vibe_check,
         # cache_examples=True,
     )
-
+    
+    
+demo.theme=gr.themes.Base()
 demo.launch()