add icons for model pools

app.py

@@ -253,6 +253,7 @@ def move_to_second_view(btn):
         info['example2'],
         info['example3'],
         info['example4'],
+        info['thumb-tiny'],        
         gr.update(choices=load_mode_list, value=load_mode_list[0]),
         "",
     )
@@ -267,6 +268,7 @@ def download_completed(
     gen_config_path,
     gen_config_sum_path,
     load_mode,
+    thumbnail_tiny,
     force_download,
 ):
     tmp_args = types.SimpleNamespace()
@@ -275,6 +277,7 @@ def download_completed(
     tmp_args.gen_config_path = gen_config_path
     tmp_args.gen_config_summarization_path = gen_config_sum_path
     tmp_args.force_download_ckpt = force_download
+    tmp_args.thumbnail_tiny = thumbnail_tiny
 
     tmp_args.mode_cpu = True if load_mode == "cpu" else False
     tmp_args.mode_mps = True if load_mode == "apple silicon" else False
@@ -645,7 +648,8 @@ def main(args):
 |          half precision        |        load_in_8bit       |         load_in_4bit      | 
 | ------------------------------ | ------------------------- | ------------------------- | 
 |   {round(7830/1024., 1)}GiB    | {round(5224/1024., 1)}GiB | {round(4324/1024., 1)}GiB |
-""")                        
+""")
+                        model_thumbnail_tiny = gr.Textbox("", visible=False)
 
                 with gr.Column():
                     gen_config_path = gr.Dropdown(
@@ -686,7 +690,7 @@ def main(args):
                         with gr.Tab("Ex4"):
                             example_showcase4 = gr.Chatbot(
                                 [("hello", "world"), ("damn", "good")]
-                            )                            
+                            )
 
                 with gr.Row():
                     back_to_model_choose_btn = gr.Button("Back")
@@ -820,7 +824,7 @@ def main(args):
                         model_image, model_name, model_params, model_base, model_ckpt,
                         model_desc, model_vram, gen_config_path, 
                         example_showcase1, example_showcase2, example_showcase3, example_showcase4,
-                        load_mode,
+                        model_thumbnail_tiny, load_mode, 
                         progress_view
                     ]
                 )
