state

app.py

@@ -4,103 +4,109 @@ import re
 from diffusers.utils import load_image
 import requests
 from awesome_chat import chat_huggingface
-from awesome_chat import set_huggingface_token, get_huggingface_token
 import os
 
-all_messages = []
-OPENAI_KEY = ""
-
 os.makedirs("public/images", exist_ok=True)
 os.makedirs("public/audios", exist_ok=True)
 os.makedirs("public/videos", exist_ok=True)
 
-def add_message(content, role):
-    message = {"role":role, "content":content}
-    all_messages.append(message)
-
-def extract_medias(message):
-    image_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(jpg|jpeg|tiff|gif|png)")
-    image_urls = []
-    for match in image_pattern.finditer(message):
-        if match.group(0) not in image_urls:
-            image_urls.append(match.group(0))
-
-    audio_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(flac|wav)")
-    audio_urls = []
-    for match in audio_pattern.finditer(message):
-        if match.group(0) not in audio_urls:
-            audio_urls.append(match.group(0))
-
-    video_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(mp4)")
-    video_urls = []
-    for match in video_pattern.finditer(message):
-        if match.group(0) not in video_urls:
-            video_urls.append(match.group(0))
-
-    return image_urls, audio_urls, video_urls
-
-def set_key(openai_key):
-    global OPENAI_KEY
-    OPENAI_KEY = openai_key
-    return OPENAI_KEY
-
-def set_token(huggingface_token):
-    set_huggingface_token(huggingface_token)
-    return huggingface_token
-
-def add_text(messages, message):
-    if len(OPENAI_KEY) == 0 or not OPENAI_KEY.startswith("sk-"):
-        return messages, "Please set your OpenAI API key or Hugging Face token first!!!"
-    add_message(message, "user")
-    messages = messages + [(message, None)]
-    image_urls, audio_urls, video_urls = extract_medias(message)
-
-    for image_url in image_urls:
-        if not image_url.startswith("http") and not image_url.startswith("public"):
-            image_url = "public/" + image_url
-        image = load_image(image_url)
-        name = f"public/images/{str(uuid.uuid4())[:4]}.jpg" 
-        image.save(name)
-        messages = messages + [((f"{name}",), None)]
-    for audio_url in audio_urls and not audio_url.startswith("public"):
-        if not audio_url.startswith("http"):
-            audio_url = "public/" + audio_url
-        ext = audio_url.split(".")[-1]
-        name = f"public/audios/{str(uuid.uuid4()[:4])}.{ext}"
-        response = requests.get(audio_url)
-        with open(name, "wb") as f:
-            f.write(response.content)
-        messages = messages + [((f"{name}",), None)]
-    for video_url in video_urls and not video_url.startswith("public"):
-        if not video_url.startswith("http"):
-            video_url = "public/" + video_url
-        ext = video_url.split(".")[-1]
-        name = f"public/audios/{str(uuid.uuid4()[:4])}.{ext}"
-        response = requests.get(video_url)
-        with open(name, "wb") as f:
-            f.write(response.content)
-        messages = messages + [((f"{name}",), None)]
-    return messages, ""
-
-def bot(messages):
-    if len(OPENAI_KEY) == 0 or not OPENAI_KEY.startswith("sk-"):
+class Client:
+    def __init__(self) -> None:
+        self.OPENAI_KEY = ""
+        self.HUGGINGFACE_TOKEN = ""
+        self.all_messages = []
+
+    def set_key(self, openai_key):
+        self.OPENAI_KEY = openai_key
+        if len(self.HUGGINGFACE_TOKEN)>0:
+            gr.update(visible = True)
+        return self.OPENAI_KEY
+
+    def set_token(self, huggingface_token):
+        self.HUGGINGFACE_TOKEN = huggingface_token
+        if len(self.OPENAI_KEY)>0:
+            gr.update(visible = True)
+        return self.HUGGINGFACE_TOKEN
+    
+    def add_message(self, content, role):
+        message = {"role":role, "content":content}
+        self.all_messages.append(message)
+
+    def extract_medias(self, message):
+        image_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(jpg|jpeg|tiff|gif|png)")
+        image_urls = []
+        for match in image_pattern.finditer(message):
+            if match.group(0) not in image_urls:
+                image_urls.append(match.group(0))
+
+        audio_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(flac|wav)")
+        audio_urls = []
+        for match in audio_pattern.finditer(message):
+            if match.group(0) not in audio_urls:
+                audio_urls.append(match.group(0))
+
+        video_pattern = re.compile(r"(http(s?):|\/)?([\.\/_\w:-])*?\.(mp4)")
+        video_urls = []
+        for match in video_pattern.finditer(message):
+            if match.group(0) not in video_urls:
+                video_urls.append(match.group(0))
+
+        return image_urls, audio_urls, video_urls
+
+    def add_text(self, messages, message):
+        if len(self.OPENAI_KEY) == 0 or not self.OPENAI_KEY.startswith("sk-") or len(self.HUGGINGFACE_TOKEN) == 0 or not self.HUGGINGFACE_TOKEN.startswith("hf_"):
+            return messages, "Please set your OpenAI API key or Hugging Face token first!!!"
+        self.add_message(message, "user")
+        messages = messages + [(message, None)]
+        image_urls, audio_urls, video_urls = self.extract_medias(message)
+
+        for image_url in image_urls:
+            if not image_url.startswith("http") and not image_url.startswith("public"):
+                image_url = "public/" + image_url
+            image = load_image(image_url)
+            name = f"public/images/{str(uuid.uuid4())[:4]}.jpg" 
+            image.save(name)
+            messages = messages + [((f"{name}",), None)]
+        for audio_url in audio_urls and not audio_url.startswith("public"):
+            if not audio_url.startswith("http"):
+                audio_url = "public/" + audio_url
+            ext = audio_url.split(".")[-1]
+            name = f"public/audios/{str(uuid.uuid4()[:4])}.{ext}"
+            response = requests.get(audio_url)
+            with open(name, "wb") as f:
+                f.write(response.content)
+            messages = messages + [((f"{name}",), None)]
+        for video_url in video_urls and not video_url.startswith("public"):
+            if not video_url.startswith("http"):
+                video_url = "public/" + video_url
+            ext = video_url.split(".")[-1]
+            name = f"public/audios/{str(uuid.uuid4()[:4])}.{ext}"
+            response = requests.get(video_url)
+            with open(name, "wb") as f:
+                f.write(response.content)
+            messages = messages + [((f"{name}",), None)]
+        return messages, ""
+
+    def bot(self, messages):
+        if len(self.OPENAI_KEY) == 0 or not self.OPENAI_KEY.startswith("sk-") or len(self.HUGGINGFACE_TOKEN) == 0 or not self.HUGGINGFACE_TOKEN.startswith("hf_"):
+            return messages
+        message = chat_huggingface(self.all_messages, self.OPENAI_KEY, self.HUGGINGFACE_TOKEN)["message"]
+        image_urls, audio_urls, video_urls = self.extract_medias(message)
+        self.add_message(message, "assistant")
+        messages[-1][1] = message
+        for image_url in image_urls:
+            image_url = image_url.replace("public/", "")
+            messages = messages + [((None, (f"public/{image_url}",)))]
+        for audio_url in audio_urls:
+            audio_url = audio_url.replace("public/", "")
+            messages = messages + [((None, (f"public/{audio_url}",)))]
+        for video_url in video_urls:
+            video_url = video_url.replace("public/", "")
+            messages = messages + [((None, (f"public/{video_url}",)))]
         return messages
