From 5b7434483e2521110bbe3cc3380d9b99a6d4e8be Mon Sep 17 00:00:00 2001
From: "Rossdan Craig rossdan@lastmileai.dev" <>
Date: Wed, 10 Jan 2024 02:03:40 -0500
Subject: [PATCH] Testing streaming outputs

Various version of this PR are for testing streaming implementations of the HuggingFace model parsers
---
 cookbooks/Gradio/hf_model_parsers.py        | 31 ++++++++++++++++
 cookbooks/Gradio/huggingface.aiconfig.json  | 40 ++++++++++++++++-----
 python/src/aiconfig/editor/server/server.py |  4 +--
 3 files changed, 65 insertions(+), 10 deletions(-)
 create mode 100644 cookbooks/Gradio/hf_model_parsers.py

diff --git a/cookbooks/Gradio/hf_model_parsers.py b/cookbooks/Gradio/hf_model_parsers.py
new file mode 100644
index 000000000..17a829baf
--- /dev/null
+++ b/cookbooks/Gradio/hf_model_parsers.py
@@ -0,0 +1,31 @@
+from aiconfig_extension_hugging_face import (
+    HuggingFaceAutomaticSpeechRecognitionTransformer,
+    HuggingFaceImage2TextTransformer,
+    HuggingFaceTextSummarizationTransformer,
+    HuggingFaceText2ImageDiffusor,
+    HuggingFaceText2SpeechTransformer,
+    HuggingFaceTextGenerationTransformer,
+    HuggingFaceTextTranslationTransformer,
+)
+from aiconfig import (AIConfigRuntime, ModelParserRegistry)
+
+def register_model_parsers() -> None:
+    """Register model parsers for HuggingFace models.
+    """
+    # Audio --> Text
+    # AIConfigRuntime.register_model_parser(HuggingFaceAutomaticSpeechRecognitionTransformer(), "AutomaticSpeechRecognition")
+
+    # # Image --> Text
+    # AIConfigRuntime.register_model_parser(HuggingFaceImage2TextTransformer(), "Image2Text")
+
+    # # Text --> Image
+    # AIConfigRuntime.register_model_parser(HuggingFaceText2ImageDiffusor(), "Text2Image")
+
+    # # Text --> Audio
+    # AIConfigRuntime.register_model_parser(HuggingFaceText2SpeechTransformer(), "Text2Speech")
+
+    # # Text --> Text
+    # AIConfigRuntime.register_model_parser(HuggingFaceTextGenerationTransformer(), "TextGeneration")
+    # AIConfigRuntime.register_model_parser(HuggingFaceTextSummarizationTransformer(), "TextSummarization")
+    ModelParserRegistry.register_model_parser(HuggingFaceTextTranslationTransformer())
+    # AIConfigRuntime.register_model_parser(HuggingFaceTextTranslationTransformer(), "TextTranslation")
diff --git a/cookbooks/Gradio/huggingface.aiconfig.json b/cookbooks/Gradio/huggingface.aiconfig.json
index 2ab54cc38..3dc5895ba 100644
--- a/cookbooks/Gradio/huggingface.aiconfig.json
+++ b/cookbooks/Gradio/huggingface.aiconfig.json
@@ -1,20 +1,44 @@
 {
-  "name": "",
+  "name": "The Tale of the Quick Brown Fox",
   "schema_version": "latest",
   "metadata": {
     "parameters": {},
-    "models": {}
+    "models": {
+      "Helsinki-NLP/opus-mt-en-fr": {
+        "model": "Helsinki-NLP/opus-mt-en-fr"
+      }
+    },
+    "default_model": "TextGeneration",
+    "model_parsers": {
+      "Helsinki-NLP/opus-mt-en-fr": "HuggingFaceTextTranslationTransformer"
+    }
   },
-  "description": "",
+  "description": "The Tale of the Quick Brown Fox",
   "prompts": [
     {
-      "name": "prompt_1",
-      "input": "",
+      "name": "Translate a story from English to French",
+      "input": "Once upon a time, in a lush and vibrant forest, there lived a magnificent creature known as the Quick Brown Fox. This fox was unlike any other, possessing incredible speed and agility that awed all the animals in the forest. With its fur as golden as the sun and its eyes as sharp as emeralds, the Quick Brown Fox was admired by everyone, from the tiniest hummingbird to the mightiest bear. The fox had a kind heart and would often lend a helping paw to those in need. The Quick Brown Fox had a particular fondness for games and challenges. It loved to test its skills against others, always seeking new adventures to satisfy its boundless curiosity. Its favorite game was called \"The Great Word Hunt,\" where it would embark on a quest to find hidden words scattered across the forest. \n\nOne day it got very old and died",
       "metadata": {
-        "model": "gpt-4",
-        "parameters": {}
+        "model": {
+          "name": "Helsinki-NLP/opus-mt-en-fr",
+          "settings": {
+            "max_length": 100,
+            "min_length": 50,
+            "num_beams": 1
+          }
+        },
+        "parameters": {
+          "city": "New York"
+        }
       },
-      "outputs": []
+      "outputs": [
+        {
+          "output_type": "execute_result",
+          "execution_count": 0,
+          "data": " Il \u00e9tait une fois, dans une for\u00eat luxuriante et vibrante, vivait une magnifique cr\u00e9ature connue sous le nom de Foxe Brun. Ce renard \u00e9tait diff\u00e9rent des autres, poss\u00e9dant une vitesse et une agilit\u00e9 incroyables qui amenaient tous les animaux dans la for\u00eat. Avec sa fourrure aussi dor\u00e9e que le soleil et ses yeux aussi aiguis\u00e9s que les \u00e9meraudes, le Foxe Brun Brun \u00e9tait admir\u00e9 par tout le monde, du plus petit colibri au",
+          "metadata": {}
+        }
+      ]
     }
   ],
   "$schema": "https://json.schemastore.org/aiconfig-1.0"
diff --git a/python/src/aiconfig/editor/server/server.py b/python/src/aiconfig/editor/server/server.py
index 9c257804a..3d54ece4a 100644
--- a/python/src/aiconfig/editor/server/server.py
+++ b/python/src/aiconfig/editor/server/server.py
@@ -293,9 +293,9 @@ def kill_thread(thread_id: int | None):
 
             # Yea I know time.sleep() isn't super accurate, but it's fine,
             # we can fix later
-            time.sleep(0.1)
+            # time.sleep(0.1)
             wait_time_in_seconds += SLEEP_DELAY_SECONDS
-            print(f"Output queue is currently empty. Waiting for {wait_time_in_seconds:.1f}s...")
+            # print(f"Output queue is currently empty. Waiting for {wait_time_in_seconds:.1f}s...")
 
         # Yield in flask is weird and you either need to send responses as a
         # string, or artificially wrap them around "[" and "]"