Update app.py

app.py

@@ -1,15 +1,20 @@
 import gradio as gr
 import numpy as np
 import openai
+import time
 
 from sentence_transformers import SentenceTransformer
 from langchain.prompts import PromptTemplate
 from collections import Counter
 
 
-def process(api, caption, category, asr, ocr):
-    openai.api_key = api
+def process(caption, category, asr, ocr):
     preference = "兴趣标签"
+    if len(category) == 0:
+        category = "空"
+    if len(asr) == 0:
+        asr = "空"
+
     example = "例如，给定一个视频，它的\"标题\"为\"长安系最便宜的轿车，4W起很多人都看不上它，但我知道车只是代步工具，又需要什么面子呢" \
               "！\"，\"类别\"为\"汽车\"，\"ocr\"为\"长安系最便宜的一款轿车\"，\"asr\"为\"我不否认现在的国产和合资还有一定的差距，" \
               "但确实是他们让我们5万开了MP V8万开上了轿车，10万开张了ICV15万开张了大七座。\"，\"{}\"生成机器人推断出合理的\"{}\"为\"" \
@@ -41,7 +46,7 @@ def process(api, caption, category, asr, ocr):
             res += ans.split('、')
 
         tag_count = Counter(res)
-        tag_count = sorted(tag_count.items(), key=lambda x: x[1], reverse=True)
+        tag_count = sorted(tag_count.items(), key=lambda x: x[1], reverse=True)[:10]
 
         tags_embed = np.load('./tag_data/tags_embed.npy')
         tags_dis = np.load('./tag_data/tags_dis.npy')
@@ -65,11 +70,9 @@ def process(api, caption, category, asr, ocr):
                 if score > 0.9:
                     final_ans.append(all_tags[j])
 
-        print(final_ans)
-
         final_ans = Counter(final_ans)
         final_ans = sorted(final_ans.items(), key=lambda x: x[1], reverse=True)[:5]
-        final_ans = [_[0] for _ in final_ans if _[1] > 1]
+        final_ans = [_[0] for _ in final_ans]
 
         return "、".join(final_ans)
 
@@ -77,17 +80,118 @@ def process(api, caption, category, asr, ocr):
         return 'api error'
 
 
+def connection(api_key):
+    openai.api_key = api_key
+    time.sleep(5)
+
+
 with gr.Blocks() as demo:
-    text_api = gr.Textbox(label='OpenAI API key')
-    text_caption = gr.Textbox(label='Caption')
-    text_category = gr.Textbox(label='Category')
-    text_asr = gr.Textbox(label='ASR')
-    text_ocr = gr.Textbox(label='OCR')
+    gr.Markdown("<h3><center>TagGPT</center></h3>")
+    gr.Markdown(
+        """
+        This is a demo to the work  [TagGPT: Large Language Models are Zero-shot Multimodal Taggers](https://github.com/TencentARC/TagGPT).<br>
+        This space connects TagGPT to provide tagging service based on the tag set (from the Kuaishou data).<br>  
+        """
+    )
+    gr.Markdown(
+        """
+        Step 1: input openai api key (sk-...) and click "connect" button.
+        """
+    )
+
+    with gr.Column(variant="panel"):
+        with gr.Row(variant="compact"):
+            text_api = gr.Textbox(
+                label='OpenAI API key',
+                placeholder="Paste your OpenAI API key here️",
+                type="password"
+            ).style(
+                container=False,
+            )
+            btn_api = gr.Button("Connect").style(full_width=False)
+
+    btn_api.click(connection, [text_api])
+
+    gr.Markdown(
+        """
+        Step 2: fall in the four items (i.e., caption, category, ASR, and OCR) in the appropriate input fields.
+        Click any item of the "Examples" to quickly see the tagging results.
+        """
+    )
+
+    text_caption = gr.Textbox(
+        label='标题(Caption)',
+        placeholder="Indispensable"
+    )
+    text_category = gr.Textbox(
+        label='类别(Category)',
+        placeholder="Dispensable"
+    )
+    text_asr = gr.Textbox(
+        label='ASR',
+        placeholder="Dispensable"
+    )
+    text_ocr = gr.Textbox(
+        label='OCR',
+        placeholder="Dispensable"
+    )
 
     text_output = gr.Textbox(value='', label='Output')
 
     btn = gr.Button(value='Submit')