-    message = chat_huggingface(all_messages, OPENAI_KEY)["message"]
-    image_urls, audio_urls, video_urls = extract_medias(message)
-    add_message(message, "assistant")
-    messages[-1][1] = message
-    for image_url in image_urls:
-        image_url = image_url.replace("public/", "")
-        messages = messages + [((None, (f"public/{image_url}",)))]
-    for audio_url in audio_urls:
-        audio_url = audio_url.replace("public/", "")
-        messages = messages + [((None, (f"public/{audio_url}",)))]
-    for video_url in video_urls:
-        video_url = video_url.replace("public/", "")
-        messages = messages + [((None, (f"public/{video_url}",)))]
-    return messages
 
 with gr.Blocks() as demo:
+    state = gr.State(value={"client": Client()})
     gr.Markdown("<h1><center>HuggingGPT</center></h1>")
     gr.Markdown("<p align='center'><img src='https://i.ibb.co/qNH3Jym/logo.png' height='25' width='95'></p>")
     gr.Markdown("<p align='center' style='font-size: 20px;'>A system to connect LLMs with ML community. See our <a href='https://github.com/microsoft/JARVIS'>Project</a> and <a href='http://arxiv.org/abs/2303.17580'>Paper</a>.</p>")
@@ -135,13 +141,25 @@ with gr.Blocks() as demo:
             ).style(container=False)
         with gr.Column(scale=0.15, min_width=0):
             btn2 = gr.Button("Send").style(full_height=True)
