Merge pull request #86 from easy-model-fusion/develop

main.py

@@ -1,21 +1,25 @@
 import argparse
-from sdk.demo import DemoTextConv, DemoTxtToImg
+from sdk.demo import DemoTextConv, DemoTextToImg, DemoTextGen
 
 
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser(description='Choose a Model type :'
-                                                 ' TextConv or TxtToImg')
+                                                 ' TextConv, TextToImg '
+                                                 'or TextGen')
     subparser = parser.add_subparsers(dest='option')
 
     conv = subparser.add_parser('TextConv')
-    img = subparser.add_parser('TxtToImg')
+    img = subparser.add_parser('TextToImg')
+    gen = subparser.add_parser('TextGen')
     args = parser.parse_args()
 
     match args.option:
         case 'TextConv':
             DemoTextConv()
-        case ('TxtToImg'):
-            DemoTxtToImg()
+        case 'TextToImg':
+            DemoTextToImg()
+        case 'TextGen':
+            DemoTextGen()
         case _:
-            DemoTxtToImg()
+            DemoTextToImg()

sdk/demo/__init__.py

@@ -1,3 +1,4 @@
 # flake8: noqa F401
 from .demo_text_conv import DemoTextConv
-from .demo_txt_to_img import DemoTxtToImg
+from .demo_text_to_img import DemoTextToImg
+from .demo_text_generation import DemoTextGen

sdk/demo/demo_text_conv.py

@@ -1,7 +1,6 @@
-from sdk.models import ModelsManagement, ModelsTextConversation
-from sdk.options import Devices, OptionsTextConversation
-from sdk.options.options_tokenizer import OptionsTokenizer
-from sdk.tokenizers.tokenizer import Tokenizer
+from sdk.models import ModelsTextConversation
+from sdk.options import Devices
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
 class DemoTextConv:
@@ -13,75 +12,20 @@ def __init__(self):
         """
         Initializes the DemoTextConv class with predefined options and models.
         """
-        # Define the model name and path
-        model_name = "Salesforce/codegen-350M-nl"
-        model_path = "Salesforce/codegen-350M-nl"
+        model_path = "microsoft/phi-2"
+        tokenizer_path = "microsoft/phi-2"
 
-        # Define options for text conversation
-        options = OptionsTextConversation(
-            prompt="Hello, what's 3 + 3 ?",
-            device=Devices.GPU,
-            model_name=model_name,
-            trust_remote_code=True,
-            max_length=100
-        )
+        model = ModelsTextConversation(model_name="model",
+                                       model_path=model_path,
+                                       tokenizer_path=tokenizer_path,
+                                       model_class=AutoModelForCausalLM,
+                                       tokenizer_class=AutoTokenizer,
+                                       device=Devices.GPU)
 
-        # Define tokenizer options
-        tokenizer_options = OptionsTokenizer(
-            device='cuda',
-            padding_side='left',
-            return_tensors="pt"
-        )
+        model.load_model()
+        model.create_new_conversation()
 
-        tokenizer = Tokenizer("Salesforce/codegen-350M-nl",
-                              "Salesforce/codegen-350M-nl",
-                              "Salesforce/codegen-350M-nl",
-                              tokenizer_options)
+        result = model.generate_prompt(
+            "I'm looking for a movie - what's your favourite one?")
 
-        # Initialize the model management
-        model_management = ModelsManagement()
-
-        # Create and load the text conversation model
-        model = ModelsTextConversation(model_name, model_path, options)
-        model.tokenizer = tokenizer
-        model_management.add_model(new_model=model, model_options=options)
-        model_management.load_model(model_name)
-
-        # Print the response to the initial prompt
-        print(model_management.generate_prompt(options.prompt))
-
-        # Generate a response to a custom prompt
-        print(model_management.generate_prompt(
-            prompt="What did I say before ?"))
-
-        # Create a new conversation with a new prompt
-        options = OptionsTextConversation(
-            prompt="Hello, what's 6 + 6 ?",
-            device=options.device,
-            model_name=model_name,
-            batch_size=1,
-            chat_id_to_use=1,
-            minimum_tokens=50,
-            create_new_conv=True
-        )
-
-        # Switch to the new conversation
-        model_management.set_model_options(model_name=model_name,
-                                           options=options)
-        print(model_management.generate_prompt(options.prompt))
-
-        print(model_management.generate_prompt("Where is Japan"))
-
-        # Switch back to the initial conversation
-        options.chat_id_to_use = 0
-        # Create a new tokenizer and use it
-        options.tokenizer_id_to_use = 1
-        model_management.set_model_options(model_name=model_name,
-                                           options=options)
-        tokenizer_options = OptionsTokenizer(
-            device='cuda',
-            padding_side='right',
-            return_tensors='pt'
-        )
-        model.tokenizer_options = tokenizer_options
-        print(model_management.generate_prompt("Bye "))
+        print(result.messages[-1]["content"])

