Upload llmdolphin.py

llmdolphin.py

@@ -6,6 +6,7 @@ from llama_cpp_agent.providers import LlamaCppPythonProvider
 from llama_cpp_agent.chat_history import BasicChatHistory
 from llama_cpp_agent.chat_history.messages import Roles
 from ja_to_danbooru.ja_to_danbooru import jatags_to_danbooru_tags
+import wrapt_timeout_decorator
 
 
 llm_models_dir = "./llm_models"
@@ -127,6 +128,7 @@ def list_uniq(l):
     return sorted(set(l), key=l.index)
 
 
+@wrapt_timeout_decorator.timeout(dec_timeout=3.5)
 def to_list_ja(s):
     import re
     s = re.sub(r'[、。]', ',', s)
@@ -359,6 +361,7 @@ def select_dolphin_language(lang: str):
     return gr.update(value=get_dolphin_sysprompt())
 
 
+@wrapt_timeout_decorator.timeout(dec_timeout=5.0)
 def get_raw_prompt(msg: str):
     import re
     m = re.findall(r'/GENBEGIN/(.+?)/GENEND/', msg, re.DOTALL)
@@ -443,12 +446,13 @@ def dolphin_respond(
 def dolphin_parse(
     history: list[tuple[str, str]],
 ):
-    if dolphin_sysprompt_mode == "Chat with LLM" or not history or len(history) < 1: "", gr.update(visible=True), gr.update(visible=True)
+    if dolphin_sysprompt_mode == "Chat with LLM" or not history or len(history) < 1:
+        return "", gr.update(visible=True), gr.update(visible=True)
     try:
         msg = history[-1][0]
+        raw_prompt = get_raw_prompt(msg)
     except Exception:
-        return ""
-    raw_prompt = get_raw_prompt(msg)
+        return "", gr.update(visible=True), gr.update(visible=True)
     prompts = []
     if dolphin_sysprompt_mode == "Japanese to Danbooru Dictionary" and is_japanese(raw_prompt):
         prompts = list_uniq(jatags_to_danbooru_tags(to_list_ja(raw_prompt)) + ["nsfw", "explicit"])
@@ -471,7 +475,6 @@ def dolphin_respond_auto(
     progress=gr.Progress(track_tqdm=True),
 ):
     #if not is_japanese(message): return [(None, None)]
-
     from pathlib import Path
     progress(0, desc="Processing...")
 
@@ -543,9 +546,9 @@ def dolphin_parse_simple(
     if dolphin_sysprompt_mode == "Chat with LLM" or not history or len(history) < 1: return message
     try:
         msg = history[-1][0]
+        raw_prompt = get_raw_prompt(msg)
     except Exception:
         return ""
-    raw_prompt = get_raw_prompt(msg)
     prompts = []
     if dolphin_sysprompt_mode == "Japanese to Danbooru Dictionary" and is_japanese(raw_prompt):
         prompts = list_uniq(jatags_to_danbooru_tags(to_list_ja(raw_prompt)) + ["nsfw", "explicit"])