Ryan Holinshead commited on
Commit
5206d2d
1 Parent(s): 352d8fc

Add local model parsers

Browse files
Files changed (1) hide show
  1. model_parsers.py +57 -0
model_parsers.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from aiconfig_extension_hugging_face import (
2
+ HuggingFaceAutomaticSpeechRecognitionTransformer,
3
+ HuggingFaceImage2TextTransformer,
4
+ HuggingFaceText2ImageDiffusor,
5
+ HuggingFaceText2SpeechTransformer,
6
+ HuggingFaceTextGenerationTransformer,
7
+ HuggingFaceTextSummarizationTransformer,
8
+ HuggingFaceTextTranslationTransformer,
9
+ )
10
+
11
+ from aiconfig import AIConfigRuntime, ModelParserRegistry
12
+
13
+
14
+ # Example of how users can register model parsers for use in the GradioWorkbook
15
+ # The implementation looks for a parsers_path (model_parsers.py by default) which
16
+ # should include a module with a register_model_parsers function.
17
+ # Here we are registering all the local HuggingFace model parsers as an example
18
+ def register_model_parsers() -> None:
19
+ """Register model parsers for local HuggingFace models."""
20
+ automatic_speech_recognition = HuggingFaceAutomaticSpeechRecognitionTransformer()
21
+ AIConfigRuntime.register_model_parser(
22
+ automatic_speech_recognition, "Automatic Speech Recognition (Local)"
23
+ )
24
+
25
+ image_to_text = HuggingFaceImage2TextTransformer()
26
+ AIConfigRuntime.register_model_parser(image_to_text, "Image-to-Text (Local)")
27
+
28
+ text_to_image = HuggingFaceText2ImageDiffusor()
29
+ AIConfigRuntime.register_model_parser(text_to_image, "Text-to-Image (Local)")
30
+
31
+ text_to_speech = HuggingFaceText2SpeechTransformer()
32
+ AIConfigRuntime.register_model_parser(text_to_speech, "Text-to-Speech (Local)")
33
+
34
+ text_generation = HuggingFaceTextGenerationTransformer()
35
+ AIConfigRuntime.register_model_parser(text_generation, "Text Generation (Local)")
36
+
37
+ text_summarization = HuggingFaceTextSummarizationTransformer()
38
+ AIConfigRuntime.register_model_parser(text_summarization, "Summarization (Local)")
39
+
40
+ text_translation = HuggingFaceTextTranslationTransformer()
41
+ AIConfigRuntime.register_model_parser(text_translation, "Translation (Local)")
42
+
43
+ # By default, model parsers will also have their own ids registered. Remove those
44
+ # since we just want the task-based names registered
45
+ parsers = [
46
+ automatic_speech_recognition,
47
+ image_to_text,
48
+ text_to_image,
49
+ text_to_speech,
50
+ text_generation,
51
+ text_summarization,
52
+ text_translation,
53
+ ]
54
+
55
+ for parser in parsers:
56
+ ModelParserRegistry.remove_model_parser(parser.id())
57
+