sdk/demo/demo_text_generation.py

@@ -0,0 +1,37 @@
+from sdk.options import Devices
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+from sdk.models import ModelTransformers
+
+
+class DemoTextGen:
+    """
+    This class demonstrates a text conversation using a chatbot model.
+    """
+
+    def __init__(self):
+        """
+        Initializes the DemoTextConv class with predefined options and models.
+        """
+        model_path = "microsoft/phi-2"
+        tokenizer_path = "microsoft/phi-2"
+
+        model_transformers = ModelTransformers(
+            model_name="model",
+            model_path=model_path,
+            tokenizer_path=tokenizer_path,
+            task="text-generation",
+            model_class=AutoModelForCausalLM,
+            tokenizer_class=AutoTokenizer,
+            device=Devices.GPU
+        )
+
+        model_transformers.load_model()
+
+        result = model_transformers.generate_prompt(
+            prompt="I'm looking for a movie - what's your favourite one?",
+            max_length=300,
+            truncation=True
+        )
+
+        print(result)

sdk/demo/demo_txt_to_img.py → sdk/demo/demo_text_to_img.py

@@ -1,27 +1,29 @@
+import torch
 from sdk.models import ModelTextToImage, ModelsManagement
-from sdk.options import Devices, OptionsTextToImage
+from sdk.options import Devices
 
 
-class DemoTxtToImg:
+class DemoTextToImg:
 
     def __init__(self):
-        options = OptionsTextToImage(
-            prompt="Astronaut in a jungle, cold color palette, "
-                   "muted colors, detailed, 8k",
-            device=Devices.GPU,
-            image_width=512,
-            image_height=512,
-        )
-
         model_stabilityai_name = "stabilityai/sdxl-turbo"
         model_stabilityai_path = "stabilityai/sdxl-turbo"
         model_management = ModelsManagement()
         model_stabilityai = ModelTextToImage(model_stabilityai_name,
-                                             model_stabilityai_path)
+                                             model_stabilityai_path,
+                                             Devices.GPU,
+                                             torch_dtype=torch.float16,
+                                             use_safetensors=True,
+                                             add_watermarker=False,
+                                             variant="fp16")
 
-        model_management.add_model(new_model=model_stabilityai,
-                                   model_options=options)
+        model_management.add_model(new_model=model_stabilityai)
         model_management.load_model(model_stabilityai_name)
 
-        image = model_management.generate_prompt()
+        image = model_management.generate_prompt(
+            prompt="Astronaut in a jungle, cold color palette, "
+                   "muted colors, detailed, 8k",
+            image_width=512,
+            image_height=512
+        )
         image.show()

sdk/downloader.py

@@ -254,12 +254,11 @@ def set_transformers_class_names(model: Model) -> None:
         # Get the configuration
         config = transformers.AutoConfig.from_pretrained(model.name)
 
-        # Map model class from model type
-        model_mapping = transformers.AutoModel._model_mapping._model_mapping
-
         # Set model class name if not already set
-        model.class_name = model.class_name or model_mapping.get(
-            config.model_type)
+        model.class_name = model.class_name or config.architectures[0] \
+            if config.architectures and config.architectures[0] else (
+            model_config_default_class_for_module[TRANSFORMERS]
+        )
 
         # Set tokenizer class name if not already set
         # and config.tokenizer_class exists

sdk/models/__init__.py

@@ -1,5 +1,6 @@
 # flake8: noqa F401
 from sdk.models.model import Model
+from sdk.models.model_transformers import ModelTransformers
 from sdk.models.models_management import ModelsManagement
 from sdk.models.model_text_to_image import ModelTextToImage
 from sdk.models.model_text_conversation import ModelsTextConversation

sdk/models/model.py

@@ -1,6 +1,4 @@
 from abc import abstractmethod
-from sdk.options import Options
-from typing import Optional
 
 
 class Model:
@@ -20,14 +18,13 @@ def __init__(self, model_name, model_path: str):
         self.model_path = model_path
 
     @abstractmethod
-    def load_model(self, option: Options) -> bool:
+    def load_model(self) -> bool:
         raise NotImplementedError
 
     @abstractmethod
     def unload_model(self) -> bool:
         raise NotImplementedError
 
     @abstractmethod
-    def generate_prompt(self, prompt: Optional[str],
-                        option: Options, **kwargs):
+    def generate_prompt(self, prompt: str, **kwargs):
         raise NotImplementedError