-
-    openai_api_key.submit(set_key, [openai_api_key], [openai_api_key])
-    txt.submit(add_text, [chatbot, txt], [chatbot, txt]).then(bot, chatbot, chatbot)
+        
+    def set_key(state, openai_api_key):
+        return state["client"].set_key(openai_api_key)
+
+    def add_text(state, chatbot, txt):
+        return state["client"].add_text(chatbot, txt)
+    
+    def set_token(state, hugging_face_token):
+        return state["client"].set_token(hugging_face_token)
+    
+    def bot(state, chatbot):
+        return state["client"].bot(chatbot)
+
+    openai_api_key.submit(set_key, [state, openai_api_key], [openai_api_key])
+    txt.submit(add_text, [state, chatbot, txt], [chatbot, txt]).then(bot, [state, chatbot], chatbot)
     hugging_face_token.submit(set_token, [hugging_face_token], [hugging_face_token])
-    btn1.click(set_key, [openai_api_key], [openai_api_key])
-    btn2.click(add_text, [chatbot, txt], [chatbot, txt]).then(bot, chatbot, chatbot)
-    btn3.click(set_token, [hugging_face_token], [hugging_face_token])
+    btn1.click(set_key, [state, openai_api_key], [openai_api_key])
+    btn2.click(add_text, [state, chatbot, txt], [chatbot, txt]).then(bot, [state, chatbot], chatbot)
+    btn3.click(set_token, [state, hugging_face_token], [hugging_face_token])
 
     gr.Examples(
         examples=["Given a collection of image A: /examples/a.jpg, B: /examples/b.jpg, C: /examples/c.jpg, please tell me how many zebras in these picture?",

awesome_chat.py

@@ -119,15 +119,6 @@ METADATAS = {}
 for model in MODELS:
     METADATAS[model["id"]] = model
 
-HUGGINGFACE_TOKEN = ""
-
-def set_huggingface_token(token):
-    global HUGGINGFACE_TOKEN
-    HUGGINGFACE_TOKEN = token
-
-def get_huggingface_token():
-    return HUGGINGFACE_TOKEN
-
 def convert_chat_to_completion(data):
     messages = data.pop('messages', [])
     tprompt = ""
@@ -343,12 +334,15 @@ def response_results(input, results, openaikey=None):
     }
     return send_request(data)
 
-def huggingface_model_inference(model_id, data, task):
-    HUGGINGFACE_HEADERS = {
-        "Authorization": f"Bearer {HUGGINGFACE_TOKEN}",
+def huggingface_model_inference(model_id, data, task, huggingfacetoken=None):
+    if huggingfacetoken is None:
+        HUGGINGFACE_HEADERS = {}
+    else:
+        HUGGINGFACE_HEADERS = {
+            "Authorization": f"Bearer {huggingfacetoken}",
     }
     task_url = f"https://api-inference.huggingface.co/models/{model_id}" # InferenceApi does not yet support some tasks
-    inference = InferenceApi(repo_id=model_id, token=HUGGINGFACE_TOKEN)
+    inference = InferenceApi(repo_id=model_id, token=huggingfacetoken)
     
     # NLP tasks
     if task == "question-answering":
@@ -573,10 +567,13 @@ def local_model_inference(model_id, data, task):
         return results
 
 
-def model_inference(model_id, data, hosted_on, task):
-    HUGGINGFACE_HEADERS = {
-        "Authorization": f"Bearer {HUGGINGFACE_TOKEN}",
-    }
+def model_inference(model_id, data, hosted_on, task, huggingfacetoken=None):
+    if huggingfacetoken:
+        HUGGINGFACE_HEADERS = {
+            "Authorization": f"Bearer {huggingfacetoken}",
+        }
+    else:
+        HUGGINGFACE_HEADERS = None
     if hosted_on == "unknown":
         r = status(model_id)
         logger.debug("Local Server Status: " + str(r.json()))
@@ -592,7 +589,7 @@ def model_inference(model_id, data, hosted_on, task):
         if hosted_on == "local":
             inference_result = local_model_inference(model_id, data, task)
         elif hosted_on == "huggingface":
-            inference_result = huggingface_model_inference(model_id, data, task)
+            inference_result = huggingface_model_inference(model_id, data, task, huggingfacetoken)
     except Exception as e:
         print(e)
         traceback.print_exc()
@@ -615,12 +612,12 @@ def get_model_status(model_id, url, headers, queue = None):
             queue.put((model_id, False, None))
         return False
 
-def get_avaliable_models(candidates, topk=10):
+def get_avaliable_models(candidates, topk=10, huggingfacetoken = None):
     all_available_models = {"local": [], "huggingface": []}
     threads = []
     result_queue = Queue()
     HUGGINGFACE_HEADERS = {
-        "Authorization": f"Bearer {HUGGINGFACE_TOKEN}",
+        "Authorization": f"Bearer {huggingfacetoken}",
     }
     for candidate in candidates:
         model_id = candidate["id"]
@@ -658,7 +655,7 @@ def collect_result(command, choose, inference_result):
     return result
 
 
-def run_task(input, command, results, openaikey = None):
+def run_task(input, command, results, openaikey = None, huggingfacetoken = None):
     id = command["id"]
     args = command["args"]
     task = command["task"]
@@ -769,11 +766,11 @@ def run_task(input, command, results, openaikey = None):
             logger.warning(f"no available models on {task} task.")
             record_case(success=False, **{"input": input, "task": command, "reason": f"task not support: {command['task']}", "op":"message"})
             inference_result = {"error": f"{command['task']} not found in available tasks."}
-            results[id] = collect_result(command, choose, inference_result)
+            results[id] = collect_result(command, "", inference_result)
             return False
 
         candidates = MODELS_MAP[task][:20]
-        all_avaliable_models = get_avaliable_models(candidates, config["num_candidate_models"])
+        all_avaliable_models = get_avaliable_models(candidates, config["num_candidate_models"], huggingfacetoken)
         all_avaliable_model_ids = all_avaliable_models["local"] + all_avaliable_models["huggingface"]
         logger.debug(f"avaliable models on {command['task']}: {all_avaliable_models}")
 
@@ -818,7 +815,7 @@ def run_task(input, command, results, openaikey = None):
                 choose_str = find_json(choose_str)
                 best_model_id, reason, choose  = get_id_reason(choose_str)
                 hosted_on = "local" if best_model_id in all_avaliable_models["local"] else "huggingface"
-    inference_result = model_inference(best_model_id, args, hosted_on, command['task'])
+    inference_result = model_inference(best_model_id, args, hosted_on, command['task'], huggingfacetoken)
 
     if "error" in inference_result:
         logger.warning(f"Inference error: {inference_result['error']}")
@@ -829,7 +826,7 @@ def run_task(input, command, results, openaikey = None):
     results[id] = collect_result(command, choose, inference_result)
     return True
 
-def chat_huggingface(messages, openaikey = None, return_planning = False, return_results = False):
+def chat_huggingface(messages, openaikey = None, huggingfacetoken = None, return_planning = False, return_results = False):
     start = time.time()
     context = messages[:-1]
     input = messages[-1]["content"]
@@ -871,7 +868,7 @@ def chat_huggingface(messages, openaikey = None, return_planning = False, return
             # logger.debug(f"d.keys(): {d.keys()}, dep: {dep}")
             if len(list(set(dep).intersection(d.keys()))) == len(dep) or dep[0] == -1:
                 tasks.remove(task)
-                thread = threading.Thread(target=run_task, args=(input, task, d, openaikey))
+                thread = threading.Thread(target=run_task, args=(input, task, d, openaikey, huggingfacetoken))
                 thread.start()
                 threads.append(thread)
         if num_threads == len(threads):

config.gradio.yaml

@@ -3,7 +3,7 @@ openai:
 huggingface:
   token: # required: huggingface token @ https://huggingface.co/settings/tokens
 dev: false
-debug: false
+debug: true
 log_file: logs/debug.log
 model: text-davinci-003 # text-davinci-003
 use_completion: true
@@ -13,7 +13,7 @@ num_candidate_models: 5
 max_description_length: 100
 proxy: 
 logit_bias:
-  parse_task: 0.1
+  parse_task: 0.5
   choose_model: 5
 tprompt:
   parse_task: >-