-    btn.click(process, inputs=[text_api, text_caption, text_category, text_asr, text_ocr], outputs=[text_output])
+    btn.click(process, inputs=[text_caption, text_category, text_asr, text_ocr], outputs=[text_output])
+
+    examples = [
+        [
+            '正确解决iCloud储存空间将满的实用技巧',
+            '高新数码',
+            '你的iCloud一直提示，iCloud储存空间将满，而且还会把你的iPhone iPad Mac等相关连设备都弹一遍，怎么办，烦死了。其实只要用iPho'
+            'ne打开设置，点击上方的头像栏，然后点击iCloud，选择管理账户储存空间，进入后，把这里不需要云备份按自己的需要进行删除，然后再回'
+            '到上一级页面，把这里的iCloud云备份关闭就可以了。',
+            'iCloud储存空间将满实用技巧'
+        ],
+        [
+            '30平米迷你小公寓如何设计？创意设计多出一间房！',
+            '房产家居',
+            '这是一个30平米的迷你公寓，入户没有鞋柜，没有餐厅，没有书桌��置，咱们可以这样设计，首先吧卫生间墙内推30公分，嵌入鞋柜，'
+            '厨房旁边装三联动推拉门，进出舒适，隔绝油烟，把原有沙发该到卧室位置，可以让出一个餐厅位置，沙发前移做一组书柜，内藏隐形'
+            '壁床，打开是卧室，收起便是小客厅，对面墙壁，再做一组书柜和书桌，各种需求都满足。',
+            '30平迷你公寓怎么设计'
+        ],
+        [
+            '好好钓鱼！千万别算账',
+            '生活',
+            '一只鱼竿四千多，用它来钓多少鱼才能回本，钓鱼是娱乐而算账却很扎心，这是一根国产标配波爬竿，我的深海标配，主要针对海绵大'
+            '型鱼类，比如金枪鱼、gt、鬼头刀等，经济价值最高的算是金枪鱼了，一不小心钓上一条鱼竿就回来了，不小心又钓上一条轮子就回来了'
+            '当你发现你的渔具全是渔获换来的，吹起牛来是不是更加有底气。一支鱼竿，一副轮子，又来一套，哈哈。你钓了几套渔具呢？',
+            '钓鱼很开心，算账很扎心'
+        ],
+        [
+            '很多朋友再问，为什么蒸出来的花卷葱没变色，今天告诉你',
+            '美食',
+            '为什么你自己在家做的画卷，蒸出来之后，它的葱总是黄色的呢？而早餐店里面的画卷蒸出来过后它的葱都是绿油油的，那么今天呢，'
+            '把方法教给你，那就是我们做花卷的时候呢，再从里面加入少量的泡打粉，拌均匀，这样做出来的画卷呢，蒸熟过后，它的葱也是绿'
+            '油油的，大家可以看一下。',
+            '为什么你蒸的葱油花卷总会变黄'
+        ],
+        [
+            '豪华七座SUV 比亚迪唐DM-i冠军版带你体验出行新方式',
+            '汽车',
+            '你听说了吗，2023年冠军版本的唐DM-i已经上市了，上市当天，汉唐总销量已经超过了8000多台，大家觉得这个成绩怎么样，那么，今天'
+            '艾琳就带大家全民啊了解一下，2023年冠军版本的唐DM-i，新款的唐DM-i和老款之间的差别并不大，但是我们的新款新增了冰川蓝的颜色'
+            '还有啊，你看这个轮毂是不是和我们老款252公里的轮毂一样呢，另外新款多了铝合金底盘，FSD可变主力悬挂，快充也升级到了40千瓦，还'
+            '有专门的快充孔，20分钟就可以让他满血复活，2023款唐DM-i一共分为三个配置，尊贵尊荣尊享，不过都是7座，那配置上的差异，艾琳'
+            '就把它放在图上了，方便大家对比。想入手的朋友记住了，武汉比亚迪找艾琳，艾琳带你早提车。',
+            '豪华标杆SUV，23款唐DM-1冠军版'
+        ]
+    ]
+
+    gr.Examples(
+        examples,
+        [text_caption, text_category, text_asr, text_ocr],
+        text_output,
+        process
+    )
 
 
 if __name__ == "__main__":