@@ -885,7 +889,7 @@ def main(args):
                 lambda: "Start downloading/loading the model...", None, txt_view
             ).then(
                 download_completed,
-                [model_name, model_base, model_ckpt, gen_config_path, gen_config_sum_path, load_mode, force_redownload],
+                [model_name, model_base, model_ckpt, gen_config_path, gen_config_sum_path, load_mode, model_thumbnail_tiny, force_redownload],
                 [progress_view2]
             ).then(
                 lambda: "Model is fully loaded...", None, txt_view

chats/alpaca.py

@@ -9,14 +9,29 @@
 
 def build_prompts(ppmanager, user_message, global_context, win_size=3):
     dummy_ppm = copy.deepcopy(ppmanager)
+
     dummy_ppm.ctx = global_context
+    for pingpong in dummy_ppm.pingpongs:
+        pong = pingpong.pong
+        first_sentence = pong.split(" ")[0]
+        if first_sentence != "" and \
+            pre.contains_image_markdown(first_sentence):
+            pong = ' '.join(pong.split(" ")[1:])
+            pingpong.pong = pong
+
     lws = CtxLastWindowStrategy(win_size)
 
     prompt = lws(dummy_ppm)
     return prompt
 
 def text_stream(ppmanager, streamer):
+    count = 0
+
     for new_text in streamer:
+        if count == 0:
+            ppmanager.append_pong(f"![]({global_vars.model_thumbnail_tiny})***[{global_vars.model_type}]:*** ")
+            count = count + 1
+
         ppmanager.append_pong(new_text)
         yield ppmanager, ppmanager.build_uis()
 

chats/alpaca_gpt4.py

@@ -9,14 +9,29 @@
 
 def build_prompts(ppmanager, user_message, global_context, win_size=3):
     dummy_ppm = copy.deepcopy(ppmanager)
+
     dummy_ppm.ctx = global_context
+    for pingpong in dummy_ppm.pingpongs:
+        pong = pingpong.pong
+        first_sentence = pong.split(" ")[0]
+        if first_sentence != "" and \
+            pre.contains_image_markdown(first_sentence):
+            pong = ' '.join(pong.split(" ")[1:])
+            pingpong.pong = pong
+
     lws = CtxLastWindowStrategy(win_size)
 
     prompt = lws(dummy_ppm)
     return prompt
 
 def text_stream(ppmanager, streamer):
+    count = 0
+
     for new_text in streamer:
+        if count == 0:
+            ppmanager.append_pong(f"![]({global_vars.model_thumbnail_tiny})***[{global_vars.model_type}]:*** ")
+            count = count + 1
+
         ppmanager.append_pong(new_text)
         yield ppmanager, ppmanager.build_uis()
 

chats/alpacoom.py

@@ -9,10 +9,18 @@
 
 def build_prompts(ppmanager, user_message, global_context, win_size=3):
     dummy_ppm = copy.deepcopy(ppmanager)
-    lws = CtxLastWindowStrategy(win_size)
 
     dummy_ppm.ctx = global_context
-
+    for pingpong in dummy_ppm.pingpongs:
+        pong = pingpong.pong
+        first_sentence = pong.split(" ")[0]
+        if first_sentence != "" and \
+            pre.contains_image_markdown(first_sentence):
+            pong = ' '.join(pong.split(" ")[1:])
+            pingpong.pong = pong
+
+    lws = CtxLastWindowStrategy(win_size)
+
     prompt = lws(dummy_ppm)
     return prompt
 

chats/baize.py

@@ -9,17 +9,33 @@
 
 def build_prompts(ppmanager, user_message, global_context, win_size=3):
     dummy_ppm = copy.deepcopy(ppmanager)
+
     dummy_ppm.ctx = global_context
+    for pingpong in dummy_ppm.pingpongs:
+        pong = pingpong.pong
+        first_sentence = pong.split(" ")[0]
+        if first_sentence != "" and \
+            pre.contains_image_markdown(first_sentence):
+            pong = ' '.join(pong.split(" ")[1:])
+            pingpong.pong = pong
+
     lws = CtxLastWindowStrategy(win_size)
 
     prompt = lws(dummy_ppm)
     return prompt
 
 def text_stream(ppmanager, streamer):
+    count = 0
+
     for new_text in streamer:
         if "[|Human|]" in new_text or \
             "[|AI|]" in new_text:
             break
+
+        if count == 0:
+            ppmanager.append_pong(f"![]({global_vars.model_thumbnail_tiny})***[{global_vars.model_type}]:*** ")
+            count = count + 1
+
         ppmanager.append_pong(new_text)
         yield ppmanager, ppmanager.build_uis()
 

chats/falcon.py

@@ -20,14 +20,29 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwa
 
 def build_prompts(ppmanager, user_message, global_context, win_size=3):
     dummy_ppm = copy.deepcopy(ppmanager)
+
     dummy_ppm.ctx = global_context
+    for pingpong in dummy_ppm.pingpongs:
+        pong = pingpong.pong
+        first_sentence = pong.split(" ")[0]
+        if first_sentence != "" and \
+            pre.contains_image_markdown(first_sentence):
+            pong = ' '.join(pong.split(" ")[1:])
+            pingpong.pong = pong
+
     lws = CtxLastWindowStrategy(win_size)
 
     prompt = lws(dummy_ppm)
     return prompt
 
 def text_stream(ppmanager, streamer):
+    count = 0
+
     for new_text in streamer:
+        if count == 0:
+            ppmanager.append_pong(f"![]({global_vars.model_thumbnail_tiny})***[{global_vars.model_type}]:*** ")
+            count = count + 1
+
         ppmanager.append_pong(new_text)
         yield ppmanager, ppmanager.build_uis()
 

chats/flan_alpaca.py

@@ -9,15 +9,29 @@
 
 def build_prompts(ppmanager, user_message, global_context, win_size=3):
     dummy_ppm = copy.deepcopy(ppmanager)
-    lws = CtxLastWindowStrategy(win_size)
-
+
     dummy_ppm.ctx = global_context
-
-    prompt = lws(dummy_ppm)  
+    for pingpong in dummy_ppm.pingpongs:
+        pong = pingpong.pong
+        first_sentence = pong.split(" ")[0]
+        if first_sentence != "" and \
+            pre.contains_image_markdown(first_sentence):
+            pong = ' '.join(pong.split(" ")[1:])
+            pingpong.pong = pong
+
+    lws = CtxLastWindowStrategy(win_size)
+
+    prompt = lws(dummy_ppm)
     return prompt
 
 def text_stream(ppmanager, streamer):
+    count = 0
+
     for new_text in streamer:
+        if count == 0:
+            ppmanager.append_pong(f"![]({global_vars.model_thumbnail_tiny})***[{global_vars.model_type}]:*** ")
+            count = count + 1
+
         ppmanager.append_pong(new_text)
         yield ppmanager, ppmanager.build_uis()
 

chats/guanaco.py

@@ -21,15 +21,29 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwa
 
 def build_prompts(ppmanager, user_message, global_context, win_size=3):
     dummy_ppm = copy.deepcopy(ppmanager)
-    lws = CtxLastWindowStrategy(win_size)
 
     dummy_ppm.ctx = global_context
-
-    prompt = lws(dummy_ppm)  
+    for pingpong in dummy_ppm.pingpongs:
+        pong = pingpong.pong
+        first_sentence = pong.split(" ")[0]
+        if first_sentence != "" and \
+            pre.contains_image_markdown(first_sentence):
+            pong = ' '.join(pong.split(" ")[1:])
+            pingpong.pong = pong
+
+    lws = CtxLastWindowStrategy(win_size)
+
+    prompt = lws(dummy_ppm)
     return prompt
 
 def text_stream(ppmanager, streamer):
+    count = 0
+
     for new_text in streamer:
+        if count == 0:
+            ppmanager.append_pong(f"![]({global_vars.model_thumbnail_tiny})***[{global_vars.model_type}]:*** ")
+            count = count + 1
+
         ppmanager.append_pong(new_text)
         yield ppmanager, ppmanager.build_uis()
 

chats/koalpaca.py

@@ -9,15 +9,29 @@
 
 def build_prompts(ppmanager, user_message, global_context, win_size=3):
     dummy_ppm = copy.deepcopy(ppmanager)
-    lws = CtxLastWindowStrategy(win_size)
 
     dummy_ppm.ctx = global_context
-
-    prompt = lws(dummy_ppm)  
+    for pingpong in dummy_ppm.pingpongs:
+        pong = pingpong.pong
+        first_sentence = pong.split(" ")[0]
+        if first_sentence != "" and \
+            pre.contains_image_markdown(first_sentence):
+            pong = ' '.join(pong.split(" ")[1:])
+            pingpong.pong = pong
+
+    lws = CtxLastWindowStrategy(win_size)
+
+    prompt = lws(dummy_ppm)
     return prompt
 
 def text_stream(ppmanager, streamer):
+    count = 0
+
     for new_text in streamer:
+        if count == 0:
+            ppmanager.append_pong(f"![]({global_vars.model_thumbnail_tiny})***[{global_vars.model_type}]:*** ")
+            count = count + 1
+
         ppmanager.append_pong(new_text)
         yield ppmanager, ppmanager.build_uis()
 

chats/mpt.py

@@ -26,15 +26,29 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwa
 
 def build_prompts(ppmanager, user_message, global_context, win_size=3):
     dummy_ppm = copy.deepcopy(ppmanager)
-    lws = CtxLastWindowStrategy(win_size)
 
     dummy_ppm.ctx = global_context
-
-    prompt = lws(dummy_ppm)  
+    for pingpong in dummy_ppm.pingpongs:
+        pong = pingpong.pong
+        first_sentence = pong.split(" ")[0]
+        if first_sentence != "" and \
+            pre.contains_image_markdown(first_sentence):
+            pong = ' '.join(pong.split(" ")[1:])
+            pingpong.pong = pong
+
+    lws = CtxLastWindowStrategy(win_size)
+
+    prompt = lws(dummy_ppm)
     return prompt
 
 def text_stream(ppmanager, streamer):
+    count = 0
+
     for new_text in streamer:
+        if count == 0:
+            ppmanager.append_pong(f"![]({global_vars.model_thumbnail_tiny})***[{global_vars.model_type}]:*** ")
+            count = count + 1
+
         ppmanager.append_pong(new_text)
         yield ppmanager, ppmanager.build_uis()
 

chats/os_stablelm.py

@@ -28,7 +28,13 @@ def build_prompts(ppmanager, user_message, global_context, win_size=3):
     return prompt
 
 def text_stream(ppmanager, streamer):
+    count = 0
+
     for new_text in streamer:
+        if count == 0:
+            ppmanager.append_pong(f"![]({global_vars.model_thumbnail_tiny})***[{global_vars.model_type}]:*** ")
+            count = count + 1
+
         ppmanager.append_pong(new_text)
         yield ppmanager, ppmanager.build_uis()
 

chats/pre.py

@@ -1,9 +1,15 @@
+import re
 import copy
 import global_vars
 from threading import Thread
 from transformers import TextIteratorStreamer
 from transformers import GenerationConfig
 
+def contains_image_markdown(string):
+    regex = re.compile(r'!\[(.*?)\]\((.*?)\)')
+    match = regex.search(string)
+    return match
+
 def build_model_inputs(prompt, return_token_type_ids):
     model_inputs = global_vars.tokenizer(
         [prompt],