[temp] Align config with parsers and prompt schemas

lastmile-ai · Jan 9, 2024 · 78c32fa · 78c32fa
1 parent 389e121
commit 78c32fa
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 18 deletions.
diff --git a/cookbooks/Gradio/hf_model_parsers.py b/cookbooks/Gradio/hf_model_parsers.py
@@ -11,8 +11,8 @@
 def register_model_parsers() -> None:
     # NOTE: Commented out models need a new version of the aiconfig_extension_hugging_face package
 
-    # AIConfigRuntime.register_model_parser(HuggingFaceTextSummarizationTransformer(), "TextSummarizationTask")
-    AIConfigRuntime.register_model_parser(HuggingFaceText2ImageDiffusor(), "Text2ImageTask")
-    # AIConfigRuntime.register_model_parser(HuggingFaceText2SpeechTransformer(), "Text2SpeechTask")
-    AIConfigRuntime.register_model_parser(HuggingFaceTextGenerationTransformer(), "TextGenerationTask")
-    # AIConfigRuntime.register_model_parser(HuggingFaceTextTranslationTransformer(), "TranslationTask")
+    # AIConfigRuntime.register_model_parser(HuggingFaceTextSummarizationTransformer(), "TextSummarization")
+    AIConfigRuntime.register_model_parser(HuggingFaceText2ImageDiffusor(), "Text2Image")
+    # AIConfigRuntime.register_model_parser(HuggingFaceText2SpeechTransformer(), "Text2Speech")
+    AIConfigRuntime.register_model_parser(HuggingFaceTextGenerationTransformer(), "TextGeneration")
+    # AIConfigRuntime.register_model_parser(HuggingFaceTextTranslationTransformer(), "Translation")
diff --git a/cookbooks/Gradio/huggingface.aiconfig.json b/cookbooks/Gradio/huggingface.aiconfig.json
@@ -4,19 +4,16 @@
   "metadata": {
     "parameters": {},
     "models": {
-      "stevhliu/my_awesome_billsum_model": {
+      "TextGeneration": {
         "model": "stevhliu/my_awesome_billsum_model",
         "min_length": 10,
         "max_length": 30
       },
-      "Salesforce/blip-image-captioning-base": {
+      "ImageToText": {
         "model": "Salesforce/blip-image-captioning-base"
       }
     },
-    "model_parsers": {
-      "suno/bark": "HuggingFaceText2SpeechTransformer"
-    },
-    "default_model": "stevhliu/my_awesome_billsum_model"
+    "default_model": "TextGeneration"
   },
   "description": "The Tale of the Quick Brown Fox",
   "prompts": [
@@ -25,7 +22,12 @@
       "input": "Tell the tale of {{topic}}",
       "outputs": [],
       "metadata": {
-        "model": "translation_en_to_fr",
+        "model": {
+          "name": "Translation",
+          "settings": {
+            "model": "translation_en_to_fr"
+          }
+        },
         "parameters": {
           "topic": "the quick brown fox"
         }
@@ -36,16 +38,23 @@
       "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.",
       "outputs": [],
       "metadata": {
-        "model": "stevhliu/my_awesome_billsum_model"
+        "model": {
+          "name": "TextGeneration",
+          "settings": {
+            "model": "stevhliu/my_awesome_billsum_model"
+          }
+        }
       }
     },
     {
       "name": "generate_audio_title",
       "input": "The Quick Brown Fox was admired by all the animals in the forest.",
       "metadata": {
         "model": {
-          "name": "suno/bark",
-          "settings": {}
+          "name": "Text2Speech",
+          "settings": {
+            "model": "suno/bark"
+          }
         }
       }
     },
@@ -60,7 +69,12 @@
         ]
       },
       "metadata": {
-        "model": "Salesforce/blip-image-captioning-base"
+        "model": {
+          "name": "ImageToText",
+          "settings": {
+            "model": "Salesforce/blip-image-captioning-base"
+          }
+        }
       }
     },
     {
@@ -84,9 +98,14 @@
         ]
       },
       "metadata": {
-        "model": "openai/whisper-small"
+        "model": {
+          "name": "AutomaticSpeechRecognition",
+          "settings": {
+            "model": "openai/whisper-small"
+          }
+        }
       }
     }
   ],
   "$schema": "https://json.schemastore.org/aiconfig-1.0"
-}
+}
diff --git a/python/src/aiconfig/editor/client/src/utils/promptUtils.ts b/python/src/aiconfig/editor/client/src/utils/promptUtils.ts
@@ -7,6 +7,7 @@ import { PaLMChatParserPromptSchema } from "../shared/prompt_schemas/PaLMChatPar
 import { HuggingFaceTextGenerationParserPromptSchema } from "../shared/prompt_schemas/HuggingFaceTextGenerationParserPromptSchema";
 import { AnyscaleEndpointPromptSchema } from "../shared/prompt_schemas/AnyscaleEndpointPromptSchema";
 import { HuggingFaceText2ImageDiffusorPromptSchema } from "../shared/prompt_schemas/HuggingFaceText2ImageDiffusorPromptSchema";
+import { HuggingFaceAutomaticSpeechRecognitionPromptSchema } from "../shared/prompt_schemas/HuggingFaceAutomaticSpeechRecognitionPromptSchema";
 
 /**
  * Get the name of the model for the specified prompt. The name will either be specified in the prompt's
@@ -84,6 +85,14 @@ export const PROMPT_SCHEMAS: Record<string, PromptSchema> = {
 
   // Local HuggingFace Parsers
   HuggingFaceText2ImageDiffusor: HuggingFaceText2ImageDiffusorPromptSchema,
+
+  // Temporary for HF example:
+  //TextGeneration: HuggingFaceTextGenerationTransformerPromptSchema,
+  Text2Image: HuggingFaceText2ImageDiffusorPromptSchema,
+  // Text2Speech: HuggingFaceText2SpeechTransformerPromptSchema,
+  AutomaticSpeechRecognition: HuggingFaceAutomaticSpeechRecognitionPromptSchema,
+  // HuggingFaceTextTranslationTransformer: HuggingFaceTextTranslationTransformerPromptSchema,
+  // Image2Text: HuggingFaceImage2TextTransformerPromptSchema,
 };
 
 export type PromptInputSchema =